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