import gradio as gr
from transformers import pipeline

# Load the models
MODEL_PATHS = {
    "Unitary Toxic BERT": "unitary/toxic-bert",
    "Martin HA Toxic Comment": "martin-ha/toxic-comment-model"
}

classifiers = {name: pipeline("text-classification", model=path, tokenizer=path) for name, path in MODEL_PATHS.items()}

def predict_toxicity(text, model_choice):
    # Get predictions
    classifier = classifiers[model_choice]
    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"),
        gr.Radio(choices=list(MODEL_PATHS.keys()), label="Choose a model", value="Toxic BERT")
    ],
    outputs=gr.Label(num_top_classes=6, label="Toxicity Scores"),
    title="Toxicity Prediction",
    description="This POC uses trained&pre-trained models to predict toxicity in text. Choose between two models: 'Toxic BERT' based and 'Martin HA Toxic Comment'.",
    examples=[
        ["Great game everyone!", "Toxic BERT"],
        ["You're such a noob, uninstall please.", "Martin HA Toxic Comment"],
        ["I hope you die in real life, loser.", "Toxic BERT"],
        ["Nice move! How did you do that?", "Martin HA Toxic Comment"],
        ["Go back to the kitchen where you belong.", "Toxic BERT"],
    ]
)

# Launch the app
iface.launch()