WaterBERT-RE-v1

WaterBERT-RE-v1

Fine-grained relation classification for wastewater- and water-treatment literature, fine-tuned from WaterBERT. It links treatment processes, pollutants, removal efficiencies, operating parameters and values, so that a removal result can be assembled into a structured record (process → efficiency → pollutant, with its value and conditions).

This is the first-generation relation model. The knowledge graph WaterKG was built with WaterBERT-RE, which uses a simpler three-relation schema (removes, removal_rate, has).

Labels

Valid entity-type pairs are listed in label_config.json.

Input format

Wrap the head entity in [E1_X]…[/E1_X] and the tail in [E2_Y]…[/E2_Y]; the markers are registered as special tokens.

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

name = "Mudi12137/WaterBERT-RE-v1"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name).eval()

text = "The [E1_REM]removal efficiency[/E1_REM] of carbamazepine reached [E2_VAL]95%[/E2_VAL] at pH 7."
with torch.no_grad():
    probs = model(**tok(text, return_tensors="pt", truncation=True)).logits.softmax(-1)[0]
print(model.config.id2label[int(probs.argmax())])   # has_value

Insert markers from right to left (by character offset) and pass the whole abstract or passage (up to 512 tokens).

Where the entity spans come from

Removal_Efficiency and Value are not entity types of WaterBERT-NER; they came from an earlier eight-type NER model. To use this model on new text, supply REM spans (efficiency expressions such as removal efficiency, removal rate, degradation) and VAL spans (numbers, percentages) from your own detector, e.g. keyword and regular-expression rules.

License

Apache-2.0.