| import gradio as gr |
| from gradio_pdf import PDF |
| from pdf2image import convert_from_path |
| from transformers import pipeline |
| from pathlib import Path |
|
|
| dir_ = Path(__file__).parent |
|
|
| p = pipeline( |
| "document-question-answering", |
| model="impira/layoutlm-document-qa", |
| ) |
|
|
| def qa(question: str, doc: str) -> str: |
| img = convert_from_path(doc)[0] |
| output = p(img, question) |
| return sorted(output, key=lambda x: x["score"], reverse=True)[0]['answer'] |
|
|
|
|
| demo = gr.Interface( |
| qa, |
| [gr.Textbox(label="Question"), PDF(label="Document")], |
| gr.Textbox(), |
| examples=[["What is the total gross worth?", str(dir_ / "invoice_2.pdf")], |
| ["Whos is being invoiced?", str(dir_ / "sample_invoice.pdf")]] |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
