import gradio as gr
#from transformers import pipeline
from simpletransformers.question_answering import QuestionAnsweringModel

# create pipeline for QA

#qa = pipeline("question-answering")
model = QuestionAnsweringModel("bert", "./fine_tunned_model",use_cuda=False)


def answer_question(context, question):
    predictions, raw_outputs = model.predict(
        [
            {
                "context": context,
                "qas": [
                    {
                        "question": question,
                        "id": "0",
                    }
                ],
            }
        ]
    )
    
    return predictions[0]['answer'][0]


iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(type="text", label="Context"),
        gr.Textbox(type="text", label="Question")
    ],
    outputs=gr.Textbox(type="text", label="Answer")
)
iface.launch()
#%%