import gradio as gr
import torch
from detoxify import Detoxify
from transformers import AutoTokenizer, AutoModelForSequenceClassification


def run_model(model_choice, text):
    if model_choice == "original":
        results = Detoxify("original").predict(text)
    elif model_choice == "unbiased":
        results = Detoxify("unbiased").predict(text)
    elif model_choice == "multilingual":
        results = Detoxify("multilingual").predict(text)
    elif model_choice == "Toxic-BERT":
        tokenizer = AutoTokenizer.from_pretrained(
            "citizenlab/distilbert-base-multilingual-cased-toxicity"
        )
        model = AutoModelForSequenceClassification.from_pretrained(
            "citizenlab/distilbert-base-multilingual-cased-toxicity"
        )
        # Perform inference with the Toxic-BERT model using tokenizer and model

        encoded_input = tokenizer(
            text, padding=True, truncation=True, max_length=512, return_tensors="pt"
        )
        logits = model(**encoded_input).logits
        probabilities = torch.sigmoid(logits).detach().cpu().numpy().tolist()[0]
        results = {"toxic": probabilities[0], "non_toxic": probabilities[1]}
        # Convert the predicted_labels to your desired output format

    return results


model_choices = ["original", "unbiased", "multilingual", "Toxic-BERT"]

input_textbox = gr.inputs.Textbox(label="Input Text")
model_choice_dropdown = gr.inputs.Dropdown(choices=model_choices, label="Model Choice")
output_text = gr.outputs.Textbox(label="Output Tags")

iface = gr.Interface(
    fn=run_model,
    inputs=[model_choice_dropdown, input_textbox],
    outputs=output_text,
    title="Toxicity Detection App",
    description="Choose a model and input text to detect toxicity.",
)

iface.launch()