app.py fix

app.py

@@ -1,16 +1,20 @@
 import gradio as gr
 from transformers import pipeline
+import sys
+
+# Toggle this to True if you want to see debug prints
+DEBUG = False
 
 # Load the toxicity classification pipeline
 print("Loading toxicity classifier pipeline...")
 toxicity_pipeline = pipeline(
     "text-classification",
     model="s-nlp/roberta_toxicity_classifier",
-    tokenizer="s-nlp/roberta_toxicity_classifier",
-    # return_all_scores=True  # This ensures we get scores for both classes
+    tokenizer="s-nlp/roberta_toxicity_classifier"
 )
 print("Pipeline loaded successfully!")
 
+
 def toxicity_classification(text: str) -> dict:
     """
     Classify the toxicity of the given text.
@@ -19,36 +23,38 @@ def toxicity_classification(text: str) -> dict:
         text (str): The text to analyze
 
     Returns:
-        dict: A dictionary containing toxicity scores and classification
+        dict: A dictionary containing toxicity classification and confidence
     """
     if not text.strip():
-        return {
-            "error": "Please enter some text to analyze"
-        }
-    
+        return {"error": "Please enter some text to analyze"}
+
     try:
-        # Get predictions using the pipeline
+        # Get the top prediction using the pipeline
         result = toxicity_pipeline(text)[0]
-        
-        # The pipeline returns a list like: [{'label': 'LABEL_0', 'score': 0.8}]
-        # LABEL_0 is non-toxic, LABEL_1 is toxic (based on the model description)
-        
-        # Convert to a more readable format
 
-        label = result['label']
-        confidence = result['score']
+        if DEBUG:
+            print(f"DEBUG - Pipeline result: {result}")
+
+        # The model returns labels like "neutral" or "toxic"
+        label = result.get("label", "neutral").lower()
+        score = result.get("score", 0.0)
+
+        # Map "neutral" (or any non-toxic) to non-toxic
+        if label == "toxic":
+            classification = "toxic"
+            confidence = score
+        else:
+            classification = "non-toxic"
+            confidence = score
 
-        classification = "non-toxic" if label == "LABEL_0" else "toxic"
-        
         return {
             "classification": classification,
             "confidence": round(confidence, 4)
         }
-        
+
     except Exception as e:
-        return {
-            "error": f"Error processing text: {str(e)}"
-        }
+        return {"error": f"Error processing text: {str(e)}"}
+
 
 # Create the Gradio interface
 demo = gr.Interface(
@@ -64,48 +70,37 @@ demo = gr.Interface(
     examples=[
         ["You are amazing!"],
         ["This is a wonderful day."],
-        ["I disagree with your opinion."],
-        ["The weather is nice today."]
+        ["I hate you so much!"],
+        ["You're such an idiot!"],
     ]
 )
 
-# Launch the interface and MCP server
+
 if __name__ == "__main__":
-    # Add debugging section for local testing
-    import sys
-    
-    # Check if running in debug mode (pass 'debug' as command line argument)
-    if len(sys.argv) > 1 and sys.argv[1] == "debug":
+    # If "debug" was passed as a command-line argument, run local tests
+    if len(sys.argv) > 1 and sys.argv[1].lower() == "debug":
+        DEBUG = True
         print("=" * 50)
-        print("DEBUG MODE - Testing toxicity classification")
+        print("DEBUG MODE - Testing toxicity classification locally")
         print("=" * 50)
-        
-        # Test cases for debugging
+
         test_cases = [
             "You are amazing!",
             "This is a wonderful day.",
             "I hate you so much!",
-            "The weather is nice today.",
             "You're such an idiot!",
             "I disagree with your opinion.",
             ""  # Empty string test
         ]
-        
-        for i, test_text in enumerate(test_cases):
+
+        for i, test_text in enumerate(test_cases, 1):
             print(f"\n--- Test Case {i} ---")
             print(f"Input: '{test_text}'")
-            
-            # Set breakpoint here for debugging
-            import pdb; pdb.set_trace()
-            
-            # Call the function
             result = toxicity_classification(test_text)
-            
             print(f"Output: {result}")
             print("-" * 30)
-        
+
         print("\nDebug testing completed!")
-        
     else:
-        # Normal Gradio mode
-        demo.launch(mcp_server=True)
+        # Normal Gradio mode: launch with MCP server enabled
+        demo.launch(mcp_server=True)