import gradio as gr
from transformers import pipeline

# Part 1: Load the model ONCE
print("Loading the MobileBERT model...")
info_extractor = pipeline("question-answering", model="csarron/mobilebert-uncased-squad-v2")
print("Model loaded successfully!")

# Part 2: Create the function that the UI will call
# This function takes the document and question from the UI,
# gets the answer from the model, and returns it.
def extract_information(context, question):
    print(f"Extracting answer for question: '{question}'")
    result = info_extractor(question=question, context=context)
    return result['answer']

# Part 3: Build and launch the Gradio Interface
print("Launching Gradio interface...")
iface = gr.Interface(
    fn=extract_information,
    inputs=[
        gr.Textbox(lines=7, label="Document", placeholder="Paste the document or text you want to ask questions about..."),
        gr.Textbox(label="Question", placeholder="What specific detail are you looking for?")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="💡 Efficient Information Extractor",
    description="Ask a question about the document below to pull out specific details using a MobileBERT model."
)

iface.launch()