Update app/main.py

app/main.py

@@ -16,33 +16,44 @@ class ModerationResponse(BaseModel):
     original_text: str
     detected_categories: dict
     is_flagged: bool
-    rewritten_text: list[str] | None  # <--- CHANGED TO LIST
+    rewritten_text: list[str] | None
 
 @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
+    # 1. Get the full result from the toxicity model
+    tox_result = toxicity_pipe(text)[0]
+    
+    # 2. Check if the LABEL is actually toxic AND the confidence is high
+    # (The model usually outputs 'toxic' or 'neutral' / 'LABEL_1' or 'LABEL_0')
+    is_toxic_by_model = (tox_result['label'].lower() in ['toxic', 'label_1']) and (tox_result['score'] > 0.5)
+    
+    # 3. Final Flag Decision
+    is_flagged = profanity_result["contains_profanity"] or is_toxic_by_model
     
     rewritten = None
     if is_flagged:
-        # Generate 5 different suggestions using sampling
+        # Added temperature=0.8 to make the 5 sentences more diverse
         outputs = detox_pipe(
             text, 
             max_length=100, 
             num_return_sequences=5, 
             do_sample=True,
+            temperature=0.8,
             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"], 
+            "model_label": tox_result['label'],
+            "model_score": tox_result['score']
+        },
         is_flagged=is_flagged,
         rewritten_text=rewritten
     )