import gradio as gr
from transformers import pipeline

# 1. Load the RoBERTa model for Question Answering
# deepset/roberta-base-squad2 is a highly popular model for this task
model_name = "deepset/roberta-base-squad2"
qa_pipeline = pipeline("question-answering", model=model_name, tokenizer=model_name)

# 2. Define the prediction function
def extract_answer(context, question):
    if not context.strip() or not question.strip():
        return "Please provide both a context and a question."
    
    try:
        # The pipeline returns a dictionary with 'score', 'start', 'end', and 'answer'
        result = qa_pipeline(question=question, context=context)
        return result['answer']
    except Exception as e:
        return f"An error occurred: {str(e)}"

# 3. Create the Gradio Interface
iface = gr.Interface(
    fn=extract_answer,
    inputs=[
        gr.Textbox(lines=10, label="Context Paragraph", placeholder="Paste the text you want the model to read here..."),
        gr.Textbox(lines=2, label="Question", placeholder="What do you want to know from the text?")
    ],
    outputs=gr.Textbox(label="Extracted Answer"),
    title="Extractive QA with RoBERTa",
    description="This application uses a pre-trained RoBERTa model to extract answers to questions based on a provided context.",
    theme="default"
)

# 4. Launch the application
iface.launch()