import gradio as gr
from transformers import pipeline

# model = AutoModelForCausalLM.from_pretrained("Karzan/ckb-gpt2-medium-base-test-1024")
model_id = "Karzan/bart-qa-ckb"
pipe = pipeline(
    "question-answering",
    model=model_id,
    max_answer_len=300,
    handle_impossible_answer=True,
    top_k=3,
)


def func(context, question):
    result = pipe(question=question, context=context)
    return result[0]["answer"], result[0]["score"], result[1]["answer"], result[1]["score"], result[2]["answer"], result[2]["score"]


interface = gr.Interface(
    fn=func,
    inputs=[
        gr.Textbox(lines=7, label="Context paragraph"),
        gr.Textbox(lines=2, label="Question"),
    ],
    outputs=[
        gr.Textbox(label="Answer"),
        gr.Textbox(label="Score"),
        gr.Textbox(label="Answer"),
        gr.Textbox(label="Score"),
        gr.Textbox(label="Answer"),
        gr.Textbox(label="Score"),
    ],
    # submit_btn=submit_btn,
)

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