Update paraphraser.py

paraphraser.py

@@ -1,55 +1,55 @@
 # paraphraser.py
 from model_loader import paraphraser_model
+import time
 
 def paraphrase_comment(comment):
     """
-    Paraphrase a toxic comment using the Granite 3.2-2B-Instruct model.
+    Paraphrase a given comment using the fine-tuned Granite 3.2-2B-Instruct model to make it non-toxic, empathetic, and professional while retaining the original intent.
     Returns the paraphrased comment.
     """
-    if not comment:
-        return None
-
     try:
+        start_time = time.time()
+        print("Starting paraphrasing...")
+
+        # Access the model and tokenizer
         model = paraphraser_model.model
         tokenizer = paraphraser_model.tokenizer
 
-        # Create a detailed prompt with guidelines and examples
+        # Define the prompt for paraphrasing
         prompt = (
-            "You are a content moderator tasked with rewriting toxic comments into neutral and constructive ones while maintaining the original meaning. "
-            "Follow these 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"
-            "- Match the length and brevity of the original toxic comment whenever possible. Keep the response short and to the point.\n\n"
-            "Examples:\n"
-            "Toxic: \"You're so dumb! You never understand anything!\"\n"
-            "Neutral: \"You might be misunderstanding this.\"\n"
-            "Toxic: \"This is the worst idea ever. Only an idiot would suggest this.\"\n"
-            "Neutral: \"I don’t think this idea works well.\"\n"
-            "Toxic: \"You’re useless.\"\n"
-            "Neutral: \"This isn’t helping much.\"\n"
-            "Toxic: \"Shut up.\"\n"
-            "Neutral: \"Let’s take a break from this.\"\n\n"
-            f"Now, rewrite this comment: \"{comment}\""
+            f"You are a content moderator tasked with paraphrasing a comment to make it non-toxic, empathetic, and professional while retaining the original intent. "
+            f"The original comment is: \"{comment}\". "
+            f"Guidelines: "
+            f"- Remove any hate speech, offensive language, or toxic elements. "
+            f"- Use a neutral or positive tone. "
+            f"- Ensure the paraphrased comment is concise and clear. "
+            f"- Maintain the core message or intent of the original comment. "
+            f"Provide the paraphrased comment only, without additional explanation."
         )
+
+        # Tokenize the prompt
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
 
-        # Generate the paraphrased comment with optimized parameters
+        # Generate the paraphrased comment
         outputs = model.generate(
-            **inputs,
-            max_new_tokens=50,  # Specify the number of new tokens to generate (excludes input length)
-            num_beams=4,  # Use beam search for faster and more consistent generation
-            early_stopping=True,  # Stop generation once a good sequence is found
-            do_sample=False  # Disable sampling to use beam search
+            inputs["input_ids"],
+            max_length=512,
+            num_beams=5,
+            no_repeat_ngram_size=2,
+            early_stopping=True
         )
 
+        # Decode the output
         paraphrased_comment = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Remove the prompt part from the output
-        paraphrased_comment = paraphrased_comment.replace(prompt, "").strip()
-        # Remove unwanted prefixes like "Neutral: "
-        if paraphrased_comment.startswith("Neutral: "):
-            paraphrased_comment = paraphrased_comment[len("Neutral: "):].strip()
-        return paraphrased_comment
+
+        # Clean up the output (remove the prompt and any extra text)
+        if prompt in paraphrased_comment:
+            paraphrased_comment = paraphrased_comment.replace(prompt, "").strip()
+        paraphrased_comment = paraphrased_comment.strip()
+
+        print(f"Paraphrasing completed in {time.time() - start_time:.2f} seconds")
+        return paraphrased_comment if paraphrased_comment else "Error: Unable to generate paraphrase."
 
     except Exception as e:
-        return f"Error paraphrasing comment: {str(e)}"
+        print(f"Error during paraphrasing: {str(e)}")
+        return "Error: Unable to generate paraphrase."