Update classifier.py

classifier.py

@@ -5,10 +5,10 @@ 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, and color for UI display.
+    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
+        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)
@@ -24,4 +24,14 @@ def classify_toxic_comment(comment):
     confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()
     label_color = "red" if label == "Toxic" else "green"
 
-    return f"Prediction: {label}", confidence, label_color
+    # 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