Update app.py

app.py

@@ -1,3 +1,25 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import torch
 
-gr.Interface.load("models/impira/layoutlm-document-qa").launch()
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+model = AutoModelForQuestionAnswering.from_pretrained("microsoft/layoutlm-base-uncased")
+
+# 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()