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("text2text-generation", model=model_id,max_answer_len=300,handle_impossible_answer=True,top_k=1)


def func(context, question):
    result = pipe(question = question, context=context)
    return result['answer']


text_input = gr.Textbox(lines=10, placeholder="Enter your text here", label="Your Text")
text_outputs = gr.Textbox(lines=10)
# submit_btn = gr.Button(value="Generate",variant="secondary")

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

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