Update app/main.py

app/main.py

@@ -1,14 +1,16 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 from glin_profanity import Filter
 import re
 
-# 1. Keep the toxicity classifier
+# 1. Toxicity classifier
 toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
 
-# 2. SWAP to an Instruction-following model for rewriting
-detox_pipe = pipeline("text2text-generation", model="google/flan-t5-base")
+# 2. UPGRADED: Use a highly capable Chat Model (Phi-3)
+model_id = "microsoft/Phi-3-mini-4k-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+llm_pipe = pipeline("text-generation", model=model_id, tokenizer=tokenizer)
 
 profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})
 
@@ -48,32 +50,36 @@ async def moderate_text(request: TextRequest):
     
     rewritten = None
     if is_flagged:
-        # 3. Create a prompt that tells the AI exactly how to rewrite it
-        prompt = f"Rewrite this sentence to be completely safe, polite, and non-threatening: '{text}'"
+        # 3. Craft a strict prompt asking for 5 different options
+        messages = [
+            {"role": "system", "content": "You are a helpful, polite assistant. Rewrite the user's toxic message into exactly 5 different safe, polite, and non-threatening sentences. Number them 1 to 5. Do not include any other text in your response."},
+            {"role": "user", "content": text}
+        ]
         
-        outputs = detox_pipe(
-            prompt, 
-            max_length=50, 
-            num_return_sequences=5, 
-            do_sample=True,
-            temperature=0.9,
-            top_p=0.9
-        )
+        # Generate response
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        outputs = llm_pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
         
-        # Clean up the outputs to make sure they are unique
-        suggestions = set()
-        for out in outputs:
-            clean_text = out['generated_text'].strip()
-            # Only add if it actually changed the text
-            if clean_text.lower() != text.lower() and clean_text:
-                suggestions.add(clean_text)
+        generated_text = outputs[0]['generated_text'][len(prompt):].strip()
         
-        # Convert the set back to a list
-        rewritten = list(suggestions)
+        # 4. Parse the output into a clean list
+        lines = generated_text.split('\n')
+        rewritten = []
+        for line in lines:
+            line = line.strip()
+            if line and line[0].isdigit(): # Only keep lines that start with a number
+                # Remove the numbering (e.g., "1. " or "1- ")
+                clean_line = re.sub(r'^\d+[\.\-\)]\s*', '', line).strip()
+                # Remove quotes if the AI added them
+                clean_line = clean_line.strip('"\'')
+                if clean_line:
+                    rewritten.append(clean_line)
         
-        # Fallback if the AI failed to generate anything different
+        # Keep only the first 5, and strip any duplicates
+        rewritten = list(dict.fromkeys(rewritten))[:5]
+
         if not rewritten:
-            rewritten = ["[AI could not generate a safe rewrite, please revise manually]"]
+            rewritten = ["[AI could not generate a rewrite, please revise manually]"]
 
     return ModerationResponse(
         original_text=text,