Support Transformers v5 cache handling

#3
by lavrenko - opened
Files changed (1) hide show
  1. custom_generate/generate.py +15 -1
custom_generate/generate.py CHANGED
@@ -169,8 +169,22 @@ def _constrained_beam_search(
169
 
170
  decoder_prompt_len = input_ids.shape[1] # record the prompt length of decoder
171
  while model._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
172
- model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)
 
 
 
 
 
 
 
 
173
 
 
 
 
 
 
 
174
  # prepare variable output controls (note: some models won't accept all output controls)
175
  model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
176
  model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
 
169
 
170
  decoder_prompt_len = input_ids.shape[1] # record the prompt length of decoder
171
  while model._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
172
+ # In cached decoding, Transformers v5 expects only the latest token
173
+ # after the first step. Keep full-sequence behavior for prefill and
174
+ # for uncached decoding.
175
+ is_first_iteration = cur_len == decoder_prompt_len
176
+ next_sequence_length = (
177
+ None
178
+ if is_first_iteration or not model_kwargs.get("use_cache", True)
179
+ else 1
180
+ )
181
 
182
+ model_inputs = model.prepare_inputs_for_generation(
183
+ input_ids,
184
+ next_sequence_length=next_sequence_length,
185
+ is_first_iteration=is_first_iteration,
186
+ **model_kwargs,
187
+ )
188
  # prepare variable output controls (note: some models won't accept all output controls)
189
  model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
190
  model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})