Update handler.py
Browse files- handler.py +40 -34
handler.py
CHANGED
@@ -1,5 +1,6 @@
|
|
1 |
-
from transformers import AutoTokenizer, AutoModelForCausalLM
|
2 |
from transformers_stream_generator import init_stream_support
|
|
|
3 |
init_stream_support()
|
4 |
|
5 |
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
|
|
14 |
Alice Gate: Motherboards, they're like puzzles and the backbone of any system.
|
15 |
{user_name}: That sounds great!
|
16 |
Alice Gate: Yeah, it's really fun. I'm lucky to be able to do this as a job.
|
17 |
-
|
18 |
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!
|
|
|
19 |
"""
|
20 |
|
21 |
class EndpointHandler():
|
22 |
|
23 |
-
def __init__(self, path=""):
|
24 |
-
quantization_config = BitsAndBytesConfig(
|
25 |
-
load_in_8bit = True,
|
26 |
-
llm_int8_threshold = 0.0,
|
27 |
-
llm_int8_enable_fp32_cpu_offload = True
|
28 |
-
)
|
29 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
30 |
self.model = AutoModelForCausalLM.from_pretrained(
|
31 |
path,
|
32 |
-
device_map = "auto"
|
33 |
-
|
34 |
-
low_cpu_mem_usage = True,
|
35 |
-
quantization_config = quantization_config
|
36 |
)
|
37 |
|
38 |
def __call__(self, data):
|
39 |
-
|
40 |
-
|
41 |
-
prompt
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
input_ids
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
2 |
from transformers_stream_generator import init_stream_support
|
3 |
+
import re
|
4 |
init_stream_support()
|
5 |
|
6 |
template = """Alice Gate's Persona: Alice Gate is a young, computer engineer-nerd with a knack for problem solving and a passion for technology.
|
|
|
15 |
Alice Gate: Motherboards, they're like puzzles and the backbone of any system.
|
16 |
{user_name}: That sounds great!
|
17 |
Alice Gate: Yeah, it's really fun. I'm lucky to be able to do this as a job.
|
18 |
+
{user_name}: Awesome!
|
19 |
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!
|
20 |
+
{user_input}
|
21 |
"""
|
22 |
|
23 |
class EndpointHandler():
|
24 |
|
25 |
+
def __init__(self, path = "."):
|
|
|
|
|
|
|
|
|
|
|
26 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
27 |
self.model = AutoModelForCausalLM.from_pretrained(
|
28 |
path,
|
29 |
+
device_map = "auto",
|
30 |
+
load_in_8bit = True
|
|
|
|
|
31 |
)
|
32 |
|
33 |
def __call__(self, data):
|
34 |
+
inputs = data.pop("inputs", data)
|
35 |
+
try:
|
36 |
+
prompt = template.format(
|
37 |
+
user_name = inputs["user_name"],
|
38 |
+
user_input = "\n".join(inputs["user_input"])
|
39 |
+
)
|
40 |
+
input_ids = self.tokenizer(
|
41 |
+
prompt,
|
42 |
+
return_tensors="pt"
|
43 |
+
) .input_ids
|
44 |
+
stream_generator = self.model.generate(
|
45 |
+
input_ids,
|
46 |
+
max_new_tokens = 50,
|
47 |
+
do_sample = True,
|
48 |
+
do_stream = True,
|
49 |
+
temperature = 0.5,
|
50 |
+
top_p = 0.9,
|
51 |
+
top_k = 0,
|
52 |
+
repetition_penalty = 1.1,
|
53 |
+
pad_token_id = 50256,
|
54 |
+
num_return_sequences = 1
|
55 |
+
)
|
56 |
+
result = []
|
57 |
+
for token in stream_generator:
|
58 |
+
result.append(self.tokenizer.decode(token))
|
59 |
+
if len(result) != 1 and result[-1] == "\n":
|
60 |
+
return {
|
61 |
+
"message": " ".join(filter(None, re.sub("\*.*?\*", "", "".join(result).strip()).split()))
|
62 |
+
}
|
63 |
+
except Exception as e:
|
64 |
+
return {
|
65 |
+
"error": str(e)
|
66 |
+
}
|