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

model_name = "distilbert-base-uncased-distilled-squad"

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

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

    inputs = tokenizer(question, context, return_tensors="pt")

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

    answer_start = torch.argmax(outputs.start_logits)
    answer_end = torch.argmax(outputs.end_logits) + 1

    answer = tokenizer.convert_tokens_to_string(
        tokenizer.convert_ids_to_tokens(
            inputs["input_ids"][0][answer_start:answer_end]
        )
    )

    return f"Answer: {answer}"

iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(lines=8, label="Context"),
        gr.Textbox(lines=2, label="Question")
    ],
    outputs=gr.Textbox(label="Predicted Answer"),
    title="Extractive Question Answering",
    description="Ask a question based on the given context."
)

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