import gradio as gr
import torch
from transformers import pipeline

app_title = "Portuguese Hate Speech Detection"

app_description = """
This app detects hate speech on Portuguese text using multiple models. You can either introduce your own sentences by filling in "Text" or click on one of the examples provided below.
"""

"""
app_examples = [
    ["As pessoas tem que perceber que ser 'panasca' não é deixar de ser homem, é deixar de ser humano 😂😂"],
    ["Isso pulhiticos merdosos, continuem a importar lixo, até Portugal deixar de ser Portugal."],
    ["Vai pá puta que te pariu seu paneleiro do caralho, virgem ofendida"],
    ["Vamo-nos unir para criar um mundo mais inclusivo e tolerante."],
    ["Eu admiro muito a coragem e a determinação da minha colega de trabalho."],
    ["O tempo está ensolarado hoje, perfeito para um passeio no parque."]
]
"""

app_examples = [
    {
        "text": "As pessoas têm que perceber que ser 'panasca' não é deixar de ser homem, é deixar de ser humano 😂😂",
        "Non Hate Speech": 0.57,
        "Hate Speech": 0.43
    },
    {
        "text": "Isso políticos merdosos, continuem a importar lixo, até Portugal deixar de ser Portugal.",
        "Non Hate Speech": 0.14,
        "Hate Speech": 0.86
    },
    {
        "text": "Vai para a puta que te pariu seu paneleiro do caralho, virgem ofendida",
        "Non Hate Speech": 0.03,
        "Hate Speech": 0.97
    },
    {
        "text": "Vamo-nos unir para criar um mundo mais inclusivo e tolerante.",
        "Non Hate Speech": 0.93,
        "Hate Speech": 0.07
    },
    {
        "text": "Eu admiro muito a coragem e a determinação da minha colega de trabalho.",
        "Non Hate Speech": 1,
        "Hate Speech": 0
    },
    {
        "text": "O tempo está ensolarado hoje, perfeito para um passeio no parque.",
        "Non Hate Speech": 1,
        "Hate Speech": 0
    }
]

model_list = [
    "knowhate/HateBERTimbau",
    "knowhate/HateBERTimbau-youtube",
    "knowhate/HateBERTimbau-twitter",
    "knowhate/HateBERTimbau-yt-tt",
]

#pipe = pipeline("text-classification", model="knowhate/HateBERTimbau")
#demo = gr.Interface.from_pipeline(pipe)
#demo.launch()

def predict(text, chosen_model):    
    
    # Initialize the pipeline with the chosen model
    model_pipeline = pipeline("text-classification", model=chosen_model)
    result = model_pipeline(text)
    
    predicted_label = result[0]['label']
    predicted_score = result[0]['score']

    non_predicted_label = "Hate Speech" if predicted_label == "Non Hate Speech" else "Non Hate Speech"
    non_predicted_score = 1 - predicted_score

    scores_dict = {
        predicted_label: predicted_score,
        non_predicted_label: non_predicted_score
    }
    
    return scores_dict#, predicted_label

inputs = [
    gr.Textbox(label="Text", value= app_examples[0]['text']),
    gr.Dropdown(label="Model", choices=model_list, value=model_list[2])
]

outputs = [
    gr.Label(label="Result"),
]


gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=app_title, 
             description=app_description, examples=app_examples).launch()