Update handler.py

handler.py

@@ -1,4 +1,4 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from transformers_stream_generator import init_stream_support
 import re
 init_stream_support()
@@ -22,45 +22,45 @@ Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when
 
 class EndpointHandler():
 
-    def __init__(self, path = "."):
+    def __init__(self, path = ""):
+        path = ""
+        # quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload = True)
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
             device_map = "auto",
-            load_in_8bit = True
+            load_in_8bit = True,
+            torch_dtype = "auto",
+            low_cpu_mem_usage = True
         )
 
     def __call__(self, data):
         inputs = data.pop("inputs", data)
-        try:
-            prompt = template.format(
-                user_name = inputs["user_name"],
-                user_input = "\n".join(inputs["user_input"])
-            )
-            input_ids = self.tokenizer(
-                prompt,
-                return_tensors="pt"
-            ) .input_ids
-            stream_generator = self.model.generate(
-                input_ids,
-                max_new_tokens = 50,
-                do_sample = True,
-                do_stream = True,
-                temperature = 0.5,
-                top_p = 0.9,
-                top_k = 0,
-                repetition_penalty = 1.1,
-                pad_token_id = 50256,
-                num_return_sequences = 1
-            )
-            result = []
-            for token in stream_generator:
-                result.append(self.tokenizer.decode(token))
-                if len(result) != 1 and result[-1] == "\n":
-                    return {
-                        "message": " ".join(filter(None, re.sub("\*.*?\*", "", "".join(result).strip()).split()))
-                    }
-        except Exception as e:
-            return {
-                "error": str(e)
-            }
+        prompt = template.format(
+            user_name = inputs["user_name"],
+            user_input = "\n".join(inputs["user_input"])
+        )
+        input_ids = self.tokenizer(
+            prompt,
+            return_tensors = "pt"
+        ).input_ids
+        stream_generator = self.model.generate(
+            input_ids,
+            max_new_tokens = 50,
+            do_sample = True,
+            do_stream = True,
+            temperature = 0.5,
+            top_p = 0.9,
+            top_k = 0,
+            repetition_penalty = 1.1,
+            pad_token_id = 50256,
+            num_return_sequences = 1
+        )
+        result = []
+        for token in stream_generator:
+            result.append(self.tokenizer.decode(token))
+            response = "".join(result).strip()
+            if len(response) != 0 and result[-1] == "\n":
+                return {
+                    "message": " ".join(filter(None, re.sub("\*.*?\*", "", response).split()))
+                }