Support Transformers v5 cache handling

#4
by lavrenko - opened
Files changed (1) hide show
  1. custom_generate/generate.py +15 -1
custom_generate/generate.py CHANGED
@@ -304,8 +304,22 @@ def _group_beam_search(
304
  batch_size * num_beams, dtype=torch.long, device=device
305
  )
306
 
 
 
 
 
 
 
 
 
 
307
  # do one decoder step on all beams of all sentences in batch
308
- model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)
 
 
 
 
 
309
 
310
  # prepare variable output controls (note: some models won't accept all output controls)
311
  model_inputs.update(
 
304
  batch_size * num_beams, dtype=torch.long, device=device
305
  )
306
 
307
+ # in cached decoding, Transformers v5 expects only the latest token
308
+ # after the first step
309
+ is_first_iteration = cur_len == decoder_prompt_len
310
+ next_sequence_length = (
311
+ None
312
+ if is_first_iteration or not model_kwargs.get("use_cache", True)
313
+ else 1
314
+ )
315
+
316
  # do one decoder step on all beams of all sentences in batch
317
+ model_inputs = model.prepare_inputs_for_generation(
318
+ input_ids,
319
+ next_sequence_length=next_sequence_length,
320
+ is_first_iteration=is_first_iteration,
321
+ **model_kwargs,
322
+ )
323
 
324
  # prepare variable output controls (note: some models won't accept all output controls)
325
  model_inputs.update(