Repository navigation
fix(utils): handle single-token attention mask in random_mask_tensor (#2525) - #2552
YuvrajSinghBhadoria2 wants to merge 2 commits into
Conversation
…uggingface#2525) random_mask_tensor() called random.randint(1, shape[-1] - 1) which raises ValueError when shape[-1] == 1 (single-token input) because it becomes random.randint(1, 0). Fix: return an all-ones mask when shape[-1] <= 1, since no padding is possible for a single token and it must remain unmasked. Add regression test test_single_token_attention_mask covering both 'pt' and 'np' frameworks with sequence_length=1. Fixes huggingface#2525
The single-token attention mask fix in optimum/utils/input_generators.py and its test in tests/utils/test_dummpy_input_generators.py belong to huggingface#2552, which has its own PR. Keeping them out of huggingface#2553 so both can be reviewed and merged independently, per review feedback on huggingface#2553.
mangeshraut712
left a comment
mangeshraut712
left a comment
There was a problem hiding this comment.
The fix is correct and minimal: with one position there's nothing to pad, so all-ones is the only valid mask. The early return mirrors the existing pt/np branches, so dtypes stay consistent.
| shape = tuple(shape) | ||
| # With a single token (shape[-1] == 1), no padding is possible: | ||
| # we must keep the token unmasked, so the mask is all ones. | ||
| if shape[-1] <= 1: |
There was a problem hiding this comment.
There's one more site with the same randint(1, n - 1) root cause that sequence_length=1 would hit, in BartDummyTextInputGenerator.generate (~L1850):
random_idx = random.randint(1, self.sequence_length - 1)
int_tensor[idx][random_idx] = self.eos_token_idThat runs when force_eos_token_id_presence is set, task == "text-classification", and the random ids don't already contain EOS. With a single token it becomes randint(1, 0) → the same ValueError. Something like random_idx = random.randint(1, self.sequence_length - 1) if self.sequence_length > 1 else 0 would cover it. You could also extend the new test with a BART config to pin both paths.
…ace#2525) Address review feedback on PR huggingface#2552: BartDummyTextInputGenerator.generate() called random.randint(1, self.sequence_length - 1) when forcing EOS token presence, which raises ValueError for sequence_length=1 (randint(1, 0)), the same root cause as random_mask_tensor. Fix: fall back to index 0 when sequence_length <= 1, placing the EOS token at the only valid position. Extend test_single_token_attention_mask with a BART config (pt and np) to pin both the random_mask_tensor and Bart EOS-forcing single-token paths.
|
Thanks for the review — the follow-up is done:
The original |
|
Hi — pinging on this PR. Happy to rebase onto latest main, add benchmarks, or split into smaller chunks if that helps unblock review. Let me know what's needed. |
Fixes #2525
random_mask_tensor() called random.randint(1, shape[-1] - 1) which
raises ValueError when shape[-1] == 1 (single-token input) because it
becomes random.randint(1, 0).
Fix: return an all-ones mask when shape[-1] <= 1, since no padding is
possible for a single token and it must remain unmasked.
Add regression test test_single_token_attention_mask covering both
'pt' and 'np' frameworks with sequence_length=1.
All tests verified passing in a CPU-only Docker container (no GPU
required).