Create handler.py

handler.py

@@ -0,0 +1,72 @@
+from utilities.setup import *
+
+import json
+import os
+
+from typing import Dict, List, Any
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        """Initialize class. Load model of interest upon init."""
+
+
+        print("Reading config")
+        self.path = path
+        self.HF_TOKEN = os.getenv("HF_TOKEN")
+        self.wd = os.getcwd()
+        self.model_name = os.path.basename(self.wd)
+
+        print("loading model")
+        self.model, self.tokenizer = self.load_model()
+
+
+    def load_model(self):
+        """Load unsloth model and tokenizer"""
+
+        model = AutoPeftModelForCausalLM.from_pretrained(
+            self.path,
+            load_in_4bit = True,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(self.path)
+        
+        return model, tokenizer
+
+
+    def prompt_formatter(self, prompt):
+        """Prompts must be formatted in alpaca style prior to API."""
+        inputs = self.tokenizer([prompt], return_tensors = "pt").to("cuda")
+        
+        return inputs, prompt
+
+
+    def infer(self, prompt, max_new_tokens=1000): # add streaming capability
+        """Bringing it all together"""
+        # load model
+        inputs, prompt_text = self.prompt_formatter(prompt)
+        outputs = self.model.generate(**inputs,
+                                max_new_tokens = max_new_tokens, 
+                                use_cache=True)
+        completion = self.tokenizer.batch_decode(outputs)
+
+        return completion
+
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `str`)
+            kwargs
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        if data["inputs"] is not None:
+            request = data['inputs']
+            
+            prediction = self.infer(request)
+            return {"prediction": prediction}
+        else:
+            return [{"Error" : "no input received."}]