import gradio as gr
from transformers import pipeline

# Load your model from the Hugging Face model hub
model_name = "sujra/insurance_Model"
qa_pipeline = pipeline("question-answering", model=model_name, tokenizer=model_name, timeout=120)


def get_answer(question):
    answer = qa_pipeline(question=question, max_length=128, batch_size=4)['answer']  # Reduce batch size and set max_length
    return answer

# Define the input and output components for the UI
question_input = gr.inputs.Textbox(lines=2, label="Enter your question")
output_text = gr.outputs.Textbox(label="Answer")

# Create the UI
gr.Interface(fn=get_answer, inputs=question_input, outputs=output_text, title="Insurance Question Answering").launch()