Create handler.py

handler.py

@@ -0,0 +1,50 @@
+from typing import  Dict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # load the model
+        self.model = transformers.AutoModelForCausalLM.from_pretrained(
+            "gpt2", torch_dtype=torch.float16, output_hidden_states=True
+        )
+        self.model = self.model.cuda()
+        self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+
+    def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
+        """
+        Args:
+            data (:obj:):
+                includes the deserialized audio file as bytes
+        Return:
+            A :obj:`dict`:. base64 encoded image
+        """
+        # process input
+        inputs = data.pop("inputs", data)
+        all_logits = []
+
+        for doc in inputs:
+          tokenized = self.tokenizer(
+              inputs,
+              return_tensors="pt",
+              padding="max_length",
+              truncation=True,
+              max_length=512,
+          )
+          token_ids, token_mask = tokens.input_ids.cuda(), tokens.attention_mask.cuda()
+          with torch.no_grad():
+              out = model(token_ids, attention_mask=token_mask)
+          meaned_logits = (out.logits * token_mask.unsqueeze(-1)).sum(1) / token_mask.sum(
+              1
+          ).unsqueeze(-1)
+          sorted_logits = torch.sort(out.logits).values
+          mean_sorted_logits = (sorted_logits * token_mask.unsqueeze(-1)).sum(
+              1
+          ) / token_mask.sum(1).unsqueeze(-1)
+          all_logits.append(meaned_logits.cpu().numpy().tolist())
+        
+        # postprocess the prediction
+        return {"logits": all_logits}