Pierce Maloney
commited on
Commit
•
965e808
1
Parent(s):
7b65a30
handler
Browse files- handler.py +69 -0
handler.py
ADDED
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import logging
|
2 |
+
from typing import Dict, List, Any
|
3 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
|
4 |
+
|
5 |
+
# Configure logging
|
6 |
+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
7 |
+
|
8 |
+
class EndpointHandler():
|
9 |
+
def __init__(self, path=""):
|
10 |
+
logging.info("Initializing EndpointHandler with model path: %s", path)
|
11 |
+
tokenizer = AutoTokenizer.from_pretrained(path)
|
12 |
+
tokenizer.pad_token = tokenizer.eos_token
|
13 |
+
self.model = AutoModelForCausalLM.from_pretrained(path)
|
14 |
+
self.tokenizer = tokenizer
|
15 |
+
self.stopping_criteria = StoppingCriteriaList([StopAtPeriodCriteria(tokenizer)])
|
16 |
+
|
17 |
+
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
18 |
+
logging.info("Starting inference")
|
19 |
+
inputs = data.pop("inputs", data)
|
20 |
+
additional_bad_words_ids = data.pop("additional_bad_words_ids", [])
|
21 |
+
|
22 |
+
# Log the input size
|
23 |
+
logging.info("Encoding inputs")
|
24 |
+
input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
|
25 |
+
logging.info("Input IDs shape: %s", input_ids.shape)
|
26 |
+
|
27 |
+
max_generation_length = 75 # Desired number of tokens to generate
|
28 |
+
max_input_length = 4092 - max_generation_length # Maximum input length to allow space for generation
|
29 |
+
|
30 |
+
# 3070, 10456, [313, 334], [29898, 1068] corresponds to "(*", and we do not want to output a comment
|
31 |
+
# 13 is a newline character
|
32 |
+
# [1976, 441, 29889], [4920, 441, 29889] is "Abort." [4920, 18054, 29889] is "Aborted."
|
33 |
+
# [2087, 29885, 4430, 29889], [3253, 29885, 4430, 29889] is "Admitted."
|
34 |
+
# [3253, 29885, 4430, 29889]
|
35 |
+
bad_words_ids = [[3070], [313, 334], [10456], [13], [1976, 441, 29889], [2087, 29885, 4430, 29889], [4920, 441], [4920, 441, 29889], [4920, 18054, 29889], [29898, 1068], [3253, 29885, 4430, 29889]]
|
36 |
+
bad_words_ids.extend(additional_bad_words_ids)
|
37 |
+
|
38 |
+
# Truncation and generation logging
|
39 |
+
if input_ids.shape[1] > max_input_length:
|
40 |
+
logging.info("Truncating input IDs to fit within max input length")
|
41 |
+
input_ids = input_ids[:, -max_input_length:]
|
42 |
+
|
43 |
+
max_length = input_ids.shape[1] + max_generation_length
|
44 |
+
|
45 |
+
logging.info("Generating output")
|
46 |
+
generated_ids = self.model.generate(
|
47 |
+
input_ids,
|
48 |
+
max_length=max_length,
|
49 |
+
bad_words_ids=bad_words_ids,
|
50 |
+
temperature=0.2,
|
51 |
+
top_k=40,
|
52 |
+
do_sample=True,
|
53 |
+
stopping_criteria=self.stopping_criteria,
|
54 |
+
)
|
55 |
+
logging.info("Finished generating output")
|
56 |
+
|
57 |
+
generated_text = self.tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
|
58 |
+
prediction = [{"generated_text": generated_text, "generated_ids": generated_ids[0][input_ids.shape[1]:].tolist()}]
|
59 |
+
logging.info("Inference complete")
|
60 |
+
return prediction
|
61 |
+
|
62 |
+
class StopAtPeriodCriteria(StoppingCriteria):
|
63 |
+
def __init__(self, tokenizer):
|
64 |
+
self.tokenizer = tokenizer
|
65 |
+
|
66 |
+
def __call__(self, input_ids, scores, **kwargs):
|
67 |
+
last_token_text = self.tokenizer.decode(input_ids[:, -1], skip_special_tokens=True)
|
68 |
+
logging.info("StopAtPeriodCriteria called. Last token text: '%s'", last_token_text)
|
69 |
+
return '.' in last_token_text
|