app.py

import gradio as gr
from transformers import pipeline

# 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
)
print("Pipeline loaded successfully!")

def toxicity_classification(text: str) -> dict:
    """
    Classify the toxicity of the given text.

    Args:
        text (str): The text to analyze

    Returns:
        dict: A dictionary containing toxicity scores and classification
    """
    if not text.strip():
        return {
            "error": "Please enter some text to analyze"
        }
    
    try:
        # Get predictions 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']

        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)}"
        }

# Create the Gradio interface
demo = gr.Interface(
    fn=toxicity_classification,
    inputs=gr.Textbox(
        placeholder="Enter text to analyze for toxicity...",
        lines=3,
        label="Input Text"
    ),
    outputs=gr.JSON(label="Toxicity Analysis Results"),
    title="Text Toxicity Classification",
    description="Analyze text toxicity using RoBERTa transformer model (s-nlp/roberta_toxicity_classifier)",
    examples=[
        ["You are amazing!"],
        ["This is a wonderful day."],
        ["I disagree with your opinion."],
        ["The weather is nice today."]
    ]
)

# 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":
        print("=" * 50)
        print("DEBUG MODE - Testing toxicity classification")
        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):
            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)