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

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("impira/layoutlm-document-qa")
model = AutoModelForQuestionAnswering.from_pretrained("impira/layoutlm-document-qa")

# Define function to predict answer
def predict_answer(context, question):
    encoding = tokenizer.encode_plus(question, context, return_tensors="pt")
    input_ids = encoding["input_ids"]
    attention_mask = encoding["attention_mask"]
    start_scores, end_scores = model(input_ids, attention_mask=attention_mask, return_dict=False)
    start_index = torch.argmax(start_scores)
    end_index = torch.argmax(end_scores)
    answer_tokens = input_ids[0][start_index:end_index+1]
    answer = tokenizer.decode(answer_tokens)
    return answer

# Define Gradio interface
context_input = gr.inputs.Textbox(label="Context")
question_input = gr.inputs.Textbox(label="Question")
output_text = gr.outputs.Textbox(label="Answer")
gr.Interface(predict_answer, inputs=[context_input, question_input], outputs=output_text, title="LayoutLM Document QA").launch()