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

model = AutoModelForQuestionAnswering.from_pretrained("i0xs0/Fine-Tuned-XLM-Question-Answering")
tokenizer = AutoTokenizer.from_pretrained("i0xs0/Fine-Tuned-XLM-Question-Answering")


def generate_answer(question, context):
    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt")
    input_ids = inputs["input_ids"].tolist()[0]

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

    answer_start = torch.argmax(answer_start_scores)
    answer_end = torch.argmax(answer_end_scores) + 1

    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end]))
    return answer

iface = gr.Interface(fn=generate_answer,
                     inputs=[gr.Textbox(lines=2, placeholder="Enter Question Here..."),
                             gr.Textbox(lines=5, placeholder="Enter Context Here...", label="Context")],
                     outputs=gr.Textbox(lines=5),
                     title="Question Answering",
                     description="Type in your question and Context, and the system will provide you with an answer.")

iface.launch()