ClaudiaIoana550
commited on
Commit
•
ca05bde
1
Parent(s):
abe5d1c
Update handler.py
Browse files- handler.py +59 -0
handler.py
CHANGED
@@ -8,6 +8,65 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
8 |
|
9 |
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
|
10 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
11 |
class EndpointHandler:
|
12 |
def __init__(self, model_path=""):
|
13 |
tokenizer=AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
|
8 |
|
9 |
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
|
10 |
|
11 |
+
from transformers import (
|
12 |
+
StoppingCriteria,
|
13 |
+
StoppingCriteriaList,
|
14 |
+
pipeline,
|
15 |
+
)
|
16 |
+
from typing import List
|
17 |
+
import torch
|
18 |
+
|
19 |
+
class StopGenerationCriteria(StoppingCriteria):
|
20 |
+
def __init__(self, max_duplicate_sequences=3, max_repeated_words=2):
|
21 |
+
self.generated_sequences = set()
|
22 |
+
self.max_duplicate_sequences = max_duplicate_sequences
|
23 |
+
self.max_repeated_words = max_repeated_words
|
24 |
+
|
25 |
+
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
|
26 |
+
tokenizer=AutoTokenizer.from_pretrained("ClaudiaIoana550/try1_deploy_falcon", trust_remote_code=True)
|
27 |
+
generated_sequence = input_ids.tolist()
|
28 |
+
|
29 |
+
if len(generated_sequence[0]) >= 50:
|
30 |
+
sequen = generated_sequence[0][-30:]
|
31 |
+
s_mare = str(generated_sequence[0]).strip("[]")
|
32 |
+
s_mic = str(sequen).strip("[]")
|
33 |
+
count2 = 0
|
34 |
+
if s_mic in s_mare:
|
35 |
+
count2 = sum(1 for i in range(len(generated_sequence[0]) - len(sequen) + 1) if generated_sequence[0][i:i + len(sequen)] == sequen)
|
36 |
+
if count2 >= 2:
|
37 |
+
return True
|
38 |
+
|
39 |
+
|
40 |
+
generated_tokens = [tokenizer.decode(token_id) for token_id in input_ids[0]]
|
41 |
+
count = 1
|
42 |
+
prev_token = None
|
43 |
+
for token in generated_tokens:
|
44 |
+
if token == prev_token:
|
45 |
+
count += 1
|
46 |
+
if count > self.max_repeated_words:
|
47 |
+
return True
|
48 |
+
else:
|
49 |
+
count = 1
|
50 |
+
prev_token = token
|
51 |
+
|
52 |
+
if len(self.generated_sequences) >= self.max_duplicate_sequences:
|
53 |
+
return True
|
54 |
+
|
55 |
+
return False
|
56 |
+
|
57 |
+
|
58 |
+
# Example usage:
|
59 |
+
# Define the maximum number of duplicate sequences and repeated words
|
60 |
+
max_duplicate_sequences = 1
|
61 |
+
max_repeated_words = 2
|
62 |
+
|
63 |
+
# Create an instance of StopGenerationCriteria
|
64 |
+
stop_criteria = StopGenerationCriteria(max_duplicate_sequences, max_repeated_words)
|
65 |
+
|
66 |
+
# Add the custom stopping criteria to a StoppingCriteriaList
|
67 |
+
stopping_criteria = StoppingCriteriaList([stop_criteria])
|
68 |
+
|
69 |
+
|
70 |
class EndpointHandler:
|
71 |
def __init__(self, model_path=""):
|
72 |
tokenizer=AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|