freddyaboulton HF staff commited on
Commit
375e731
1 Parent(s): 7232d54

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from gradio_pdf import PDF
3
+ from pdf2image import convert_from_path
4
+ from transformers import pipeline
5
+ from pathlib import Path
6
+
7
+ dir_ = Path(__file__).parent
8
+
9
+ p = pipeline(
10
+ "document-question-answering",
11
+ model="impira/layoutlm-document-qa",
12
+ )
13
+
14
+ def qa(doc: str, question: str) -> str:
15
+ img = convert_from_path(doc)[0]
16
+ output = p(img, question)
17
+ return sorted(output, key=lambda x: x["score"], reverse=True)[0]['answer']
18
+
19
+
20
+ demo = gr.Interface(
21
+ qa,
22
+ [PDF(label="Document"), gr.Textbox()],
23
+ gr.Textbox(),
24
+ examples=[[str(dir_ / "invoice_2.pdf"), "What is the total gross worth?"]]
25
+ )
26
+
27
+ demo.launch()