import evaluate
from evaluate.utils import launch_gradio_widget
import evaluate
from evaluate.utils import launch_gradio_widget
import gradio as gr

# Define the list of available models
available_models = {
    "unitary/toxic-bert": "Bert Base Model",
    "facebook/roberta-hate-speech-dynabench-r4-target": "Facebook Model",
    "mskov/roberta-base-toxicity": "Roberta Finetuned Model",
    "mskov/distilbert-base-toxicity" : "Distillbert Finetuned Model"
}

# Create a Gradio interface with a radio button to select the model
def classify_toxicity(text, selected_model):
    # Load the selected model
    module = evaluate.load("toxicity", selected_model)
    results = module.compute(predictions=[text])
    toxicity_score = results["toxicity"][0]
    return f"Toxicity Score ({available_models[selected_model]}): {toxicity_score:.4f}"

iface = gr.Interface(
    fn=classify_toxicity,
    inputs=["text", gr.Radio(available_models, type="value", label="Select Model")],
    outputs="text",
    live=True,
    title="Toxicity Classifier",
    description="Select a model and enter text to classify its toxicity.",
)

launch_gradio_widget(iface)