import gradio as gr
from transformers import pipeline

# **1. Load a better QA model** – using RoBERTa-large for higher accuracy.
# (You can switch to 'deepset/roberta-base-squad2-distilled' for speed or others as needed.)
MODEL_NAME = "deepset/roberta-large-squad2"
qa_pipeline = pipeline(
    "question-answering", 
    model=MODEL_NAME, 
    tokenizer=MODEL_NAME
    # You can add device=0 here if using a GPU for faster inference
)

# Define the QA function for Gradio
def answer_question(question, context):
    # **2. Use the pipeline with improved parameters**
    result = qa_pipeline(
        question=question, 
        context=context, 
        handle_impossible_answer=True,  # allow "no answer" if applicable
        top_k=1,                        # we only want the best answer
        max_answer_len=30               # increase if expecting longer answers
    )
    answer = result.get("answer", "").strip()
    score = result.get("score", 0.0)
    # **3. Handle cases where no answer is found or confidence is low**
    if answer == "" or score < 0.1:
        # If the model found no answer or is very unsure, return a fallback message
        return "🤔 I’m not sure – the model couldn’t find a clear answer in the text."
    return answer

# **4. Set up Gradio interface** with appropriate input/output components
interface = gr.Interface(
    fn=answer_question, 
    inputs=[
        gr.components.Textbox(lines=2, label="Question"), 
        gr.components.Textbox(lines=10, label="Context")
    ],
    outputs=gr.components.Textbox(label="Answer"),
    title="Question Answering Demo",
    description="Ask a question and get an answer from the provided context. " \
                "Supports unanswerable questions."
)

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