Guard `cache_position` against `None` in `prepare_inputs_for_generation`

#52
by wittjeff - opened

Guard cache_position against None in prepare_inputs_for_generation

prepare_inputs_for_generation subscripts cache_position unconditionally:

if cache_position[0] == 0:
    model_inputs["pixel_values"] = pixel_values

Recent transformers releases pass cache_position=None on some generation paths, which crashes with TypeError: 'NoneType' object is not subscriptable at this line. The failure is content-dependent and version-dependent, which makes it confusing to hit in practice:

  • transformers 5.14: fails on a subset of inputs — in our document-OCR benchmarking, 4 of 25 OmniDocBench pages crashed while the rest processed normally.
  • transformers 5.12: fails immediately on the first generation call.

Full traceback (5.12, but the failing frame is identical on 5.14):

File ".../transformers/generation/utils.py", line 3824, in _prefill
  model_inputs = self.prepare_inputs_for_generation(
File ".../modeling_dots_ocr.py", line 128, in prepare_inputs_for_generation
  if cache_position[0] == 0:
TypeError: 'NoneType' object is not subscriptable

This PR adds the standard guard used by in-tree transformers models:

if cache_position is not None and cache_position[0] == 0:

Semantics are preserved: when cache_position is None, pixel_values is simply not re-attached for that step, matching the behavior of the equivalent guard in upstream transformers VLM implementations. Verified on the previously-crashing pages — they now generate and parse normally.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment