Update handler.py
Browse files- handler.py +10 -18
handler.py
CHANGED
@@ -9,30 +9,22 @@ LOGGER = logging.getLogger(__name__)
|
|
9 |
|
10 |
class EndpointHandler():
|
11 |
def __init__(self, path=""):
|
12 |
-
self.model = AutoModelForCausalLM.from_pretrained("
|
13 |
-
self.tokenizer = AutoTokenizer.from_pretrained("
|
14 |
# Load the Lora model
|
15 |
|
16 |
-
def __call__(self, data
|
17 |
"""
|
18 |
Args:
|
19 |
data (Dict): The payload with the text prompt and generation parameters.
|
20 |
"""
|
21 |
-
|
22 |
-
|
23 |
-
prompt = data.pop("inputs", None)
|
24 |
-
parameters = data.pop("parameters", None)
|
25 |
-
if prompt is None:
|
26 |
-
raise ValueError("Missing prompt.")
|
27 |
-
# Preprocess
|
28 |
-
input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
29 |
# Forward
|
30 |
LOGGER.info(f"Start generation.")
|
31 |
-
|
32 |
-
|
33 |
-
else:
|
34 |
-
output = self.model.generate(input_ids=input_ids)
|
35 |
# Postprocess
|
36 |
-
|
37 |
-
|
38 |
-
return
|
|
|
9 |
|
10 |
class EndpointHandler():
|
11 |
def __init__(self, path=""):
|
12 |
+
self.model = AutoModelForCausalLM.from_pretrained("Ozgur98/pushed_model_mosaic_small", load_in_8bit=True, device_map='auto')
|
13 |
+
self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
14 |
# Load the Lora model
|
15 |
|
16 |
+
def __call__(self, data):
|
17 |
"""
|
18 |
Args:
|
19 |
data (Dict): The payload with the text prompt and generation parameters.
|
20 |
"""
|
21 |
+
print("CALLED")
|
22 |
+
LOGGER.info(data)
|
|
|
|
|
|
|
|
|
|
|
|
|
23 |
# Forward
|
24 |
LOGGER.info(f"Start generation.")
|
25 |
+
tokenized_example = tokenizer(data, return_tensors='pt')
|
26 |
+
outputs = self.model.generate(tokenized_example['input_ids'].to('cuda:0'), max_new_tokens=100, do_sample=True, top_k=10, top_p = 0.95)
|
|
|
|
|
27 |
# Postprocess
|
28 |
+
answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
29 |
+
prompt = answer[0].rstrip()
|
30 |
+
return prompt
|