diff --git a/tokenizers/src/decoders/strip.rs b/tokenizers/src/decoders/strip.rs index 9aeffec647..099787acd8 100644 --- a/tokenizers/src/decoders/strip.rs +++ b/tokenizers/src/decoders/strip.rs @@ -42,7 +42,11 @@ impl Decoder for Strip { } let mut stop_cut = chars.len(); - for i in 0..self.stop { + // Bound the trailing strip by the token length: `self.stop` is a raw + // config value, so without `.min(chars.len())` a token made entirely + // of `content` and shorter than `stop` keeps decrementing `index` past + // 0, underflowing `usize` (debug panic / out-of-bounds index in release). + for i in 0..self.stop.min(chars.len()) { let index = chars.len() - i - 1; if chars[index] == self.content { stop_cut = index; @@ -52,6 +56,11 @@ impl Decoder for Strip { } } + // The leading and trailing windows can overlap when a short token is + // stripped from both ends (start_cut > stop_cut), which would make the + // slice range reversed and panic. Clamp so an over-stripped token + // collapses to an empty string instead. + let stop_cut = stop_cut.max(start_cut); let new_token: String = chars[start_cut..stop_cut].iter().collect(); new_token }) @@ -77,4 +86,30 @@ mod tests { .unwrap(); assert_eq!(res, vec!["He", " friend!"]); } + + #[test] + fn short_all_content_token_does_not_panic() { + // A token made entirely of `content` and shorter than `stop` used to + // underflow `chars.len() - i - 1` (debug: subtract overflow; release: + // out-of-bounds index). It should strip what is there and stop. + let decoder = Strip::new('_', 0, 2); + let res = decoder.decode_chain(vec!["_".into()]).unwrap(); + assert_eq!(res, vec![""]); + + let decoder = Strip::new('a', 0, 5); + let res = decoder + .decode_chain(vec!["aa".into(), "aab".into()]) + .unwrap(); + assert_eq!(res, vec!["", "aab"]); + } + + #[test] + fn overlapping_start_stop_windows_collapse_to_empty() { + // When the leading and trailing windows overlap on a short token, + // start_cut can exceed stop_cut; the slice used to panic on a reversed + // range and now collapses to an empty string. + let decoder = Strip::new('H', 2, 1); + let res = decoder.decode_chain(vec!["HH".into()]).unwrap(); + assert_eq!(res, vec![""]); + } }