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

# Load model
model_name = "distilbert-base-cased-distilled-squad"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

def answer_question(question, context):
    if not question or not context:
        return "Please enter both question and context.", ""

    inputs = tokenizer(question, context, return_tensors="pt")
    outputs = model(**inputs)

    start = torch.argmax(outputs.start_logits)
    end = torch.argmax(outputs.end_logits) + 1

    answer = tokenizer.convert_tokens_to_string(
        tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start:end])
    )

    score = float(torch.max(outputs.start_logits).item())

    return answer, f"Confidence Score: {score:.2f}"


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("# 📘 Extractive Question Answering System")
    gr.Markdown(
        "Enter a **context paragraph** and ask a **question**. "
        "The model will extract the exact answer from the text."
    )

    with gr.Row():
        with gr.Column():
            question = gr.Textbox(
                label="❓ Question",
                placeholder="e.g., Who created Python?"
            )

            context = gr.Textbox(
                label="📄 Context",
                lines=8,
                placeholder="Paste your paragraph here..."
            )

            submit_btn = gr.Button("🔍 Get Answer", variant="primary")
            clear_btn = gr.Button("🧹 Clear")

        with gr.Column():
            answer_output = gr.Textbox(label="✅ Extracted Answer")
            score_output = gr.Textbox(label="📊 Confidence Score")

    # Example inputs (very important for demo)
    gr.Examples(
        examples=[
            [
                "Who created Python?",
                "Python is a programming language created by Guido van Rossum and first released in 1991."
            ],
            [
                "What does Hugging Face do?",
                "Hugging Face is a company that develops tools for natural language processing and machine learning."
            ]
        ],
        inputs=[question, context],
    )

    submit_btn.click(
        fn=answer_question,
        inputs=[question, context],
        outputs=[answer_output, score_output]
    )

    clear_btn.click(
        fn=lambda: ("", "", "", ""),
        inputs=[],
        outputs=[question, context, answer_output, score_output]
    )

demo.launch()