Update app/main.py

app/main.py

@@ -3,53 +3,50 @@ from pydantic import BaseModel
 from transformers import pipeline
 from glin_profanity import Filter
 
-# Initialize your models
-toxicity_pipe = pipeline(
-    "text-classification", model="textdetox/twitter-xlmr-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")
-profanity_filter = Filter(
-    {"languages": ["english", "spanish", "french"], "replace_with": "***"}
-)
+profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})
 
 app = FastAPI(title="Multilingual Moderation API")
 
-
 class TextRequest(BaseModel):
     text: str
 
-
 class ModerationResponse(BaseModel):
     original_text: str
     detected_categories: dict
     is_flagged: bool
-    rewritten_text: str | None
-
+    rewritten_text: list[str] | None  # <--- CHANGED TO LIST
 
 @app.post("/moderate", response_model=ModerationResponse)
 async def moderate_text(request: TextRequest):
     text = request.text
+    
     profanity_result = profanity_filter.check_profanity(text)
     toxicity_scores = toxicity_pipe(text)[0]
-
-    is_flagged = (
-        profanity_result["contains_profanity"] or toxicity_scores["score"] > 0.5
-    )
+    
+    is_flagged = profanity_result["contains_profanity"] or toxicity_scores['score'] > 0.5
+    
     rewritten = None
     if is_flagged:
-        rewritten = detox_pipe(text, max_length=100)[0]["generated_text"]
-
+        # Generate 5 different suggestions using sampling
+        outputs = detox_pipe(
+            text, 
+            max_length=100, 
+            num_return_sequences=5, 
+            do_sample=True,
+            top_p=0.9
+        )
+        # Extract just the text from the output dictionaries
+        rewritten = [out['generated_text'] for out in outputs]
+    
     return ModerationResponse(
         original_text=text,
-        detected_categories={
-            "profanity": profanity_result["contains_profanity"],
-            "toxicity_model": toxicity_scores,
-        },
+        detected_categories={"profanity": profanity_result["contains_profanity"], "toxicity_model": toxicity_scores},
         is_flagged=is_flagged,
-        rewritten_text=rewritten,
+        rewritten_text=rewritten
     )
 
-
 @app.get("/health")
 async def health_check():
-    return {"status": "ok"}
+    return {"status": "ok"}