fix: sanitize invalid ids before fast decode

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 -----------------------------------------------------
 

python/src/lib.rs

@@ -450,6 +450,12 @@ struct PyTokenizer {
 }
 
 impl PyTokenizer {
+    fn sanitize_decode_ids(ids: Vec<i64>) -> Vec<u32> {
+        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<u32>, skip_special_tokens: bool) -> PyResult<String> {
+    fn decode(&self, ids: Vec<i64>, skip_special_tokens: bool) -> PyResult<String> {
+        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<Vec<u32>>,
+        sentences: Vec<Vec<i64>>,
         skip_special_tokens: bool,
     ) -> PyResult<Vec<String>> {
         let state = self.read();
-        let refs: Vec<&[u32]> = sentences.iter().map(Vec::as_slice).collect();
+        let sanitized: Vec<Vec<u32>> = 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]
+        );
+    }
 }
 
 // ---------------------------------------------------------------------------

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