Create handler.py

handler.py

@@ -0,0 +1,64 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        #quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
+
+        # device_map = {
+        #     "transformer.word_embeddings": 0,
+        #     "transformer.word_embeddings_layernorm": 0,
+        #     "lm_head": "cpu",
+        #     "transformer.h": 0,
+        #     "transformer.ln_f": 0,
+        # }
+        #path = "anon8231489123/gpt4-x-alpaca-13b-native-4bit-128g"
+
+        self.model =  AutoModelForCausalLM.from_pretrained(
+            path, 
+            device_map="auto",
+            load_in_8bit=True,
+            #kwargs="--wbits 4 --groupsize 128",
+            #device_map=device_map, 
+            #quantization_config=quantization_config
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.pipeline = pipeline("conversational", model = self.model, tokenizer=self.tokenizer, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16)
+    
+        #rep= "anon8231489123/gpt4-x-alpaca-13b-native-4bit-128g"
+        # tokenizer = AutoTokenizer.from_pretrained(rep)
+        #model = AutoModelForCausalLM.from_pretrained(rep)
+
+        # inputs = tokenizer(["Today is"], return_tensors="pt")
+
+        # reply_ids = model.generate(**inputs, max_new_tokens=590) # return_dict_in_generate=True, output_scores=True
+        # outputs = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
+        # print(outputs)
+
+        #modelPath = "/"
+        
+        #self.pipeline = pipeline("conversational", model=modelPath)
+
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+        print("end")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", None)
+
+        # preprocess 
+        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids      
+        
+        # pass inputs with all kwargs in data
+        if parameters is not None:
+            outputs = self.model.generate(input_ids, **parameters)
+        else:
+            outputs = self.model.generate(input_ids)
+        
+        # postprocess the prediction
+        prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return [{"generated_text": prediction}]