Fix out-of-bounds error
Merged
AP
5mo ago+11+11+01234567890123456789-1-1-0123456789
This pull request introduces important safety checks to prevent out-of-bounds access in the logits processing code for transformers. The main updates ensure that token IDs are validated against the vocabulary size before being used, which improves robustness and prevents potential crashes.
Safety and robustness improvements:
- Added bounds checking for token IDs in the
RepetitionPenaltyLogitsProcessor<T>::Processmethod to ensure only valid IDs are used when accessingbeam_token_scores. - Added bounds checking for token IDs in the
NoRepeatNGramLogitsProcessor<T>::Processmethod to prevent out-of-bounds writes tobeam_token_scores. - Updated the
NextTokenScores::SetScoremethod to return early if the providedtoken_idis out of bounds, replacing the previous assert with a safe check.