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/ruanchaves\">Ruan Chaves Rodrigues</a>. Read more about our <a href=\"https://github.com/ruanchaves/eplm\">research on the evaluation of Portuguese language models</a>."

app_title = "Text simplification (Simplificação Textual)"

app_description = """
Given two versions of the same sentence, this app determines which one is more simple.
You can either introduce your own sentences by filling in "Sentence A" and "Sentence B" or click on one of the example pairs provided below.

(Dado duas versões da mesma frase, este aplicativo determina qual é a mais simples.
Você pode introduzir suas próprias frases preenchendo os campos "Sentence A" e "Sentence B" ou clicar em um dos pares de exemplo fornecidos abaixo.)
"""

app_examples = [
    ["O preço para instalar um DVD player no carro fica entre R$ 2 mil e R$ 5 mil.", "Instalar um DVD player no carro tem preço médio entre R$ 2 mil e R$ 5 mil."],
    ["Especialista alerta para cuidados com os olhos.", "Especialista alerta para cuidados com os olhos."],
    ["Para evitar um novo enfraquecimento político, o governo não pretende influir para que os senadores do PMDB contrários à CPMF sejam substituídos na CCJ.", "O governo não pretende influir para que os senadores do PMDB contrários à CPMF sejam substituídos na CCJ."]
]

output_textbox_component_description = """
Output will appear here once the app has finished analyzing the answer.

(A saída aparecerá aqui assim que o aplicativo terminar de analisar a resposta.)
"""

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.)
""" }

score_descriptions = {
    0: "Sentence A is more simple than Sentence B.",
    1: "The two sentences are equally simple.", 
    2: "Sentence B is more simple than Sentence A.",
}

score_descriptions_pt = {
     0: "(A frase A é mais simples que a frase B.)",
     1: "(As duas frases são igualmente simples.)",
     2: "(A frase B é mais simples que a frase A.)",
}

model_list = [
    "ruanchaves/mdeberta-v3-base-porsimplessent",
    "ruanchaves/bert-base-portuguese-cased-porsimplessent",
]

user_friendly_name = {
    "ruanchaves/mdeberta-v3-base-porsimplessent": "mDeBERTa-v3 (PorSimplesSent)",
    "ruanchaves/bert-base-portuguese-cased-porsimplessent": "BERTimbau base (PorSimplesSent)",
}

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, s2, 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], [s2]), 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
    
    scores = { get_description(k):v for k,v in enumerate(logits) }

    return scores


inputs = [
    gr.Textbox(label="Sentence A", value=app_examples[0][0]),
    gr.Textbox(label="Sentence B", value=app_examples[0][1]),
    gr.Dropdown(label="Model", choices=user_friendly_name_list, value=user_friendly_name_list[0])
]

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


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