Ankur Goyal commited on
Commit
8bd074d
1 Parent(s): 64aa075

Update DocQuery

Browse files
Files changed (2) hide show
  1. app.py +4 -4
  2. requirements.txt +1 -1
app.py CHANGED
@@ -8,7 +8,7 @@ import traceback
8
  import gradio as gr
9
 
10
  import torch
11
- from docquery.pipeline import get_pipeline
12
  from docquery.document import load_bytes, load_document, ImageDocument
13
  from docquery.ocr_reader import get_ocr_reader
14
 
@@ -29,19 +29,19 @@ CHECKPOINTS = {
29
  PIPELINES = {}
30
 
31
 
32
- def construct_pipeline(model):
33
  global PIPELINES
34
  if model in PIPELINES:
35
  return PIPELINES[model]
36
 
37
  device = "cuda" if torch.cuda.is_available() else "cpu"
38
- ret = get_pipeline(checkpoint=CHECKPOINTS[model], device=device)
39
  PIPELINES[model] = ret
40
  return ret
41
 
42
 
43
  def run_pipeline(model, question, document, top_k):
44
- pipeline = construct_pipeline(model)
45
  return pipeline(question=question, **document.context, top_k=top_k)
46
 
47
 
 
8
  import gradio as gr
9
 
10
  import torch
11
+ from docquery import pipeline
12
  from docquery.document import load_bytes, load_document, ImageDocument
13
  from docquery.ocr_reader import get_ocr_reader
14
 
 
29
  PIPELINES = {}
30
 
31
 
32
+ def construct_pipeline(task, model):
33
  global PIPELINES
34
  if model in PIPELINES:
35
  return PIPELINES[model]
36
 
37
  device = "cuda" if torch.cuda.is_available() else "cpu"
38
+ ret = pipeline(task=task, model=CHECKPOINTS[model], device=device)
39
  PIPELINES[model] = ret
40
  return ret
41
 
42
 
43
  def run_pipeline(model, question, document, top_k):
44
+ pipeline = construct_pipeline("document-question-answering", model)
45
  return pipeline(question=question, **document.context, top_k=top_k)
46
 
47
 
requirements.txt CHANGED
@@ -1,4 +1,4 @@
1
  torch
2
  git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
3
- git+https://github.com/impira/docquery.git@c9ba34072b980b04992e5c063805feeaf0480871
4
  sentencepiece
 
1
  torch
2
  git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
3
+ git+https://github.com/impira/docquery.git@8d92692c36f63ef652f3c84cccedd5674ee7b383
4
  sentencepiece