import gradio as gr
from transformers import pipeline

# Load lightweight QA model
qa_pipeline = pipeline(
    "question-answering",
    model="distilbert-base-cased-distilled-squad"
)

# Function
def answer(context, question):
    if context.strip() == "" or question.strip() == "":
        return "Please enter context and question."

    result = qa_pipeline(
        question=question,
        context=context
    )

    return result["answer"]

# UI
demo = gr.Interface(
    fn=answer,
    inputs=[
        gr.Textbox(lines=8, label="Context"),
        gr.Textbox(lines=2, label="Question")
    ],
    outputs="text",
    title="🔍 RAG-QA System",
    description="Ask questions from given context.",
    theme=gr.themes.Soft()
)

demo.launch(server_name="0.0.0.0", server_port=7860)