import gradio as gr
import numpy as np
import torch
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

from transformers import DistilBertTokenizer, TFDistilBertForQuestionAnswering

# Load model & tokenizer
# model_name = "deepset/roberta-base-squad2"
# model_name = "AmazonScience/qanlu"
model_name = 'distilbert-base-cased-distilled-squad'
# tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# tokenizer = DistilBertTokenizer.from_pretrained(model_name)
# model = TFDistilBertForQuestionAnswering.from_pretrained(model_name)

# model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

def QA_function(context, question):
    inputs = tokenizer(question, context, add_special_tokens=True, return_tensors="pt") # pt for PyTorch tf for tensorflow
    input_ids = inputs["input_ids"].tolist()[0]

    outputs = model(**inputs)
    answer_start_scores = outputs.start_logits
    answer_end_scores = outputs.end_logits

    ## following for torch
    # Get the most likely beginning of answer with the argmax of the score
    answer_start_index = torch.argmax(answer_start_scores)
    # Get the most likely end of answer with the argmax of the score
    answer_end_index = torch.argmax(answer_end_scores) + 1

    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start_index:answer_end_index]))
    return answer

# gradio_ui = gr.Interface(QA_function, [gr.inputs.Textbox(lines=7, label="Context"), gr.inputs.Textbox(label="Question")], gr.outputs.Textbox(label="Answer"))



# a) Get predictions
nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)

def get_answer(context,question):
    QA_input = {
        'question': question,
        'context': context 
    }
    res = nlp(QA_input)
    return res['answer']

gradio_ui = gr.Interface(get_answer, [gr.inputs.Textbox(lines=7, label="Context"), gr.inputs.Textbox(label="Question")], gr.outputs.Textbox(label="Answer"))

gradio_ui.launch()