Repository navigation
DataCollatorForLanguageModeling fails with whole-word masking and dynamic padding #49971
Description
System Info
transformersversion: 5.21.0.dev0- Platform: Linux-6.18.54.1-microsoft-standard-WSL2-x86_64-with-glibc2.39
- Python version: 3.12.12
- Huggingface_hub version: 1.33.0
- Safetensors version: 0.8.0
- Accelerate version: 1.15.0
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.1+cpu (NA)
- Using distributed or parallel set-up in script?: No
Who can help?
Reproduction
On main, whole-word masking fails when collating sequences of different lengths. tokenizer.pad pads the token inputs but leaves offset_mapping unpadded.
from transformers import AutoTokenizer, DataCollatorForLanguageModeling
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
examples = [
tokenizer(text, return_offsets_mapping=True, return_special_tokens_mask=True)
for text in ["hello world", "hello"]
]
collator = DataCollatorForLanguageModeling(
tokenizer,
whole_word_mask=True,
mask_replace_prob=1.0,
random_replace_prob=0.0,
)
collator(examples)Traceback
Traceback (most recent call last):
File "/home/xiao/transformers/src/transformers/tokenization_utils_base.py", line 740, in convert_to_tensors
tensor = as_tensor(value)
^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/tokenization_utils_base.py", line 699, in as_tensor
return torch.tensor(value, dtype=dtype)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: expected sequence of length 4 at dim 1 (got 3)
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/xiao/transformers/reproduce_whole_word_padding.py", line 14, in <module>
collator(examples)
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 42, in __call__
return self.torch_call(features)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 776, in torch_call
batch = pad_without_fast_tokenizer_warning(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 63, in pad_without_fast_tokenizer_warning
padded = tokenizer.pad(*pad_args, **pad_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/tokenization_utils_base.py", line 2788, in pad
return BatchEncoding(batch_outputs, tensor_type=return_tensors)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/tokenization_utils_base.py", line 249, in __init__
self.convert_to_tensors(tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis)
File "/home/xiao/transformers/src/transformers/tokenization_utils_base.py", line 756, in convert_to_tensors
raise ValueError(
ValueError: Unable to create tensor, you should probably activate truncation and/or padding with 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your features (`offset_mapping` in this case) have excessive nesting (inputs type `list` where type `int` is expected).The same problem occurs with equal-length inputs when the collator adds padding:
from transformers import AutoTokenizer, DataCollatorForLanguageModeling
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
examples = [
tokenizer(text, return_offsets_mapping=True, return_special_tokens_mask=True)
for text in ["hello world", "world hello"]
]
collator = DataCollatorForLanguageModeling(
tokenizer,
whole_word_mask=True,
mask_replace_prob=1.0,
random_replace_prob=0.0,
pad_to_multiple_of=8,
)
collator(examples)Traceback
Traceback (most recent call last):
File "/home/xiao/transformers/reproduce_whole_word_pad_to_multiple.py", line 15, in <module>
collator(examples)
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 42, in __call__
return self.torch_call(features)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 788, in torch_call
batch["input_ids"], batch["labels"] = self.torch_mask_tokens(
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 815, in torch_mask_tokens
word_ids, no_mask_mask = self._calc_word_ids_and_prob_mask(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xiao/transformers/src/transformers/data/data_collator.py", line 995, in _calc_word_ids_and_prob_mask
is_new_word = (~special_tokens_mask) & ((token_starts != prev_token_ends) | (prev_token_special == 1))
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~
ValueError: operands could not be broadcast together with shapes (2,4) (2,8) Expected behavior
Both examples should produce a padded batch with whole-word masking applied, without errors.
I have a small local fix in tokenizer._pad that pads offset_mapping alongside the token inputs, with regression tests. Would a PR for this fix be welcome?
Related issue: #18681.
Activity
Rocketknight1 commentedon Oct 27, 2026
Yes, this looks legit, and you can open a PR for it! I think
_padis the right placeyuanx749 commentedon Oct 27, 2026
Thanks for confirming. I'll open a PR with the
_padfix and tests in a few days.