# -*- coding: utf-8 -*-
from transformers import pipeline
import gradio as gr
from datasets import load_dataset

dataset = load_dataset("xjlulu/ntu_adl_QA", "Preprocess", split="validation")

inputs = [
    gr.inputs.Textbox(lines=3, label="Question"),
    gr.inputs.Textbox(lines=12, label="Context paragraph"),
]

examples = [{"Question": data["question"], "Context paragraph": data["context"]} for data in dataset[2:6]]

model_name = "cluecorpussmall"
question_answering = pipeline("question-answering", model=f"xjlulu/ntu_adl_span_selection_{model_name}")

def generate_answer(question, context):
        result = question_answering(question=question, context=context)
        return result['answer']

iface = gr.Interface(
    fn=generate_answer,
    inputs=inputs,
    outputs=gr.Textbox(lines=10, label="Answer"),
    title="Question Answering",
    description="Answer questions based on a given context paragraph",
    live=False,
    examples=examples,
    share=True,
    theme="huggingface"
)

iface.launch()