Upload handler.py

handler.py

@@ -1,4 +1,5 @@
 import numpy as np
+import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from typing import Dict, List, Any
 
@@ -20,18 +21,13 @@ class EndpointHandler:
     def __init__(self, path=""):
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForSequenceClassification.from_pretrained(path)
-        self.id2label = {
-            i: label for i, label in enumerate(self.model.config.id2label.values())
-        }
         self.MAX_LENGTH = 512  # or any other max length you prefer
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # get inputs
         inputs = data.pop("inputs", data)
 
-        encodings = self.tokenizer(
-            inputs, padding=False, truncation=False, return_tensors="pt"
-        )
+        encodings = self.tokenizer(inputs, padding=False, truncation=False)
 
         truncated_input_ids = middle_truncate(
             encodings["input_ids"][0].tolist(), self.MAX_LENGTH, self.tokenizer
@@ -46,21 +42,12 @@ class EndpointHandler:
             "attention_mask": torch.tensor([attention_masks]),
         }
 
+        truncated_encodings.set_format("torch")
+
         outputs = self.model(**truncated_encodings)
 
         # transform logits to probabilities and apply threshold
         probs = 1 / (1 + np.exp(-outputs.logits.detach().cpu().numpy()))
-        predictions = (probs >= 0.5).astype(float)
-
-        # transform predicted id's into actual label names
-        predicted_labels = [
-            self.id2label[idx]
-            for idx, label in enumerate(predictions[0])
-            if label == 1.0
-        ]
 
         # You can return it in any format you like, here's an example:
-        return [
-            {"label": label, "score": prob}
-            for label, prob in zip(predicted_labels, probs[0])
-        ]
+        return [{"scores": probs}]