Update app/main.py

app/main.py

@@ -4,8 +4,12 @@ from transformers import pipeline
 from glin_profanity import Filter
 import re
 
+# 1. Keep the toxicity classifier
 toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
-detox_pipe = pipeline("text2text-generation", model="s-nlp/mt0-xl-detox-orpo")
+
+# 2. SWAP to an Instruction-following model for rewriting
+detox_pipe = pipeline("text2text-generation", model="google/flan-t5-base")
+
 profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})
 
 app = FastAPI(title="Multilingual Moderation API")
@@ -19,30 +23,23 @@ class ModerationResponse(BaseModel):
     is_flagged: bool
     rewritten_text: list[str] | None
 
-# 1. Define explicit threat keywords (You can expand this list)
 THREAT_KEYWORDS = [
     r"\bguns?\b", r"\bweapons?\b", r"\bbombs?\b", r"\bexplosives?\b", 
     r"\bkill\b", r"\bmurder\b", r"\bshoot\b", r"\bstab\b", 
     r"\bdie\b", r"\bdeath\b", r"\bhurt\b"
 ]
-# Compile the regex pattern for fast matching
 threat_pattern = re.compile("|".join(THREAT_KEYWORDS), re.IGNORECASE)
 
 @app.post("/moderate", response_model=ModerationResponse)
 async def moderate_text(request: TextRequest):
     text = request.text
     
-    # Check 1: Profanity filter
     profanity_result = profanity_filter.check_profanity(text)
-    
-    # Check 2: AI Toxicity model
     tox_result = toxicity_pipe(text)[0]
-    is_toxic_by_model = (tox_result['label'].lower() in ['toxic', 'label_1']) and (tox_result['score'] > 0.5)
     
-    # Check 3: Threat Keyword Regex
+    is_toxic_by_model = (tox_result['label'].lower() in ['toxic', 'label_1']) and (tox_result['score'] > 0.5)
     is_threat_by_keyword = bool(threat_pattern.search(text))
     
-    # Final Decision: Flag if ANY check is true
     is_flagged = (
         profanity_result["contains_profanity"] or 
         is_toxic_by_model or 
@@ -51,16 +48,33 @@ 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}'"
+        
         outputs = detox_pipe(
-            text, 
-            max_length=100, 
+            prompt, 
+            max_length=50, 
             num_return_sequences=5, 
             do_sample=True,
-            temperature=0.8,
+            temperature=0.9,
             top_p=0.9
         )
-        rewritten = [out['generated_text'] for out in outputs]
-    
+        
+        # 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)
+        
+        # Convert the set back to a list
+        rewritten = list(suggestions)
+        
+        # Fallback if the AI failed to generate anything different
+        if not rewritten:
+            rewritten = ["[AI could not generate a safe rewrite, please revise manually]"]
+
     return ModerationResponse(
         original_text=text,
         detected_categories={