import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

REPO_ID = "aurelianvolturi/rubert-tiny2-multitask-toxicity"
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForSequenceClassification.from_pretrained(
    REPO_ID, trust_remote_code=True
).eval()

def predict(text):
    batch = tokenizer(
        text, truncation=True, max_length=model.config.max_length,
        return_tensors="pt"
    )
    with torch.inference_mode():
        probabilities = torch.sigmoid(model(**batch).logits)[0].tolist()
    return {
        label: {
            "detected": probability >= model.config.thresholds[label],
            "probability": probability,
        }
        for label, probability in zip(model.config.labels, probabilities)
    }