Update handler.py

handler.py

@@ -1,5 +1,6 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers_stream_generator import init_stream_support
+import re
 init_stream_support()
 
 template = """Alice Gate's Persona: Alice Gate is a young, computer engineer-nerd with a knack for problem solving and a passion for technology.
@@ -14,47 +15,52 @@ Alice Gate: I love exploring, going out with friends, watching movies, and playi
 Alice Gate: Motherboards, they're like puzzles and the backbone of any system.
 {user_name}: That sounds great!
 Alice Gate: Yeah, it's really fun. I'm lucky to be able to do this as a job.
-<END>
+{user_name}: Awesome!
 Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when she sees you. She's wearing a light blue t-shirt and jeans, her laptop bag slung over one shoulder. She takes a seat next to you, her enthusiasm palpable in the air* Hey! I'm so excited to finally meet you. I've heard so many great things about you and I'm eager to pick your brain about computers. I'm sure you have a wealth of knowledge that I can learn from. *She grins, eyes twinkling with excitement* Let's get started!
+{user_input}
 """
 
 class EndpointHandler():
 
-    def __init__(self, path=""):
-        quantization_config = BitsAndBytesConfig(
-            load_in_8bit = True,
-            llm_int8_threshold = 0.0,
-            llm_int8_enable_fp32_cpu_offload = True
-        )
+    def __init__(self, path = "."):
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
-            device_map = "auto"
-            torch_dtype = "auto",
-            low_cpu_mem_usage = True,
-            quantization_config = quantization_config
+            device_map = "auto",
+            load_in_8bit = True
         )
 
     def __call__(self, data):
-        prompt += data.pop("inputs", data)
-        input_ids = self.tokenizer(
-            prompt,
-            return_tensors="pt"
-        ) .input_ids
-        stream_generator = self.model.generate(
-            input_ids,
-            max_new_tokens = 70,
-            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 result[-1] == "\n": 
-                return "".join(result).strip()
+        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)
+            }