import gradio as gr
import torch
from transformers import pipeline

app_title = "kNOwHATE Prototype: OHS Detection (Not for All Audiences)"

app_description = """
This prototype from the kNOwHATE project aims to detect online hate speech (OHS) in European Portuguese. We collected 24,739 YouTube comments and 29,846 tweets, annotated by experts, and trained our prototype on this data. 

We invite you to try it out. You can just enter a sentence below and submit it to see if it contains hate speech.

For more, visit our [website](https://knowhate.eu) and [Hugging Face page](https://huggingface.co/knowhate).
"""

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-twitter"],
    ["Quando não são ciganos são africanos quando não são africanos são brasileiros.  Sempre os mesmos", "knowhate/HateBERTimbau-youtube"],
    ["O tempo está ensolarado hoje, perfeito para um passeio no parque.", "knowhate/HateBERTimbau-yt-tt"]
]

model_list = [
    "knowhate/HateBERTimbau-twitter",
    "knowhate/HateBERTimbau-youtube",
    "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

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, theme=gr.themes.Base(primary_hue="red")).launch()