Update handler.py

handler.py

@@ -1,25 +1,34 @@
+from typing import  Dict, List, Any
 from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import AutoTokenizer
-import torch
+from transformers import AutoTokenizer, pipeline
 
 
 class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
-        self.model = ORTModelForSequenceClassification.from_pretrained(path)
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-
-
-    def __call__(self, data):
-        answers = data.pop("answers")
-        paraphrases = data.pop("paraphrases")
-
-        inputs = self.tokenizer(answers, paraphrases, max_length=253, padding=True, truncation=True, return_tensors='pt')
-
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-
-        logits = outputs.logits
-        predictions = torch.argmax(logits, dim=-1).numpy()
-
-        return list(predictions)
+        model = ORTModelForSequenceClassification.from_pretrained(path)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+
+        self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
+
+
+    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