Upload handler.py

handler.py

@@ -30,16 +30,20 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
 
         encodings = self.tokenizer(
-            inputs, padding=False, truncation=False, max_length=514
+            inputs, padding=False, truncation=False, return_tensors="pt"
         )
+
         truncated_input_ids = middle_truncate(
-            encodings["input_ids"], 514, self.tokenizer
+            encodings["input_ids"][0].tolist(), self.MAX_LENGTH, self.tokenizer
         )
-        truncated_input_ids_array = np.array(truncated_input_ids)
-        attention_masks = (truncated_input_ids_array != 1).astype(int)
+
+        attention_masks = [
+            int(token_id != self.tokenizer.pad_token_id)
+            for token_id in truncated_input_ids
+        ]
         truncated_encodings = {
-            "input_ids": truncated_input_ids,
-            "attention_mask": attention_masks,
+            "input_ids": torch.tensor([truncated_input_ids]),
+            "attention_mask": torch.tensor([attention_masks]),
         }
 
         outputs = self.model(**truncated_encodings)