currently running

.gitignore

@@ -0,0 +1,10 @@
+__pycache__
+*.ipynb
+*.pdf
+
+test_endpoint.py
+test_handler_local.py
+
+setup
+upload_to_hf
+requirements.txt

README.md

@@ -15,3 +15,10 @@ LayoutLMv2 is an improved version of LayoutLM with new pre-training tasks to mod
 
 [LayoutLMv2: Multi-modal Pre-training for Visually-Rich Document Understanding](https://arxiv.org/abs/2012.14740)
 Yang Xu, Yiheng Xu, Tengchao Lv, Lei Cui, Furu Wei, Guoxin Wang, Yijuan Lu, Dinei Florencio, Cha Zhang, Wanxiang Che, Min Zhang, Lidong Zhou, ACL 2021
+
+
+Examples & Guides
+
+- https://github.com/NielsRogge/Transformers-Tutorials/blob/master/LayoutLMv2/DocVQA/Fine_tuning_LayoutLMv2ForQuestionAnswering_on_DocVQA.ipynb
+
+- https://mccormickml.com/2020/03/10/question-answering-with-a-fine-tuned-BERT/

handler.py

@@ -0,0 +1,146 @@
+import torch
+from typing import Any
+# from transformers import LayoutLMForTokenClassification
+
+from transformers import LayoutLMv2ForQuestionAnswering
+from transformers import LayoutLMv2Processor
+from transformers import LayoutLMv2FeatureExtractor
+from transformers import LayoutLMv2ImageProcessor
+from transformers import LayoutLMv2TokenizerFast
+
+from PIL import Image, ImageDraw, ImageFont
+from subprocess import run
+import pdf2image
+
+from pprint import pprint
+
+# set device
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# install tesseract-ocr and pytesseract
+# run("apt install -y tesseract-ocr", shell=True, check=True)
+
+feature_extractor = LayoutLMv2FeatureExtractor()
+
+class NoOCRReaderFound(Exception):
+    def __init__(self, e):
+        self.e = e
+
+    def __str__(self):
+        return f"Could not load OCR Reader: {self.e}"
+
+# helper function to unnormalize bboxes for drawing onto the image
+def unnormalize_box(bbox, width, height):
+    return [
+        width * (bbox[0] / 1000),
+        height * (bbox[1] / 1000),
+        width * (bbox[2] / 1000),
+        height * (bbox[3] / 1000),
+    ]
+
+def pdf_to_image(b: bytes):
+    # First, try to extract text directly
+    # TODO: This library requires poppler, which is not present everywhere.
+    # We should look into alternatives. We could also gracefully handle this
+    # and simply fall back to _only_ extracted text
+    images = [x.convert("RGB") for x in pdf2image.convert_from_bytes(b)]
+    encoded_inputs = feature_extractor(images)
+    print('feature_extractor: ', encoded_inputs.keys())
+    data = {}
+    data['image'] = encoded_inputs.pixel_values
+    data['words'] = encoded_inputs.words
+    data['boxes'] = encoded_inputs.boxes
+    return data
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # self.model = LayoutLMForTokenClassification.from_pretrained(path).to(device)
+        # self.processor = LayoutLMv2Processor.from_pretrained(path)
+        self.image_processor = LayoutLMv2ImageProcessor()  # apply_ocr is set to True by default
+        self.tokenizer = LayoutLMv2TokenizerFast.from_pretrained("microsoft/layoutlmv2-base-uncased")
+        # self.processor = LayoutLMv2Processor(self.image_processor, self.tokenizer)
+        self.processor = LayoutLMv2Processor.from_pretrained("microsoft/layoutlmv2-base-uncased")
+        # processor = AutoProcessor.from_pretrained("microsoft/layoutlmv2-base-uncased")
+
+        self.model = LayoutLMv2ForQuestionAnswering.from_pretrained("microsoft/layoutlmv2-base-uncased")
+
+    def __call__(self, data: dict[str, bytes]):
+        """
+        Args:
+            data (:obj:):
+                includes the deserialized image file as PIL.Image
+        """
+        image = data.pop("inputs", data)
+
+        # image = pdf_to_image(image)
+        images = [x.convert("RGB") for x in pdf2image.convert_from_bytes(image)]
+        for image in images:
+            question = "what is the invoice date"
+            encoding = self.processor(
+                image, 
+                question, 
+                return_tensors="pt",
+            )
+            # print(encoding.keys())
+
+            outputs = self.model(**encoding)
+            # print(outputs.keys())
+            predicted_start_idx = outputs.start_logits.argmax(-1).item()
+            predicted_end_idx = outputs.end_logits.argmax(-1).item()
+
+            predicted_answer_tokens = encoding.input_ids.squeeze()[predicted_start_idx : predicted_end_idx + 1]
+            predicted_answer = self.processor.tokenizer.decode(predicted_answer_tokens)
+            print('answer: ',predicted_answer)
+
+            target_start_index = torch.tensor([7])
+            target_end_index = torch.tensor([14])
+
+            outputs = self.model(**encoding, start_positions=target_start_index, end_positions=target_end_index)
+            predicted_answer_span_start = outputs.start_logits.argmax(-1).item()
+            predicted_answer_span_end = outputs.end_logits.argmax(-1).item()
+            print(predicted_answer_span_start, predicted_answer_span_end)
+
+        # pprint(image)
+        # for image, words, boxes in zip(image['image'], image['words'], image['boxes']):
+            # print(image, words, boxes)
+
+            # question = "what is the invoice date"
+            # encoding = self.processor(
+            #     image, 
+            #     question, 
+            #     words,
+            #     boxes=boxes,
+            #     return_tensors="pt",
+            #     # apply_ocr=False
+            # )
+            # print(encoding.keys())
+        
+
+        # 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)
+
+        # # 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}
+        return ''