import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from collections import Counter
from scipy.special import softmax

article_string = "Author: <a href=\"https://huggingface.co./FpOliveira\">Felipe Ramos de Oliveira</a>. Read more about our <a href=\"https://github.com/Silly-Machine/TuPi-Portuguese-Hate-Speech-Dataset\">The Portuguese hate speech dataset (TuPI) </a>."

app_title = "Portuguese hate speech classifier (Binary) - Classificador de discurso de ódio em português (Binário)"

app_description = """
EN: This application employs multiple natural language models to identify hate speech in portuguese. You have the option to enter your own phrases by filling in the "Text" field or choosing one of the examples provided below.
\nPT: Esta aplicativo emprega múltiplos modelos de linguagem natural para identificar discuros de odio em português. Você tem a opção de inserir suas próprias frases preenchendo o campo "Text" ou escolhendo um dos exemplos abaixo
"""

app_examples = [
    ["bom dia flor do dia!!!"],
    ["o ódio é muito grande no coração da ex-deputada federal joise hasselmann contra a família bolsonaro"],
    ["mano deus me livre q nojo da porra!🤮🤮🤮🤮🤮"],
    ["obrigada princesa, porra, tô muito feliz snrsss 🤩🤩🤩❤️"],
    ["mds mas o viado vir responder meus status falando q a taylor foi racista foi o auge 😂😂"],
    ["Pra ser minha inimiga no mínimo tem que ter um rostinho bonito e delicado, não se considere minha rival com essa sua cara de cavalo não, feia, cara de traveco, cabeçuda, queixo quadrado 🤣🤣"]
]

output_textbox_component_description = """
EN: This box will display hate speech results based on the average score of multiple models.
PT: Esta caixa exibirá resultados da classicação de discurso de ódio com base na pontuação média de vários modelos.
"""

output_json_component_description = { "breakdown": """
This box presents a detailed breakdown of the evaluation for each model.
""",
"detalhamento": """
(Esta caixa apresenta um detalhamento da avaliação para cada modelo.)
""" }

short_score_descriptions = {
   0: "Not hate",
   1: "Hate"
}

score_descriptions = {
    0: "This text is not a hate speech.",
    1: "This text is a hate speech.",
}

score_descriptions_pt = {
    1: "Este texto contem discurso de ódio",
    0: "Este texto não contem discurso de ódio",
}

model_list = [
    "FpOliveira/tupi-bert-large-portuguese-cased",
    "FpOliveira/tupi-bert-base-portuguese-cased",
    "FpOliveira/tupi-gpt2-small",
]

user_friendly_name = {
    "FpOliveira/tupi-bert-large-portuguese-cased": "BERTimbau large (TuPi)",
    "FpOliveira/tupi-bert-base-portuguese-cased":  "BERTimbau base (TuPi)",
    "FpOliveira/tupi-gpt2-small":"GPT2 small (TuPi)",
}

reverse_user_friendly_name = { v:k for k,v in user_friendly_name.items() }

user_friendly_name_list = list(user_friendly_name.values())

model_array = []

for model_name in model_list:
    row = {}
    row["name"] = model_name
    row["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
    row["model"] = AutoModelForSequenceClassification.from_pretrained(model_name)
    model_array.append(row)
 
def most_frequent(array):
    occurence_count = Counter(array)
    return occurence_count.most_common(1)[0][0]


def predict(s1, chosen_model):
    if not chosen_model:
      chosen_model = user_friendly_name_list[0]
    scores = {}
    full_chosen_model_name = reverse_user_friendly_name[chosen_model]
    for row in model_array:
        name = row["name"]
        if name != full_chosen_model_name:
          continue
        else:
          tokenizer = row["tokenizer"]
          model = row["model"]
          model_input = tokenizer(*([s1],), padding=True, return_tensors="pt")
          with torch.no_grad():
              output = model(**model_input)
              logits = output[0][0].detach().numpy()
              logits = softmax(logits).tolist()
              break
    def get_description(idx):
      description = score_descriptions[idx]
      description_pt = score_descriptions_pt[idx]
      final_description = description + "\n \n" + description_pt
      return final_description
    
    max_pos = logits.index(max(logits))
    markdown_description = get_description(max_pos)
    scores = { short_score_descriptions[k]:v for k,v in enumerate(logits) }

    return scores, markdown_description


inputs = [
    gr.Textbox(label="Text", value=app_examples[0][0]),
    gr.Dropdown(label="Model", choices=user_friendly_name_list, value=user_friendly_name_list[0])
]

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


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