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. You can either introduce your own sentences by filling in the text box 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 😂😂", "knowhate/HateBERTimbau-youtube"],
    ["Vamo-nos unir para criar um mundo mais inclusivo e tolerante.", "knowhate/HateBERTimbau-twitter"],
    ["Isso pulhiticos merdosos, continuem a importar lixo, até Portugal deixar de ser Portugal.", "knowhate/HateBERTimbau-yt-tt"],
    ["Eu admiro muito a coragem e a determinação da minha colega de trabalho.", "knowhate/HateBERTimbau"],
    ["Vai pá puta que te pariu seu paneleiro do caralho, virgem ofendida", "knowhate/HateBERTimbau-youtube"],
    ["O tempo está ensolarado hoje, perfeito para um passeio no parque.", "knowhate/HateBERTimbau-twitter"]
]

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

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()