milanchndr/Toxic-Comment-Classifier-Explainer

1from fastapi import FastAPI2from pydantic import BaseModel3from transformers import AutoTokenizer, AutoModelForSequenceClassification4import torch5import numpy as np6import re7 8# ----------------------------------------------------------------------9# CONFIG10# ----------------------------------------------------------------------11MODEL_ID = "milanchndr/toxicity-classifier-mdeberta" 12LABELS = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]13MAX_LEN = 25614 15device = torch.device("cuda" if torch.cuda.is_available() else "cpu")16 17# ----------------------------------------------------------------------18# LOAD MODEL19# ----------------------------------------------------------------------20tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)21model = AutoModelForSequenceClassification.from_pretrained(22    MODEL_ID, output_attentions=True23).to(device)24model.eval()25 26# ----------------------------------------------------------------------27# INPUT BODY28# ----------------------------------------------------------------------29class TextInput(BaseModel):30    text: str31 32# ----------------------------------------------------------------------33# FASTAPI APP34# ----------------------------------------------------------------------35app = FastAPI(title="Toxicity Classification API")36 37# ----------------------------------------------------------------------38# UTILS39# ----------------------------------------------------------------------40def clean_text(text):41    """Basic text normalization."""42    text = text.lower()43    text = re.sub(r"http\S+|www\.\S+", " ", text)44    text = re.sub(r"\s+", " ", text).strip()45    return text46 47def clean_token_text(token: str) -> str:48    """49    Removes tokenizer artifacts (like Ġ or leading underscores)50    to produce clean, readable text for the UI.51    """52    # Remove RoBERTa/DeBERTa specific 'Ġ' char53    t = token.replace("Ġ", "")54    # Remove standard whitespace or underscores55    t = t.strip().lstrip("_")56    return t57 58# ----------------------------------------------------------------------59# WORD IMPORTANCE (Leave-one-out)60# ----------------------------------------------------------------------61def compute_word_importance(words, base_pred, label_idx):62    """63    Computes importance by masking one word at a time.64    """65    scores = []66    67    # Limit computation to prevent timeouts on very long texts68    # If text is huge, you might want to limit 'words' to first 10069    70    for i in range(len(words)):71        # Construct text without the i-th word72        perturbed = words[:i] + words[i+1:]73        perturbed_text = " ".join(perturbed)74 75        inputs = tokenizer(76            perturbed_text,77            truncation=True,78            padding="max_length",79            max_length=MAX_LEN,80            return_tensors="pt"81        ).to(device)82 83        with torch.no_grad():84            out = torch.sigmoid(model(**inputs).logits)85            perturbed_pred = out[0, label_idx].item()86 87        # Importance = Drop in confidence when word is removed88        # (Positive score = word contributed to the label)89        scores.append(base_pred - perturbed_pred)90 91    return scores92 93# ----------------------------------------------------------------------94# MAIN ENDPOINT95# ----------------------------------------------------------------------96@app.post("/predict")97async def predict(input_data: TextInput):98    # 1. Prepare Input99    raw_text = input_data.text100    text = clean_text(raw_text)101 102    inputs = tokenizer(103        text,104        truncation=True,105        padding="max_length",106        max_length=MAX_LEN,107        return_tensors="pt",108        return_attention_mask=True109    ).to(device)110 111    # 2. Model Inference112    with torch.no_grad():113        outputs = model(**inputs)114        logits = outputs.logits115        probs = torch.sigmoid(logits)[0].cpu().numpy()116        117        # Extract Attention (Last Layer, Average Heads)118        attentions = outputs.attentions[-1] 119        attn = attentions[0].mean(dim=0).cpu().numpy()120 121    # 3. Process Logic & Labels122    label_idx = int(np.argmax(probs))123    label_name = LABELS[label_idx]124    base_pred = float(probs[label_idx])125    126    # Determine Final Label (highest prob)127    final_label = label_name if base_pred > 0.5 else "neutral"128 129    # 4. Process Tokens & Attention Matrix130    # Get valid length (ignore padding)131    valid_len = int(inputs["attention_mask"][0].sum().item())132    133    # Raw tokens (e.g. ['[CLS]', 'Ġyou', 'Ġare', ...])134    raw_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])[:valid_len]135    136    # Clean tokens for UI (e.g. ['[CLS]', 'you', 'are', ...])137    clean_tokens = [clean_token_text(t) for t in raw_tokens]138    139    # Slice matrix to NxN valid size140    attn_matrix = attn[:valid_len, :valid_len]141 142    # 5. Calculate Top-K Attention Links (Optimization B)143    # Instead of sending full matrix for rendering, send sparse top connections144    attention_links = {}145    for i, token in enumerate(clean_tokens):146        # Skip special tokens for clarity if desired, or keep them147        if token in ["[CLS]", "[SEP]"]: 148            continue149            150        row = attn_matrix[i]151        # Get top 3 indices (excluding self-attention if desired, but keeping here)152        top_indices = row.argsort()[-4:][::-1] 153        154        links = []155        for idx in top_indices:156            # Filter low attention noise157            score = float(row[idx])158            if score > 0.05: 159                links.append([clean_tokens[idx], score])160        161        attention_links[token] = links162 163    # 6. Calculate Word Importance164    # Note: We use space-split words for Leave-One-Out as per original logic165    words = text.split()166    importance_scores = compute_word_importance(words, base_pred, label_idx)167    168    # 7. Calculate Top Contributors (Optimization A)169    # Combine words and scores, sort by score descending170    contributors = [171        {"token": w, "score": s} 172        for w, s in zip(words, importance_scores)173    ]174    # Sort by absolute impact or positive impact? usually positive for "toxicity"175    top_contributors = sorted(contributors, key=lambda x: x["score"], reverse=True)[:5]176 177    # 8. Construct Response178    return {179        "input_text": raw_text,180        "final_label": final_label,181        182        "probabilities": {label: float(p) for label, p in zip(LABELS, probs)},183        184        # Clean tokens for the main display185        "tokens": clean_tokens, 186 187        "attention": {188            "tokens": clean_tokens,189            # Full matrix (good for heatmaps)190            "matrix": [[float(v) for v in row] for row in attn_matrix],191            # Sparse links (good for arc diagrams / fast UI)192            "links": attention_links193        },194 195        "word_importance": {196            "label": str(label_name),197            "tokens": words, # The exact words used for the importance calc198            "importance_scores": [float(v) for v in importance_scores],199            "base_value": float(base_pred - sum(importance_scores)),200            "prediction": float(base_pred),201            "top_contributors": top_contributors202        }203    }