Update handler.py

handler.py

@@ -1,34 +1,40 @@
 from typing import  Dict, List, Any
-from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import pipeline, AutoTokenizer
-
+from transformers import AutoTokenizer, LayoutLMForSequenceClassification
+import torch
 
 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)
+        self.tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+        self.model = LayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased")        # load the optimized model
 
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-        """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
-                - "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)
-        # postprocess the prediction
-        return prediction
+        words = ["Hello", "world"]
+        normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782]
+        
+        token_boxes = []
+        for word, box in zip(words, normalized_word_boxes):
+            word_tokens = tokenizer.tokenize(word)
+            token_boxes.extend([box] * len(word_tokens))
+        # add bounding boxes of cls + sep tokens
+        token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+        
+        encoding = tokenizer(" ".join(words), return_tensors="pt")
+        input_ids = encoding["input_ids"]
+        attention_mask = encoding["attention_mask"]
+        token_type_ids = encoding["token_type_ids"]
+        bbox = torch.tensor([token_boxes])
+        sequence_label = torch.tensor([1])
+        
+        outputs = self.model(
+            input_ids=input_ids,
+            bbox=bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            labels=sequence_label,
+        )
+        
+        loss = outputs.loss
+        logits = outputs.logits
+        return {"logits": logits}
+