Spaces:
Runtime error
Runtime error
| 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() |