Update DocQuery

app.py

@@ -8,7 +8,7 @@ import traceback
 import gradio as gr
 
 import torch
-from docquery.pipeline import get_pipeline
+from docquery import pipeline
 from docquery.document import load_bytes, load_document, ImageDocument
 from docquery.ocr_reader import get_ocr_reader
 
@@ -29,19 +29,19 @@ CHECKPOINTS = {
 PIPELINES = {}
 
 
-def construct_pipeline(model):
+def construct_pipeline(task, model):
     global PIPELINES
     if model in PIPELINES:
         return PIPELINES[model]
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    ret = get_pipeline(checkpoint=CHECKPOINTS[model], device=device)
+    ret = pipeline(task=task, model=CHECKPOINTS[model], device=device)
     PIPELINES[model] = ret
     return ret
 
 
 def run_pipeline(model, question, document, top_k):
-    pipeline = construct_pipeline(model)
+    pipeline = construct_pipeline("document-question-answering", model)
     return pipeline(question=question, **document.context, top_k=top_k)
 
 

requirements.txt

@@ -1,4 +1,4 @@
 torch
 git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
-git+https://github.com/impira/docquery.git@c9ba34072b980b04992e5c063805feeaf0480871
+git+https://github.com/impira/docquery.git@8d92692c36f63ef652f3c84cccedd5674ee7b383
 sentencepiece