fix: Restore explicit .keys() calls for TensorDict compatibility

[pankajbaid567]

Fix: Restore explicit .keys() calls for TensorDict compatibility

Fixes: PreTrainedTokenizerBase.pad() broken on tensordict.TensorDict #42370
Open

What does this PR do?

This PR restores TensorDict compatibility that was broken by PR #37283, which removed explicit .keys() calls from dictionary iteration patterns. The issue caused RuntimeError: generator raised StopIteration when using TensorDict objects with data collators and tokenization utilities.

Problem

TensorDict objects implement __iter__() to iterate over batch dimensions rather than dictionary keys (as standard dicts do). When PR #37283 replaced explicit .keys() calls with implicit iteration (e.g., for key in dict[0]), it broke compatibility with TensorDict and similar Mapping implementations.

Error encountered:

RuntimeError: generator raised StopIteration

This affected:

DataCollatorWithPadding
DataCollatorForLanguageModeling
tokenizer.pad() method
Any code path that converts a list of dict-like objects into a dict of lists
Solution
Restored explicit .keys() calls in 5 critical locations where dict-list conversion happens during padding operations. This ensures compatibility with any object implementing the Mapping protocol, regardless of how it implements iter().

Pattern applied:

# Before (broken for TensorDict):
encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0]}

# After (works for all Mapping objects):
encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()}

Also added a length check to prevent IndexError on empty batches:
if isinstance(encoded_inputs, (list, tuple)) and len(encoded_inputs) > 0 and isinstance(encoded_inputs[0], Mapping):
Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

• Did you read the contributor guideline, Pull Request section?

• Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.

• Did you make sure to update the documentation with your changes? (Added inline comments explaining the fix)

• Did you write any new necessary tests? (Yes - comprehensive test suite with 7 test cases)

Who can review?
@ArthurZucker @Rocketknight1 - This affects tokenizers and data collators @SunMarc - This affects the trainer data collators
Additional Context:

This is a minimal, surgical fix that only touches the specific lines where dict-list conversion happens. The fix is defensive and generalizes well to any Mapping implementation that might have non-standard iter() behavior.

[ligz08]

Imo these tests could be added to tests/tokenization/test_tokenization_utils.py and/or tests/trainer/test_data_collator.py, rather than creating a new .py file

[ligz08]

I don't think the same comment needs to be repeated everywhere. A short inline one like "call .keys() explicitly to avoid issue #42370" is probably better for readability.

[pankajbaid567]

have done the required changes , sir

[Rocketknight1]

Hey, thank you for the PR! As far as I know, we aren't using tensordict anywhere in the codebase, so I don't know if we want tensordict-specific tests. Can we drop the tests and the require_tensordict / is_tensordict_available code, but keep the extra .keys()? The comment should just explain that it needs to stay for tensordict rather than referring to the issue number, since that forces users to look it up on Github to understand why it's there.

[ligz08]

Hey, thank you for the PR! As far as I know, we aren't using tensordict anywhere in the codebase, so I don't know if we want tensordict-specific tests. Can we drop the tests and the require_tensordict / is_tensordict_available code, but keep the extra .keys()?

In my humble opinion it's good to keep some minimal tests against TensorDict to avoid breaking a not-uncommon use case, despite no internal use or official support. Consider my opinion non-blocking though, maintainer’s judgment should take precedence.

The comment should just explain that it needs to stay for tensordict rather than referring to the issue number, since that forces users to look it up on Github to understand why it's there.

Agreed, I take back my previous comment.

[pankajbaid567]

The failing test test_training_gradient_checkpointing_use_reentrant_false in Longformer is unrelated to this PR. Here's why:

Scope of changes: This PR only modifies tokenization padding logic (adding explicit .keys() calls) in 5 tokenizer files. No changes to model architecture or gradient computation.

Luke's gradient checkpointing tests are already skipped: The Luke model (one of the modified tokenizers) has all three gradient checkpointing tests explicitly disabled with @unittest.skip because "This architecture seem to not compute gradients properly when using GC" (see test_modeling_luke.py).

No connection between changes and failure:

Tokenization happens during data preprocessing
Gradient checkpointing happens during model forward/backward pass
These are completely separate phases with no interaction

[pankajbaid567]

@ligz08 @Rocketknight1 any more changes required ?

[Rocketknight1]

Hmn, I'd still remove the tests and is_tensordict_available! It's not a dependency anywhere, so our CI is unlikely to ever have it installed. As a result, those tests are just going to get skipped anyway.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: luke, mluke

[pankajbaid567]

@Rocketknight1 I have made the required changes sir