safetensors

handler.py

@@ -1,6 +1,6 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import re
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from transformers_stream_generator import init_stream_support
+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.
 <START>
@@ -14,46 +14,38 @@ 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.
-{user_name}: Definetly.
 <END>
 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}
-Alice Gate:"""
+"""
 
 class EndpointHandler():
 
-    def __init__(self, path = ""):
+    def __init__(self, path=""):
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit = True,
+            llm_int8_threshold = 0.0,
+            llm_int8_enable_fp32_cpu_offload = True
+        )
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
-            low_cpu_mem_usage = True, 
-            trust_remote_code = False,
-            torch_dtype = torch.float16,
-        ).to('cuda')
-    
-    def response(self, result, user_name):
-        result = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
-        try:
-            result = result[:[m.start() for m in re.finditer(r'[.!?]', result)][-1]+1]
-        except Exception: pass
-        return {
-            "message": result
-        }
+            device_map = "auto"
+            torch_dtype = "auto",
+            low_cpu_mem_usage = True,
+            quantization_config = quantization_config
+        )
 
     def __call__(self, data):
-        inputs = data.pop("inputs", data)
-        user_name = inputs["user_name"]
-        user_input = "\n".join(inputs["user_input"])
+        prompt += data.pop("inputs", data)
         input_ids = self.tokenizer(
-            template.format(
-                user_name = user_name,
-                user_input = user_input
-            ),
-            return_tensors = "pt"
-        ).to("cuda")
-        generator = self.model.generate(
-            input_ids["input_ids"],
-            max_new_tokens = 50,
+            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,
@@ -61,4 +53,8 @@ class EndpointHandler():
             pad_token_id = 50256,
             num_return_sequences = 1
         )
-        return self.response(self.tokenizer.decode(generator[0], skip_special_tokens=True), user_name)
+        result = []
+        for token in stream_generator:
+            result.append(self.tokenizer.decode(token))
+            if result[-1] == "\n": 
+                return "".join(result).strip()

requirements.txt

@@ -1,3 +1,5 @@
 accelerate==0.18.0
 bitsandbytes==0.37.2
+safetensors==0.3.1
+transformers-stream-generator==0.0.4
 transformers @ git+https://github.com/huggingface/transformers.git@151425ddb29d4ad1a121e8cce62000a2ac52d3ba