added more files

.gitignore

@@ -1,7 +1,10 @@
 __pycache__
 .mypy_cache
 *.pdf
+*.png
 
 main.py
 setup.md
+invoice_example.png
 initial_files
+test_*

handler.py

@@ -3,9 +3,14 @@ from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor
 import torch
 from subprocess import run
 
-# install tesseract-ocr and pytesseract
 run("apt install -y tesseract-ocr", shell=True, check=True)
-run("pip install pytesseract", shell=True, check=True)
+
+class HugEndpointException(Exception):
+    def __init__(self, e):
+        self.e = e
+
+    def __str__(self):
+        return f'Custom Endpoint Exception: {self.e}'
 
 # helper function to unnormalize bboxes for drawing onto the image
 def unnormalize_box(bbox, width, height):
@@ -37,28 +42,31 @@ class EndpointHandler:
         # process image
         encoding = self.processor(image, return_tensors="pt")
 
-        # run prediction
-        with torch.inference_mode():
-            outputs = self.model(
-                input_ids=encoding.input_ids.to(device),
-                bbox=encoding.bbox.to(device),
-                attention_mask=encoding.attention_mask.to(device),
-                token_type_ids=encoding.token_type_ids.to(device),
-            )
-            predictions = outputs.logits.softmax(-1)
+        try:
+            # run prediction
+            with torch.inference_mode():
+                outputs = self.model(
+                    input_ids=encoding.input_ids.to(device),
+                    bbox=encoding.bbox.to(device),
+                    attention_mask=encoding.attention_mask.to(device),
+                    token_type_ids=encoding.token_type_ids.to(device),
+                )
+                predictions = outputs.logits.softmax(-1)
 
-        # post process output
-        result = []
-        for item, inp_ids, bbox in zip(
-            predictions.squeeze(0).cpu(), 
-            encoding.input_ids.squeeze(0).cpu(), 
-            encoding.bbox.squeeze(0).cpu()
-        ):
-            label = self.model.config.id2label[int(item.argmax().cpu())]
-            if label == "O":
-                continue
-            score = item.max().item()
-            text = self.processor.tokenizer.decode(inp_ids)
-            bbox = unnormalize_box(bbox.tolist(), image.width, image.height)
-            result.append({"label": label, "score": score, "text": text, "bbox": bbox})
-        return {"predictions": result}
+            # post process output
+            result = []
+            for item, inp_ids, bbox in zip(
+                predictions.squeeze(0).cpu(), 
+                encoding.input_ids.squeeze(0).cpu(), 
+                encoding.bbox.squeeze(0).cpu()
+            ):
+                label = self.model.config.id2label[int(item.argmax().cpu())]
+                if label == "O":
+                    continue
+                score = item.max().item()
+                text = self.processor.tokenizer.decode(inp_ids)
+                bbox = unnormalize_box(bbox.tolist(), image.width, image.height)
+                result.append({"label": label, "score": score, "text": text, "bbox": bbox})
+            return {"predictions": result}
+        except Exception as e:
+            raise HugEndpointException(e)

preprocessor_config.json

@@ -0,0 +1,9 @@
+{
+  "apply_ocr": true,
+  "do_resize": true,
+  "feature_extractor_type": "LayoutLMv2FeatureExtractor",
+  "ocr_lang": null,
+  "processor_class": "LayoutLMv2Processor",
+  "resample": 2,
+  "size": 224
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f31380262cd4f276be211189196f190c0268e9cece977d500886a4e4c16fc07
+size 450606565

requirements.txt

@@ -0,0 +1 @@
+pytesseract

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,39 @@
+{
+  "additional_special_tokens": null,
+  "apply_ocr": false,
+  "cls_token": "[CLS]",
+  "cls_token_box": [
+    0,
+    0,
+    0,
+    0
+  ],
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "name_or_path": "microsoft/layoutlmv2-base-uncased",
+  "never_split": null,
+  "only_label_first_subword": true,
+  "pad_token": "[PAD]",
+  "pad_token_box": [
+    0,
+    0,
+    0,
+    0
+  ],
+  "pad_token_label": -100,
+  "processor_class": "LayoutLMv2Processor",
+  "sep_token": "[SEP]",
+  "sep_token_box": [
+    1000,
+    1000,
+    1000,
+    1000
+  ],
+  "special_tokens_map_file": null,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "LayoutLMv2Tokenizer",
+  "unk_token": "[UNK]"
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c56fc4a68a8102016f0d13df85e3cef173b08bfd50400f2f88c520a325d11676
+size 3375