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 parágrafo e faça perguntas sobre o texto. Quão bem o modelo responde? (este 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='' target='_blank'>NLP nas empresas | Como eu treinei um modelo T5 na tarefa QA em português no Google Colab</a> (26/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

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

question = "Quando começou a pandemia de Covid-19 no mundo?"

def qa(question, context):
  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(placeholder="Digite uma questao aqui:", lines=5),gr.inputs.Textbox(placeholder="Digite um contexto aqui:", lines=5)],
    outputs=[gr.outputs.HTML(label="QA com T5 base")],
    examples=examples
)

iface.launch()