Include invoices

app.py

@@ -21,6 +21,7 @@ def ensure_list(x):
 
 CHECKPOINTS = {
     "LayoutLMv1 🦉": "impira/layoutlm-document-qa",
+    "LayoutLMv1 for Invoices 🧾": "impira/layoutlm-invoices",
     "Donut 🍩": "naver-clova-ix/donut-base-finetuned-docvqa",
 }
 
@@ -148,13 +149,12 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
             # prediction for now
             break
 
-        if "start" in p and "end" in p:
+        if "word_ids" in p:
             image = pages[p["page"]]
             draw = ImageDraw.Draw(image, "RGBA")
+            word_boxes = lift_word_boxes(document, p["page"])
             x1, y1, x2, y2 = normalize_bbox(
-                expand_bbox(
-                    lift_word_boxes(document, p["page"])[p["start"] : p["end"] + 1]
-                ),
+                expand_bbox([word_boxes[i] for i in p["word_ids"]]),
                 image.width,
                 image.height,
             )