gradiopdf / app.py
freddyaboulton's picture
Create app.py
375e731
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(doc: str, question: 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,
[PDF(label="Document"), gr.Textbox()],
gr.Textbox(),
examples=[[str(dir_ / "invoice_2.pdf"), "What is the total gross worth?"]]
)
demo.launch()