import os
import torch
import gradio as gr

from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering
)

MODEL_NAME = "harishforaiandml/my-pretrained-qa-model"

# Load tokenizer + model manually (SAFE for Spaces)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)

model.eval()


def ask(context, question):

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

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

    start_logits = outputs.start_logits
    end_logits = outputs.end_logits

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

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

    if answer.strip() == "":
        return "Not a valid question"

    return answer


demo = gr.Interface(
    fn=ask,

    inputs=[
        gr.Textbox(lines=6, label="Context"),
        gr.Textbox(label="Question")
    ],

    outputs="text",

    title="QA Model"
)

demo.launch()