Update paraphraser.py

paraphraser.py

@@ -1,48 +1,91 @@
-# paraphraser.py
+# classifier.py
 import torch
-from model_loader import paraphrase_model, paraphrase_tokenizer
+from model_loader import classifier_model
+from paraphraser import paraphrase_comment
+from metrics import compute_semantic_similarity, compute_emotion_shift, compute_empathy_score, compute_bleu_score, compute_rouge_score, compute_entailment_score
 
-def paraphrase_comment(comment):
+def classify_toxic_comment(comment):
     """
-    Paraphrase a toxic comment using the Granite 3.2-2B-Instruct model.
-    Returns the paraphrased 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 additional Stage 3 metrics.
+    Returns the prediction label, confidence, color, toxicity score, bias score, paraphrased comment (if applicable), and its metrics.
     """
-    # Define the paraphrasing prompt with system instruction, guidelines, examples, and the task
-    prompt = (
-        "You are a content moderator tasked with rewriting toxic comments into neutral and constructive ones while maintaining the original meaning.\n"
-        "Guidelines:\n"
-        "- Remove explicit hate speech, personal attacks, or offensive language.\n"
-        "- Keep the response neutral and professional.\n"
-        "- Ensure the rewritten comment retains the original intent but in a constructive tone.\n"
-        "Examples:\n"
-        "Toxic: \"You're so dumb! You never understand anything!\"\n"
-        "Neutral: \"I think there's some misunderstanding. Let's clarify things.\"\n"
-        "Toxic: \"This is the worst idea ever. Only an idiot would suggest this.\"\n"
-        "Neutral: \"I don't think this idea works well. Maybe we can explore other options.\"\n"
-        "Now, rewrite this comment: \"{comment}\""
-    )
-
-    # Format the prompt with the input comment
-    prompt = prompt.format(comment=comment)
-
-    # Tokenize the prompt
-    inputs = paraphrase_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
-
-    # Generate the paraphrased output
+    if not comment.strip():
+        return "Error: Please enter a comment.", None, None, None, None, 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
+    inputs = tokenizer(comment, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+    # Run inference
     with torch.no_grad():
-        outputs = paraphrase_model.generate(
-            **inputs,
-            max_length=512,
-            num_return_sequences=1,
-            do_sample=True,
-            top_p=0.95,
-            temperature=0.7
-        )
-
-    # Decode the generated output
-    paraphrased_comment = paraphrase_tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Remove the prompt part from the output (if the model includes it)
-    paraphrased_comment = paraphrased_comment.replace(prompt, "").strip()
-
-    return paraphrased_comment
+        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)
+
+    # If the comment is toxic, paraphrase it and compute additional metrics
+    paraphrased_comment = None
+    paraphrased_prediction = None
+    paraphrased_confidence = None
+    paraphrased_color = None
+    paraphrased_toxicity_score = None
+    paraphrased_bias_score = None
+    semantic_similarity = None
+    original_emotion = None
+    paraphrased_emotion = None
+    emotion_shift_positive = None
+    empathy_score = None
+    bleu_score = None
+    rouge_scores = None
+    entailment_score = None
+
+    if label == "Toxic":
+        # Paraphrase the comment
+        paraphrased_comment = paraphrase_comment(comment)
+
+        # Re-evaluate the paraphrased comment
+        paraphrased_inputs = tokenizer(paraphrased_comment, return_tensors="pt", truncation=True, padding=True, max_length=512)
+        with torch.no_grad():
+            paraphrased_outputs = model(**paraphrased_inputs)
+            paraphrased_logits = paraphrased_outputs.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)
+
+        # Compute additional Stage 3 metrics
+        semantic_similarity = compute_semantic_similarity(comment, paraphrased_comment)
+        original_emotion, paraphrased_emotion, emotion_shift_positive = compute_emotion_shift(comment, paraphrased_comment)
+        empathy_score = compute_empathy_score(paraphrased_comment)
+        bleu_score = compute_bleu_score(comment, paraphrased_comment)
+        rouge_scores = compute_rouge_score(comment, paraphrased_comment)
+        entailment_score = compute_entailment_score(comment, paraphrased_comment)
+
+    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, f"Original: {original_emotion}, Paraphrased: {paraphrased_emotion}, Positive Shift: {emotion_shift_positive}" if original_emotion else None,
+        empathy_score, bleu_score, rouge_scores, entailment_score
+    )