Update handler.py
Browse files- handler.py +10 -10
handler.py
CHANGED
@@ -22,16 +22,16 @@ class EndpointHandler():
|
|
22 |
print("tokenizer created ", datetime.now())
|
23 |
|
24 |
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
|
36 |
self.generate_text = transformers.pipeline(
|
37 |
model=self.model,
|
|
|
22 |
print("tokenizer created ", datetime.now())
|
23 |
|
24 |
|
25 |
+
stop_token_ids = self.tokenizer.convert_tokens_to_ids(["<|endoftext|>"])
|
26 |
+
|
27 |
+
class StopOnTokens(StoppingCriteria):
|
28 |
+
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
|
29 |
+
for stop_id in stop_token_ids:
|
30 |
+
if input_ids[0][-1] == stop_id:
|
31 |
+
return True
|
32 |
+
return False
|
33 |
+
|
34 |
+
stopping_criteria = StoppingCriteriaList([StopOnTokens()])
|
35 |
|
36 |
self.generate_text = transformers.pipeline(
|
37 |
model=self.model,
|