import gradio as gr
import time
import pandas as pd
from PIL import Image
import matplotlib as plt

# device
import torch
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

# models
model_name_bb = "pierreguillou/bert-base-cased-squad-v1.1-portuguese"
model_name_bl = "pierreguillou/bert-large-cased-squad-v1.1-portuguese"

# load models
from transformers import pipeline
qa_bb = pipeline("question-answering", model_name_bb, device=device)

from optimum.pipelines import pipeline
qa_bb_better = pipeline("question-answering", model_name_bb, accelerator="bettertransformer", device=device)

from transformers import pipeline
qa_bl = pipeline("question-answering", model_name_bl, device=device)

from optimum.pipelines import pipeline
qa_bl_better = pipeline("question-answering", model_name_bl, accelerator="bettertransformer", device=device)

# function to get results
def get_answer(context, question):

  # get predictions
  start = time.perf_counter()
  answer_bl = qa_bl(question=question, context=context)
  end = time.perf_counter()
  diff_bl = round(end - start, 2)
  answer_bl["time (s)"] = diff_bl
  del answer_bl["start"]
  del answer_bl["end"]

  start = time.perf_counter()
  answer_bl_better = qa_bl_better(question=question, context=context)
  end = time.perf_counter()
  diff_bl_better = round(end - start, 2)
  answer_bl_better["time (s)"] = diff_bl_better
  del answer_bl_better["start"]
  del answer_bl_better["end"]

  start = time.perf_counter()
  answer_bb = qa_bb(question=question, context=context)
  end = time.perf_counter()
  diff_bb = round(end - start, 2)
  answer_bb["time (s)"] = diff_bb
  del answer_bb["start"]
  del answer_bb["end"]

  start = time.perf_counter()
  answer_bb_better = qa_bb_better(question=question, context=context)
  end = time.perf_counter()
  diff_bb_better = round(end - start, 2)
  answer_bb_better["time (s)"] = diff_bb_better
  del answer_bb_better["start"]
  del answer_bb_better["end"]

  answer = dict()
  answer["BERT large"] = answer_bl
  answer["BERT large (BetterTransformer)"] = answer_bl_better
  answer["BERT base"] = answer_bb
  answer["BERT base (BetterTransformer)"] = answer_bb_better

  # get image of prediction times
  df = pd.DataFrame.from_dict({"Method":["BERT base (BetterTransformer)", "BERT base", "BERT large (BetterTransformer)", "BERT large"], 
                               "Time (seconds)":  [answer["BERT base (BetterTransformer)"]["time (s)"], answer["BERT base"]["time (s)"], answer["BERT large (BetterTransformer)"]["time (s)"], answer["BERT large"]["time (s)"]]})
  ax = df.plot.barh(x='Method', title=f'Prediction times on {str(device).replace("cuda:0", "GPU").replace("cpu", "CPU")}')
  ax.figure.savefig("img.png", bbox_inches='tight')
  image = Image.open('img.png')

  return image, answer

title = "QA in Portuguese with BetterTransformer (this App runs on " + str(device).replace("cuda:0", "GPU").replace("cpu", "CPU") + ")"
description = '<p>(20/11/2022) Forneça seu próprio parágrafo e faça perguntas sobre o texto. Quão bem os modelos respondem?<br />(este aplicativo usa os modelos <a href="https://huggingface.co/pierreguillou/bert-base-cased-squad-v1.1-portuguese">pierreguillou/bert-base-cased-squad-v1.1-portuguese</a> and <a href="https://huggingface.co/pierreguillou/bert-large-cased-squad-v1.1-portuguese">pierreguillou/bert-large-cased-squad-v1.1-portuguese</a> and their versions <a href="https://huggingface.co/docs/optimum/bettertransformer/overview">BetterTransformer</a>)</p><p>Blog post sobre BetterTransformer: <a href="https://medium.com/@pierre_guillou/ia-empresas-diminua-o-tempo-de-infer%C3%AAncia-de-modelos-transformer-com-bettertransformer-584a5a7702c8">IA & empresas | Diminua o tempo de inferência de modelos Transformer com BetterTransformer</a></p>'
examples = [
    ["Dom Pedro II foi o segundo e último monarca do Império do Brasil, reinando por mais de 58 anos.", "Quem foi Dom Pedro II?"],
    ["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.", "Quando começou a pandemia de Covid-19 no mundo?"],
    ["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.", "Onde começou a pandemia de Covid-19?"]
    ]

demo = gr.Interface(
    fn=get_answer,
    inputs=[
        gr.Textbox(lines=7, label="Context"), 
        gr.Textbox(lines=2, label="Question")
        ],
    outputs=[
        gr.Image(label="Prediction times", type="pil"),
        gr.JSON(label="Results"),
        ],
    title=title,
    description=description,
    examples=examples,
    allow_flagging="never")

if __name__ == "__main__":
    demo.launch()