import gradio as gr
from transformers import pipeline

# Load the model
MODEL_PATH = "unitary/toxic-bert"
classifier = pipeline("text-classification", model=MODEL_PATH, tokenizer=MODEL_PATH)

def predict_toxicity(text):
    # Get predictions
    predictions = classifier(text, return_all_scores=True)[0]
    
    # Format results
    results = {}
    for pred in predictions:
        results[pred['label']] = f"{pred['score']:.4f}"
    
    return results

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_toxicity,
    inputs=gr.Textbox(lines=5, label="Enter text to analyze"),
    outputs=gr.Label(num_top_classes=6, label="Toxicity Scores"),
    title="Toxicity Prediction",
    description="This POC uses the model based onm toxic-bert to predict toxicity in text. Multi-class response.",
    examples=[
        ["Great game everyone!"],
        ["You're such a noob, uninstall please."],
        ["I hope you die in real life, loser."],
        ["Nice move! How did you do that?"],
        ["Go back to the kitchen where you belong."],
    ]
)

# Launch the app
iface.launch()