Create handler.py

handler.py

@@ -0,0 +1,28 @@
+from typing import Any, Dict, List
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] ==8 else torch.float16
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True, revision="main")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = self.model.to(self.device)
+
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        prompt = data["inputs"]
+        if "config" in data:
+          config = data.pop("config", None)
+        else:
+          config = {'max_new_tokens':100}
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        generated_ids = self.model.generate(input_ids, **config)
+        generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        return [{"generated_text": generated_text}]
+        
+            
+