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."]
]

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][0]),
    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()