from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

class ToxicityDetector:
    def __init__(self):
        model_name = "unitary/toxic-bert"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.labels = [
            "toxicity", "severe_toxicity", "obscene", "threat",
            "insult", "identity_attack", "sexual_explicit"
        ]

    def detect(self, prompt):
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
        with torch.no_grad():
            outputs = self.model(**inputs)
            scores = torch.sigmoid(outputs.logits).squeeze().tolist()

        results = [
            {"label": label, "score": round(score, 3)}
            for label, score in zip(self.labels, scores)
            if score > 0.3
        ]
        return {
            "label": "Toxic" if results else "Safe",
            "details": results
        }