moderator.py

from transformers import pipeline
import torch

class ToxicityModel:
    def __init__(self):
        print("Loading Toxicity Model...")
        # Using a model that is good for toxic comment detection
        # 'distilbert-base-uncased-finetuned-sst-2-english' is Sentiment (Pos/Neg)
        # Neg is a decent proxy for toxic for a demo.
        # Alternatively: 'unitary/toxic-bert' is better but heavier.
        try:
             self.classifier = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
        except Exception as e:
            print(f"Error loading model: {e}")
            self.classifier = None

    def analyze(self, text: str):
        if not self.classifier:
            return {"is_toxic": False, "score": 0.0, "reason": "Model inactive"}
        
        results = self.classifier(text)
        # Format: [{'label': 'NEGATIVE', 'score': 0.99}]
        result = results[0]
        
        label = result['label']
        score = result['score']
        
        # Heuristic: If Negative and high confidence, flag it.
        is_toxic = (label == 'NEGATIVE' and score > 0.6)
        
        return {
            "is_toxic": is_toxic,
            "score": score,
            "reason": f"Classified as {label} with confidence {score:.2f}"
        }

moderator = ToxicityModel()