import gradio as gr
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# model & tokenizer
model_name = "pierreguillou/t5-base-qa-squad-v1.1-portuguese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# parameters
max_target_length=32
num_beams=1
early_stopping=True

title = "QA App | T5 base finetuned on SQuAD 1.1 in Portuguese"
description = "Forneça seu próprio texto e faça perguntas sobre ele. Quão bem o modelo responde? (esse aplicativo usa o modelo https://huggingface.co/pierreguillou/t5-base-qa-squad-v1.1-portuguese)"
article = f"<div style='text-align: center; font-size: 90%;'><p>Link para o modelo <a style='color:blue;' href='https://huggingface.co/pierreguillou/t5-base-qa-squad-v1.1-portuguese' target='_blank'>T5 base QA in Portuguese</a></p><p>Blog post: <a style='color:blue;' href='https://medium.com/@pierre_guillou/nlp-nas-empresas-como-eu-treinei-um-modelo-t5-em-portugu%C3%AAs-na-tarefa-qa-no-google-colab-e8eb0dc38894' target='_blank'>NLP nas empresas | Como eu treinei um modelo T5 em português na tarefa QA no Google Colab</a> (27/01/2022)</p><p>BERT QA in Portuguese App by <a style='color:blue;' href='https://www.linkedin.com/in/pierreguillou/' target='_blank'>Pierre GUILLOU</a> | <a style='color:blue;' href='https://github.com/piegu/language-models#language-models' target='_blank'>Github Repo</a></p></div>"
allow_screenshot = False
allow_flagging = False

question = "Quando começou a pandemia de Covid-19 no mundo?"
context = "A pandemia de COVID-19, também conhecida como pandemia de coronavírus, é uma pandemia em curso de COVID-19, uma doença respiratória aguda causada pelo coronavírus da síndrome respiratória aguda grave 2 (SARS-CoV-2). A doença foi identificada pela primeira vez em Wuhan, na província de Hubei, República Popular da China, em 1 de dezembro de 2019, mas o primeiro caso foi reportado em 31 de dezembro do mesmo ano."

def qa(context, question):
  input_text = "question: " + question + "context: " + context
  inputs = tokenizer(input_text, return_tensors="pt")

  outputs = model.generate(inputs["input_ids"],
                             max_length=max_target_length, 
                             num_beams=num_beams, 
                             early_stopping=early_stopping
                            )
  pred = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
  return pred
             
# interface gradio
iface = gr.Interface(
    title=title,
    description=description,
    article=article,
    allow_screenshot=allow_screenshot,
    allow_flagging=allow_flagging,
    fn=qa,
    inputs=[gr.inputs.Textbox(label="Texto:", default=context, lines=6),gr.inputs.Textbox(label="Questão:", default=question, lines=1)],
    outputs=[gr.outputs.HTML(label="Resposta do modelo T5 base")],
)

iface.launch()