gradio_pdf / _app.py
freddyaboulton's picture
Upload folder using huggingface_hub
3311c9f verified
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()