Update handler.py

handler.py

@@ -4,26 +4,17 @@ from transformers import AutoTokenizer, TextGenerationPipeline
 from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
 # check for GPU
-device = 0 if torch.cuda.is_available() else -1
+#device = 0 if torch.cuda.is_available() else -1
 
 #print(f"cuda: {device}")
 
 class EndpointHandler():
-    def __init__(self, path=""):
-        quantize_config = BaseQuantizeConfig(**{
-                                        "bits": 4,
-                                        "group_size": 128,
-                                        "damp_percent": 0.01,
-                                        "desc_act": False,
-                                        "static_groups": False,
-                                        "sym": True,
-                                        "true_sequential": True
-                                      })        
+    def __init__(self, path=""):     
         # load the optimized model
-        model = AutoGPTQForCausalLM.from_quantized(path, device="cuda:0", quantize_config=quantize_config, use_safetensors=True) #file_name="model-quantized.onnx")
-        tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoGPTQForCausalLM.from_quantized(path, device_map="auto", use_safetensors=True) #file_name="model-quantized.onnx")
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
         # or you can also use pipeline
-        self.generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
+        #self.generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -36,10 +27,15 @@ class EndpointHandler():
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
 
+        input_ids = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
+
         # pass inputs with all kwargs in data
         if parameters is not None:
-            prediction = self.generator(inputs, **parameters)
+            #prediction = self.generator(inputs, **parameters)
+            outputs = self.model.generate(**input_ids, **parameters)
         else:
-            prediction = self.generator(inputs)
+            outputs = self.model.generate(**input_ids)
 
-        return prediction
+        prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return [{"generated_text": prediction}]