diff --git a/python/fastokens/_compat.py b/python/fastokens/_compat.py index 7c70fb4..8157f56 100644 --- a/python/fastokens/_compat.py +++ b/python/fastokens/_compat.py @@ -15,6 +15,7 @@ from fastokens._native import Encoding, Tokenizer + # Backwards-compatibility alias used by tests and any code that imports # _Encoding directly from this module. _Encoding = Encoding @@ -300,7 +301,9 @@ def decode_batch( sequences: list[list[int]], skip_special_tokens: bool = True, ) -> list[str]: - return self._fast.decode_batch(sequences, skip_special_tokens=skip_special_tokens) + return self._fast.decode_batch( + sequences, skip_special_tokens=skip_special_tokens + ) # -- Vocabulary ----------------------------------------------------- diff --git a/python/src/lib.rs b/python/src/lib.rs index e68852c..2843b3a 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -450,6 +450,12 @@ struct PyTokenizer { } impl PyTokenizer { + fn sanitize_decode_ids(ids: Vec) -> Vec { + ids.into_iter() + .filter_map(|id| (0..=u32::MAX as i64).contains(&id).then_some(id as u32)) + .collect() + } + fn read(&self) -> std::sync::RwLockReadGuard<'_, TokenizerState> { self.state.read().expect("PyTokenizer state lock poisoned") } @@ -1041,7 +1047,8 @@ impl PyTokenizer { /// Decode token IDs back into text. #[pyo3(signature = (ids, skip_special_tokens = false))] - fn decode(&self, ids: Vec, skip_special_tokens: bool) -> PyResult { + fn decode(&self, ids: Vec, skip_special_tokens: bool) -> PyResult { + let ids = Self::sanitize_decode_ids(ids); self.read() .inner .decode(&ids, skip_special_tokens) @@ -1052,11 +1059,15 @@ impl PyTokenizer { #[pyo3(signature = (sentences, skip_special_tokens = false))] fn decode_batch( &self, - sentences: Vec>, + sentences: Vec>, skip_special_tokens: bool, ) -> PyResult> { let state = self.read(); - let refs: Vec<&[u32]> = sentences.iter().map(Vec::as_slice).collect(); + let sanitized: Vec> = sentences + .into_iter() + .map(Self::sanitize_decode_ids) + .collect(); + let refs: Vec<&[u32]> = sanitized.iter().map(Vec::as_slice).collect(); state .inner .decode_batch(&refs, skip_special_tokens) @@ -1142,6 +1153,15 @@ mod tests { assert_eq!(enc.attention_mask, vec![0u32, 0, 1, 1, 1]); assert_eq!(enc.type_ids, vec![7u32, 7, 0, 0, 0]); } + + #[test] + fn sanitize_decode_ids_filters_out_of_range_values() { + let ids = vec![1_i64, -1, 2, i64::from(u32::MAX) + 1, 3]; + assert_eq!( + PyTokenizer::sanitize_decode_ids(ids), + vec![1_u32, 2, 3] + ); + } } // --------------------------------------------------------------------------- diff --git a/python/tests/test_patch_transformers.py b/python/tests/test_patch_transformers.py index dbb6ebe..0bcc906 100644 --- a/python/tests/test_patch_transformers.py +++ b/python/tests/test_patch_transformers.py @@ -44,6 +44,27 @@ def test_encode_decode_through_shim(): assert "Hello" in decoded, f"unexpected decode: {decoded!r}" +def test_decode_ignores_ids_outside_uint32(): + """Fast decode should drop invalid IDs instead of raising.""" + import fastokens + + fastokens.patch_transformers() + + tok = transformers.AutoTokenizer.from_pretrained(MODEL) + valid = tok("Hello, world!")["input_ids"] + mixed = valid[:2] + [-1, 2**32] + valid[2:] + + assert tok.decode(mixed, skip_special_tokens=True) == tok.decode( + valid, skip_special_tokens=True + ) + + batch_valid = [valid, valid] + batch_mixed = [mixed, valid[:1] + [-1] + valid[1:]] + assert tok.batch_decode(batch_mixed, skip_special_tokens=True) == tok.batch_decode( + batch_valid, skip_special_tokens=True + ) + + def test_unpatch_restores_backend(): """After unpatching, from_pretrained should return the original backend.""" import fastokens