Added handler for FUNSD example

handler.py

@@ -1,15 +1,25 @@
+import numpy as np
 from typing import  Dict, List, Any
-from optimum.onnxruntime import ORTModelForSequenceClassification
+from PIL import Image, ImageDraw, ImageFont
+from transformers import LayoutLMv2Processor, LayoutLMv2ForTokenClassification
 from transformers import pipeline, AutoTokenizer
 
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # load the optimized model
-        model = ORTModelForSequenceClassification.from_pretrained(path)
-        tokenizer = AutoTokenizer.from_pretrained(path)
-        # create inference pipeline
-        self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
+        # load the processor and model
+
+        self.processor = LayoutLMv2Processor.from_pretrained("microsoft/layoutlmv2-base-uncased")
+        self.model = LayoutLMv2ForTokenClassification.from_pretrained("nielsr/layoutlmv2-finetuned-funsd")   
+        self.id2label = {
+                        0: 'O',
+                        1: 'B-HEADER',
+                        2: 'I-HEADER',
+                        3: 'B-QUESTION',
+                        4: 'I-QUESTION',
+                        5: 'B-ANSWER',
+                        6: 'I-ANSWER'
+            }
 
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -22,13 +32,34 @@ class EndpointHandler():
                 - "label": A string representing what the label/class is. There can be multiple labels.
                 - "score": A score between 0 and 1 describing how confident the model is for this label/class.
         """
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
-
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
+
+        def unnormalize_box(bbox, width, height):
+            return [
+                width * (bbox[0] / 1000),
+                height * (bbox[1] / 1000),
+                width * (bbox[2] / 1000),
+                height * (bbox[3] / 1000),
+            ]
+
+        image = data.pop("inputs", data)
+        # encode
+        encoding = self.processor(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
+        offset_mapping = encoding.pop('offset_mapping')
+
+        # forward pass
+        outputs = self.model(**encoding)
+
+        # get predictions
+        predictions = outputs.logits.argmax(-1).squeeze().tolist()
+        token_boxes = encoding.bbox.squeeze().tolist()
+
+        # only keep non-subword predictions
+        #is_subword = np.array(offset_mapping.squeeze().tolist())[:,0] != 0
+        width, height = image.size
+
+        true_predictions = [self.id2label[prediction]  for prediction in predictions]
+        true_boxes = [unnormalize_box(box, width, height) for box in token_boxes]
+        is_subword = np.array(offset_mapping.squeeze().tolist())[:,0] != 0
+
         # postprocess the prediction
-        return prediction
+        return {"labels": true_predictions, "boxes": true_boxes, "is_subword": is_subword}