Upload generate.py
Browse files- generate.py +219 -0
generate.py
ADDED
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
def custom_generate(
|
2 |
+
self,
|
3 |
+
input_ids,
|
4 |
+
attention_mask=None,
|
5 |
+
max_new_tokens=None,
|
6 |
+
min_length=None,
|
7 |
+
do_sample=None,
|
8 |
+
early_stopping=None,
|
9 |
+
num_beams=None,
|
10 |
+
temperature=None,
|
11 |
+
top_k=None,
|
12 |
+
top_p=None,
|
13 |
+
repetition_penalty=None,
|
14 |
+
bad_words_ids=None,
|
15 |
+
bos_token_id=None,
|
16 |
+
pad_token_id=None,
|
17 |
+
eos_token_id=None,
|
18 |
+
streamer=None,
|
19 |
+
length_penalty=None,
|
20 |
+
no_repeat_ngram_size=None,
|
21 |
+
num_return_sequences=None,
|
22 |
+
decoder_start_token_id=None,
|
23 |
+
use_cache=None,
|
24 |
+
num_beam_groups=None,
|
25 |
+
diversity_penalty=None,
|
26 |
+
prefix_allowed_tokens_fn=None,
|
27 |
+
output_attentions=None,
|
28 |
+
output_hidden_states=None,
|
29 |
+
output_scores=None,
|
30 |
+
return_dict_in_generate=None,
|
31 |
+
forced_bos_token_id=None,
|
32 |
+
forced_eos_token_id=None,
|
33 |
+
remove_invalid_values=None,
|
34 |
+
synced_gpus=None,
|
35 |
+
**kwargs,
|
36 |
+
):
|
37 |
+
if input_ids is None or input_ids.nelement() == 0:
|
38 |
+
# If input_ids is None or an empty tensor, create a default input tensor
|
39 |
+
input_ids = torch.LongTensor([[self.tokenizer.bos_token_id]]).to(self.device)
|
40 |
+
attention_mask = torch.ones_like(input_ids).to(self.device)
|
41 |
+
|
42 |
+
device = input_ids.device
|
43 |
+
with torch.no_grad():
|
44 |
+
batch_size = input_ids.shape[0]
|
45 |
+
finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
46 |
+
generated_token_ids = torch.full((batch_size, max_new_tokens), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
|
47 |
+
|
48 |
+
for cur_token_idx in range(max_new_tokens):
|
49 |
+
# Sample the next token
|
50 |
+
new_ids = self(
|
51 |
+
input_ids[~finished_generating],
|
52 |
+
attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
|
53 |
+
**kwargs
|
54 |
+
)['logits']
|
55 |
+
|
56 |
+
# Mask out the start and end thought tokens so we don't accidentally sample them
|
57 |
+
new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
|
58 |
+
|
59 |
+
for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
|
60 |
+
# Find the index of the last token that is not padding
|
61 |
+
base_answer_ids = input_ids[answer_idx]
|
62 |
+
new_answer_ids = new_ids[list_idx]
|
63 |
+
last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
|
64 |
+
|
65 |
+
new_ids_sampled = torch.multinomial(
|
66 |
+
torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temperature, dim=-1), 1)
|
67 |
+
|
68 |
+
# Assign the new id to the last token
|
69 |
+
if last_token_idx + 1 >= len(base_answer_ids):
|
70 |
+
# Add padding everywhere
|
71 |
+
new_padding = torch.full((batch_size, 1), self.tokenizer.pad_token_id, dtype=torch.long,
|
72 |
+
device=device)
|
73 |
+
input_ids = torch.cat([input_ids, new_padding], dim=-1)
|
74 |
+
if attention_mask is not None:
|
75 |
+
attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
|
76 |
+
|
77 |
+
if attention_mask is not None:
|
78 |
+
attention_mask[answer_idx, last_token_idx + 1] = 1
|
79 |
+
input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
|
80 |
+
generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
|
81 |
+
|
82 |
+
if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
|
83 |
+
finished_generating[answer_idx] = 1
|
84 |
+
|
85 |
+
# Check if the end token is generated
|
86 |
+
if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
|
87 |
+
finished_generating[answer_idx] = 1
|
88 |
+
|
89 |
+
if finished_generating.all():
|
90 |
+
break
|
91 |
+
|
92 |
+
if streamer is not None:
|
93 |
+
streamer.put(new_ids_sampled)
|
94 |
+
|
95 |
+
return generated_token_ids
|
96 |
+
|
97 |
+
|
98 |
+
def generate(
|
99 |
+
self,
|
100 |
+
input_ids,
|
101 |
+
attention_mask=None,
|
102 |
+
max_new_tokens=None,
|
103 |
+
min_length=None,
|
104 |
+
do_sample=None,
|
105 |
+
early_stopping=None,
|
106 |
+
num_beams=None,
|
107 |
+
temperature=1.1,
|
108 |
+
streamer=None,
|
109 |
+
top_k=None,
|
110 |
+
top_p=None,
|
111 |
+
repetition_penalty=None,
|
112 |
+
bad_words_ids=None,
|
113 |
+
bos_token_id=None,
|
114 |
+
pad_token_id=None,
|
115 |
+
eos_token_id=None,
|
116 |
+
length_penalty=None,
|
117 |
+
no_repeat_ngram_size=None,
|
118 |
+
num_return_sequences=None,
|
119 |
+
decoder_start_token_id=None,
|
120 |
+
use_cache=None,
|
121 |
+
num_beam_groups=None,
|
122 |
+
diversity_penalty=None,
|
123 |
+
prefix_allowed_tokens_fn=None,
|
124 |
+
output_attentions=None,
|
125 |
+
output_hidden_states=None,
|
126 |
+
output_scores=None,
|
127 |
+
return_dict_in_generate=None,
|
128 |
+
forced_bos_token_id=None,
|
129 |
+
forced_eos_token_id=None,
|
130 |
+
remove_invalid_values=None,
|
131 |
+
synced_gpus=None,
|
132 |
+
n_ahead=4,
|
133 |
+
n_ahead_talk=4,
|
134 |
+
merged_talk_heads=True,
|
135 |
+
merged_lm_and_talk_heads=False,
|
136 |
+
merged_lm_and_think_heads=True,
|
137 |
+
use_concat_talk_head=True,
|
138 |
+
use_shallow_think=True,
|
139 |
+
use_shallow_talk=False,
|
140 |
+
use_complex_think_head=False,
|
141 |
+
use_complex_talk_head=True,
|
142 |
+
use_weighted_talk_head=True,
|
143 |
+
trust_remote_code=True,
|
144 |
+
torch_dtype=torch.bfloat16,
|
145 |
+
**model_kwargs,
|
146 |
+
):
|
147 |
+
|
148 |
+
if max_new_tokens is None:
|
149 |
+
max_new_tokens = 128
|
150 |
+
|
151 |
+
# Set model attributes
|
152 |
+
self.max_thoughts = n_ahead + n_ahead_talk + 1
|
153 |
+
self.merged_talk_heads = merged_talk_heads
|
154 |
+
self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
|
155 |
+
self.merged_lm_and_think_heads = merged_lm_and_think_heads
|
156 |
+
self.use_concat_talk_head = use_concat_talk_head
|
157 |
+
self.use_shallow_think = use_shallow_think
|
158 |
+
self.use_shallow_talk = use_shallow_talk
|
159 |
+
self.use_complex_think_head = use_complex_think_head
|
160 |
+
self.use_complex_talk_head = use_complex_talk_head
|
161 |
+
self.use_weighted_talk_head = use_weighted_talk_head
|
162 |
+
|
163 |
+
# Set model properties
|
164 |
+
self.use_end_thought_token = True
|
165 |
+
self.use_start_thought_token = True
|
166 |
+
self.n_ahead = n_ahead
|
167 |
+
self.n_passes = 1
|
168 |
+
self.eval_mode = True
|
169 |
+
self.first_run = False
|
170 |
+
self.rm_initialized = True
|
171 |
+
self.original_mode = False
|
172 |
+
|
173 |
+
# Check if the input is a string (for compatibility with text-generation-webui)
|
174 |
+
if isinstance(input_ids, str):
|
175 |
+
input_ids = self.tokenizer.encode(input_ids, return_tensors='pt')
|
176 |
+
|
177 |
+
# Move input_ids and attention_mask to the same device as the model
|
178 |
+
input_ids = input_ids.to(self.device)
|
179 |
+
if attention_mask is not None:
|
180 |
+
attention_mask = attention_mask.to(self.device)
|
181 |
+
|
182 |
+
generated_token_ids = custom_generate(
|
183 |
+
self,
|
184 |
+
input_ids=input_ids,
|
185 |
+
attention_mask=attention_mask,
|
186 |
+
max_new_tokens=max_new_tokens,
|
187 |
+
min_length=min_length,
|
188 |
+
do_sample=do_sample,
|
189 |
+
early_stopping=early_stopping,
|
190 |
+
num_beams=num_beams,
|
191 |
+
temperature=temperature,
|
192 |
+
top_k=top_k,
|
193 |
+
top_p=top_p,
|
194 |
+
repetition_penalty=repetition_penalty,
|
195 |
+
bad_words_ids=bad_words_ids,
|
196 |
+
bos_token_id=bos_token_id,
|
197 |
+
pad_token_id=pad_token_id,
|
198 |
+
eos_token_id=eos_token_id,
|
199 |
+
length_penalty=length_penalty,
|
200 |
+
no_repeat_ngram_size=no_repeat_ngram_size,
|
201 |
+
num_return_sequences=num_return_sequences,
|
202 |
+
decoder_start_token_id=decoder_start_token_id,
|
203 |
+
use_cache=use_cache,
|
204 |
+
num_beam_groups=num_beam_groups,
|
205 |
+
diversity_penalty=diversity_penalty,
|
206 |
+
prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
|
207 |
+
output_attentions=output_attentions,
|
208 |
+
output_hidden_states=output_hidden_states,
|
209 |
+
output_scores=output_scores,
|
210 |
+
return_dict_in_generate=return_dict_in_generate,
|
211 |
+
forced_bos_token_id=forced_bos_token_id,
|
212 |
+
forced_eos_token_id=forced_eos_token_id,
|
213 |
+
remove_invalid_values=remove_invalid_values,
|
214 |
+
synced_gpus=synced_gpus,
|
215 |
+
streamer=streamer,
|
216 |
+
**model_kwargs,
|
217 |
+
)
|
218 |
+
|
219 |
+
return generated_token_ids
|