File size: 1,634 Bytes
bd229ab d6b5249 bd229ab d6b5249 bd229ab d6b5249 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
# classifier.py
import torch
from model_loader import model, tokenizer
def classify_toxic_comment(comment):
"""
Classify a comment as toxic or non-toxic using the fine-tuned XLM-RoBERTa model.
Returns the prediction label, confidence, color, toxicity score, and bias score for UI display.
"""
if not comment.strip():
return "Error: Please enter a comment.", None, None, None, None
# Tokenize the input comment
inputs = tokenizer(comment, return_tensors="pt", truncation=True, padding=True, max_length=512)
# Run inference
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# Get the predicted class (0 = non-toxic, 1 = toxic)
predicted_class = torch.argmax(logits, dim=1).item()
label = "Toxic" if predicted_class == 1 else "Non-Toxic"
confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()
label_color = "red" if label == "Toxic" else "green"
# Simulate Toxicity Score (in a real scenario, use a model like Detoxify)
# For now, we'll approximate it based on the confidence of the toxic class
toxicity_score = torch.softmax(logits, dim=1)[0][1].item() # Probability of toxic class
toxicity_score = round(toxicity_score, 2)
# Simulate Bias Score (in a real scenario, use a bias detection model like WEAT)
# For now, we'll use a placeholder value (since the example comment is non-toxic)
bias_score = 0.01 if label == "Non-Toxic" else 0.15 # Placeholder logic
bias_score = round(bias_score, 2)
return f"Prediction: {label}", confidence, label_color, toxicity_score, bias_score |