import gradio as gr
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch

# Load model manually (NO pipeline task)
model_name = "distilbert-base-cased-distilled-squad"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

def answer_question(question, context):
    try:
        inputs = tokenizer(question, context, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)

        start_idx = torch.argmax(outputs.start_logits)
        end_idx = torch.argmax(outputs.end_logits) + 1

        answer = tokenizer.decode(inputs["input_ids"][0][start_idx:end_idx])

        return f"Answer: {answer}"

    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks() as app:
    gr.Markdown("# 🧠 Question Answering (Manual Mode)")

    q = gr.Textbox(value="Where do I work?", label="Question")
    c = gr.Textbox(
        value="My name is Sylvain and I work at Hugging Face in Brooklyn.",
        label="Context",
        lines=4
    )

    out = gr.Textbox(label="Answer")

    gr.Button("Get Answer").click(
        fn=answer_question,
        inputs=[q, c],
        outputs=out
    )

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