Create handler.py

handler.py

@@ -0,0 +1,27 @@
+from typing import Dict, List, Any
+from time import time
+
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+#
+
+
+class EndpointHandler:
+    def __init__(self, path=''):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModel.from_pretrained(path, load_in_8bit=True)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        inputs = data.pop('inputs', data)
+        parameters = data.pop('parameters', {})
+
+        starting_time = time.time()
+
+        tokenized = self.tokenizer(inputs, return_tensors='pt')
+        out = self.model.generate(tokenized.to('cuda'), **parameters).to('cpu')
+        detokenized = self.tokenizer.batch_decode(out)
+
+        ending_time = time.time()
+
+        return [{'generated_text': detokenized, 'generation_time': ending_time-starting_time}]