WaterBERT-RE

WaterBERT-RE

Relation classification for wastewater- and water-treatment literature, fine-tuned from WaterBERT. Given a text with two marked entities, it predicts one of four labels:

This is the relation model used to build the WaterKG knowledge graph. A finer-grained first-generation model with five relations (process, pollutant, removal efficiency, parameter and value) is available as WaterBERT-RE-v1.

Only these entity-type pairs are valid inputs (label_config.json): (Process, Pollutant), (Reactor, Pollutant), (Pollutant, Value), (Process, Parameter), (Reactor, Parameter).

Input format

Wrap the head entity in [E1_X]…[/E1_X] and the tail in [E2_Y]…[/E2_Y], where X/Y is the type code. The markers are already registered as special tokens in the tokenizer.

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

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

text = "[E1_WTP]Ozonation[/E1_WTP] effectively removed [E2_POL]carbamazepine[/E2_POL] from secondary effluent."
with torch.no_grad():
    probs = model(**tok(text, return_tensors="pt", truncation=True)).logits.softmax(-1)[0]
print(model.config.id2label[int(probs.argmax())])   # removes

Insert markers from right to left (by character offset) so earlier offsets stay valid. Pass the whole abstract or passage (up to 512 tokens), not only the sentence, so that cross-sentence relations remain possible.

Where the entity spans come from

POL, WTP, RCT and TRP spans can come from WaterBERT-NER. VAL spans (numbers with units or percentages) are not produced by WaterBERT-NER; they must come from a separate value detector — a regular expression for percentages is enough for removal_rate.

Training

Fine-tuned from WaterBERT on entity pairs drawn from manually annotated abstracts (128,221 candidate pairs, of which 16,160 carry a relation), with a class-weighted loss. Details are in the accompanying paper.

License

Apache-2.0.