Update classifier.py

classifier.py

@@ -4,7 +4,14 @@ import time
 from model_loader import classifier_model
 from paraphraser import paraphrase_comment
 from metrics import compute_semantic_similarity, compute_empathy_score, compute_bias_score, compute_hallucination_score
+from metrics import compute_reward_scores
+import numpy as np
 
+def softmax(logits):
+    exp_logits = np.exp(logits - np.max(logits))
+    return exp_logits / exp_logits.sum()
+
+    
 def compute_reward_scores(original, paraphrased):
     """
     Compute all reward scores for a paraphrase.
@@ -43,88 +50,44 @@ def compute_reward_scores(original, paraphrased):
 
 def classify_toxic_comment(comment):
     """
-    Classify a comment as toxic or non-toxic using the fine-tuned XLM-RoBERTa model.
-    If toxic, paraphrase the comment, re-evaluate, and compute essential metrics.
-    Returns the prediction label, confidence, color, toxicity score, bias score, paraphrased comment (if applicable), and its metrics.
+    Classify a comment for toxicity and compute additional metrics.
+    Returns a dictionary with classification results and scores.
     """
-    start_total = time.time()
-    print("Starting classification...")
-
-    if not comment.strip():
-        return "Error: Please enter a comment.", None, None, None, None, None, None, None, None, None, None, None, None
-
-    # Access the model and tokenizer
-    model = classifier_model.model
-    tokenizer = classifier_model.tokenizer
-
-    # Tokenize the input comment
-    start_classification = time.time()
-    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"
-
-    # Compute Toxicity Score (approximated as the probability of the toxic class)
-    toxicity_score = torch.softmax(logits, dim=1)[0][1].item()
-    toxicity_score = round(toxicity_score, 2)
-
-    # Simulate Bias Score (placeholder)
-    bias_score = 0.01 if label == "Non-Toxic" else 0.15
-    bias_score = round(bias_score, 2)
-    print(f"Classification took {time.time() - start_classification:.2f} seconds")
-
-    # If the comment is toxic, paraphrase it and compute essential metrics
-    paraphrased_comment = None
-    paraphrased_prediction = None
-    paraphrased_confidence = None
-    paraphrased_color = None
-    paraphrased_toxicity_score = None
-    paraphrased_bias_score = None
-    semantic_similarity = None
-    empathy_score = None
-
-    if label == "Toxic":
-        # Paraphrase the comment
-        start_paraphrase = time.time()
-        paraphrased_comment = paraphrase_comment(comment)
-        print(f"Paraphrasing took {time.time() - start_paraphrase:.2f} seconds")
-
-        # Re-evaluate the paraphrased comment
-        start_reclassification = time.time()
-        paraphrased_inputs = tokenizer(paraphrased_comment, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    try:
+        start_time = time.time()
+        print("Starting classification...")
+
+        # Tokenize the comment
+        inputs = classifier_model.tokenizer(
+            comment,
+            return_tensors="pt",
+            truncation=True,
+            padding=True,
+            max_length=512
+        ).to(classifier_model.device)
+
+        # Classify using the toxicity classifier
         with torch.no_grad():
-            paraphrased_outputs = model(**paraphrased_inputs)
-            paraphrased_logits = paraphrased_outputs.logits
+            outputs = classifier_model.model(**inputs)
+            logits = outputs.logits.cpu().numpy()[0]
+            probs = softmax(logits)
 
-        paraphrased_predicted_class = torch.argmax(paraphrased_logits, dim=1).item()
-        paraphrased_label = "Toxic" if paraphrased_predicted_class == 1 else "Non-Toxic"
-        paraphrased_confidence = torch.softmax(paraphrased_logits, dim=1)[0][paraphrased_predicted_class].item()
-        paraphrased_color = "red" if paraphrased_label == "Toxic" else "green"
-        paraphrased_toxicity_score = torch.softmax(paraphrased_logits, dim=1)[0][1].item()
-        paraphrased_toxicity_score = round(paraphrased_toxicity_score, 2)
-        paraphrased_bias_score = 0.01 if paraphrased_label == "Non-Toxic" else 0.15  # Placeholder
-        paraphrased_bias_score = round(paraphrased_bias_score, 2)
-        print(f"Reclassification of paraphrased comment took {time.time() - start_reclassification:.2f} seconds")
+        toxicity = probs[1]  # Assuming label 1 is toxic
+        print(f"Classification took {time.time() - start_time:.2f} seconds")
 
-        # Compute essential metrics
-        start_metrics = time.time()
-        semantic_similarity = compute_semantic_similarity(comment, paraphrased_comment)
-        empathy_score = compute_empathy_score(paraphrased_comment)
-        print(f"Metrics computation took {time.time() - start_metrics:.2f} seconds")
+        # Compute additional metrics (empathy, bias, hallucination, reward)
+        scores = compute_reward_scores(comment, comment)  # Use comment as both original and paraphrase for classification
+        scores["toxicity"] = toxicity  # Override toxicity with classifier result
 
-    print(f"Total processing time: {time.time() - start_total:.2f} seconds")
+        print(f"Total processing time: {time.time() - start_time:.2f} seconds")
+        return scores
 
-    return (
-        f"Prediction: {label}", confidence, label_color, toxicity_score, bias_score,
-        paraphrased_comment, f"Prediction: {paraphrased_label}" if paraphrased_comment else None,
-        paraphrased_confidence, paraphrased_color, paraphrased_toxicity_score, paraphrased_bias_score,
-        semantic_similarity, empathy_score
-    )
+    except Exception as e:
+        print(f"Error during classification: {str(e)}")
+        return {
+            "empathy": 0.0,
+            "toxicity": 1.0,
+            "bias": 1.0,
+            "hallucination": 1.0,
+            "reward": 0.0
+        }