diff --git a/src/lib.rs b/src/lib.rs index ea54eac8..f088d9e1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -358,18 +358,32 @@ impl CoreBPE { } pub fn encode_ordinary(&self, text: &str) -> Vec { + self.encode_ordinary_native(text) + .expect("Regex error while tokenizing") + } + + /// Fallible core of [`encode_ordinary`](Self::encode_ordinary). + /// + /// The regex engine (`fancy_regex`) can return a recoverable `Err` + /// (e.g. `RuntimeError::StackOverflow`) on pathological input instead of + /// blowing the native stack. This method surfaces that as an + /// [`EncodeError`] rather than panicking. + fn encode_ordinary_native(&self, text: &str) -> Result, EncodeError> { // This is the core of the encoding logic; the other functions in here // just make things complicated :-) let regex = self._get_tl_regex(); let mut ret = vec![]; for mat in regex.find_iter(text) { - let piece = mat.unwrap().as_str().as_bytes(); + let mat = mat.map_err(|e| EncodeError { + message: format!("Regex error while tokenizing: {e}"), + })?; + let piece = mat.as_str().as_bytes(); match self.encoder.get(piece) { Some(token) => ret.push(*token), None => ret.extend(&byte_pair_encode(piece, &self.encoder)), } } - ret + Ok(ret) } pub fn encode( @@ -388,7 +402,12 @@ impl CoreBPE { let mut start_find = start; loop { // Find the next allowed special token, if any - next_special = special_regex.find_from_pos(text, start_find).unwrap(); + next_special = + special_regex + .find_from_pos(text, start_find) + .map_err(|e| EncodeError { + message: format!("Regex error while tokenizing: {e}"), + })?; match next_special { Some(m) => { if allowed_special.contains(&text[m.start()..m.end()]) { @@ -485,11 +504,24 @@ impl CoreBPE { text: &str, allowed_special: &HashSet<&str>, ) -> (Vec, HashSet>) { - let (tokens, last_piece_token_len) = self.encode(text, allowed_special).unwrap(); + self._encode_unstable_native_impl(text, allowed_special) + .expect("Regex error while tokenizing") + } + + /// Fallible core of [`_encode_unstable_native`](Self::_encode_unstable_native). + /// + /// Delegates to [`encode`](Self::encode), which surfaces recoverable + /// `fancy_regex` runtime errors as [`EncodeError`] instead of panicking. + fn _encode_unstable_native_impl( + &self, + text: &str, + allowed_special: &HashSet<&str>, + ) -> Result<(Vec, HashSet>), EncodeError> { + let (tokens, last_piece_token_len) = self.encode(text, allowed_special)?; if last_piece_token_len == 0 { // If last_piece_token_len is zero, the last token was a special token and we have // no unstable bytes - return (tokens, HashSet::new()); + return Ok((tokens, HashSet::new())); } let (mut tokens, last_piece_token_len) = self._increase_last_piece_token_len(tokens, last_piece_token_len); @@ -505,7 +537,7 @@ impl CoreBPE { let mut completions = HashSet::new(); if unstable_bytes.is_empty() { - return (tokens, completions); + return Ok((tokens, completions)); } // This is the easy bit. Just find all single tokens that start with unstable_bytes @@ -544,7 +576,7 @@ impl CoreBPE { // So convert to UTF-8 and do regex splitting. // E.g. with cl100k_base " !" gets split to " " + " !", // but byte_pair_encode(" !") != byte_pair_encode(" ") - Ok(s) => self.encode_ordinary(s), + Ok(s) => self.encode_ordinary_native(s)?, // Technically, whether or not this arm is correct depends on whether there // would be a regex split before the UTF-8 truncation point. @@ -595,7 +627,7 @@ impl CoreBPE { } } - (tokens, completions) + Ok((tokens, completions)) } pub fn new( @@ -671,16 +703,19 @@ impl CoreBPE { pub fn encode_with_special_tokens(&self, text: &str) -> Vec { let allowed_special = self.special_tokens(); - self.encode(text, &allowed_special).unwrap().0 + self.encode(text, &allowed_special) + .expect("Regex error while tokenizing") + .0 } } #[cfg(test)] mod tests { - use fancy_regex::Regex; + use std::collections::HashSet; + use rustc_hash::FxHashMap as HashMap; - use crate::{Rank, byte_pair_split}; + use crate::{CoreBPE, Rank, byte_pair_split}; fn setup_ranks() -> HashMap, Rank> { HashMap::from_iter([(b"ab".to_vec(), 0), (b"cd".to_vec(), 1)]) @@ -699,4 +734,88 @@ mod tests { let res = byte_pair_split(b"abab", &ranks); assert_eq!(res, vec![b"ab", b"ab"]); } + + // The cl100k_base regex (as used by e.g. gpt-4), whose `\s+(?!\S)` alternative + // can drive fancy_regex's backtracking engine to return a recoverable + // `RuntimeError::StackOverflow` on large whitespace-heavy input. See + // https://github.com/openai/tiktoken/issues/400. + const CL100K_PAT: &str = r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s"; + + // Build a byte-level CoreBPE (one token per possible byte) with the given + // pattern. This is enough to exercise the encoding paths without downloading + // a real vocabulary. + fn byte_level_bpe(pattern: &str) -> CoreBPE { + let encoder: HashMap, Rank> = + (0u16..256).map(|b| (vec![b as u8], b as Rank)).collect(); + CoreBPE::new_internal(encoder, HashMap::default(), pattern) + .expect("failed to build test CoreBPE") + } + + // Input that deterministically triggers fancy_regex's + // `RuntimeError::StackOverflow` with the cl100k_base pattern: a long run of + // spaces followed by a non-space so the engine must attempt (and backtrack + // out of) the `\s+(?!\S)` alternative. The threshold is not empirical: it is + // fancy_regex 0.17.0's hard-coded `MAX_STACK = 1_000_000` backtracking-step + // counter (see `vm.rs`), so 2M chars is a safe, deterministic margin over it. + // This runs in a few tens of milliseconds in a release build; under a plain + // debug `cargo test` it is closer to a second. + fn pathological_input() -> String { + let mut s = " ".repeat(2_000_000); + s.push('x'); + s + } + + // Characterization test (not a regression for this change): `encode` already + // returns `Err(EncodeError)` on unmodified `main` — its inner `find_iter` + // error handling was added upstream in eedc856. This just pins that behavior + // and the shared error message the sibling paths now reuse. + #[test] + fn test_encode_returns_err_on_pathological_input() { + let bpe = byte_level_bpe(CL100K_PAT); + let text = pathological_input(); + let res = bpe.encode(&text, &HashSet::new()); + assert!( + res.is_err(), + "expected encode to surface a recoverable regex error, got Ok" + ); + let msg = res.unwrap_err().message; + assert!( + msg.contains("Regex error while tokenizing"), + "unexpected error message: {msg}" + ); + } + + // The actual regression this change fixes: `encode_ordinary` (and the other + // sibling paths) panicked on `main` because they never got eedc856's + // treatment. Its fallible core must now surface the error instead. The public + // `encode_ordinary` delegates here and only then `.expect()`s, preserving its + // infallible signature for Rust callers. + #[test] + fn test_encode_ordinary_native_returns_err_on_pathological_input() { + let bpe = byte_level_bpe(CL100K_PAT); + let text = pathological_input(); + let res = bpe.encode_ordinary_native(&text); + assert!( + res.is_err(), + "expected encode_ordinary_native to surface a recoverable regex error, got Ok" + ); + } + + // Sanity: on ordinary input the fixed paths still succeed and agree with each + // other, so the error handling didn't change normal behavior. + #[test] + fn test_encode_paths_agree_on_normal_input() { + let bpe = byte_level_bpe(CL100K_PAT); + let text = "hello world, this is a test 123"; + let ordinary = bpe.encode_ordinary(text); + let ordinary_native = bpe + .encode_ordinary_native(text) + .expect("normal input should encode without error"); + let (via_encode, _) = bpe + .encode(text, &HashSet::new()) + .expect("normal input should encode without error"); + assert_eq!(ordinary, ordinary_native); + assert_eq!(ordinary, via_encode); + assert!(!ordinary.is_empty()); + } } diff --git a/src/py.rs b/src/py.rs index 60b62452..ccf79bcc 100644 --- a/src/py.rs +++ b/src/py.rs @@ -27,8 +27,11 @@ impl CoreBPE { // ==================== #[pyo3(name = "encode_ordinary")] - fn py_encode_ordinary(&self, py: Python, text: &str) -> Vec { - py.detach(|| self.encode_ordinary(text)) + fn py_encode_ordinary(&self, py: Python, text: &str) -> PyResult> { + py.detach(|| { + self.encode_ordinary_native(text) + .map_err(|e| PyErr::new::(e.message)) + }) } #[pyo3(name = "encode")] @@ -69,17 +72,20 @@ impl CoreBPE { buffer.into_py_any(py) } - fn _encode_bytes(&self, py: Python, bytes: &[u8]) -> Vec { + fn _encode_bytes(&self, py: Python, bytes: &[u8]) -> PyResult> { py.detach(|| { match std::str::from_utf8(bytes) { // Straightforward case - Ok(text) => self.encode_ordinary(text), + Ok(text) => self + .encode_ordinary_native(text) + .map_err(|e| PyErr::new::(e.message)), // Oops, don't actually have UTF-8. But we need to do the regex splitting in // Unicode space, so we make our best guess at where we would have splits Err(e) => { let text = unsafe { std::str::from_utf8_unchecked(&bytes[..e.valid_up_to()]) }; - let (tokens, last_piece_token_len) = - self.encode(text, &HashSet::new()).unwrap(); + let (tokens, last_piece_token_len) = self + .encode(text, &HashSet::new()) + .map_err(|e| PyErr::new::(e.message))?; let (mut tokens, last_piece_token_len) = self._increase_last_piece_token_len(tokens, last_piece_token_len); @@ -108,7 +114,7 @@ impl CoreBPE { } } } - tokens + Ok(tokens) } } }) @@ -124,8 +130,9 @@ impl CoreBPE { let (tokens, completions): (Vec, HashSet>) = py.detach(|| { let allowed_special: HashSet<&str> = allowed_special.iter().map(|s| s.as_ref()).collect(); - self._encode_unstable_native(text, &allowed_special) - }); + self._encode_unstable_native_impl(text, &allowed_special) + .map_err(|e| PyErr::new::(e.message)) + })?; let py_completions = PyList::new(py, completions.into_iter())?; Ok((tokens, py_completions.into())) } diff --git a/tests/test_encoding.py b/tests/test_encoding.py index b77ca135..e67d8b58 100644 --- a/tests/test_encoding.py +++ b/tests/test_encoding.py @@ -124,6 +124,32 @@ def test_catastrophically_repetitive(make_enc: Callable[[], tiktoken.Encoding]): assert big_value == enc.decode(enc.encode(big_value)) +def test_pathological_regex_input_raises(): + # A long run of spaces followed by a non-space drives fancy_regex's + # backtracking engine (via cl100k_base's `\s+(?!\S)` alternative) past its + # hard-coded MAX_STACK = 1_000_000 counter, which it reports as a recoverable + # runtime error. tiktoken must surface that as a catchable ValueError rather + # than letting the panic escape as a PanicException (which is not catchable + # as a normal `except Exception`). See + # https://github.com/openai/tiktoken/issues/400. + # + # `encode()` itself already handles this upstream; the regression here is for + # the sibling paths that previously called `.unwrap()`. 2M chars is a + # comfortable margin over the 1M threshold and still returns in tens of + # milliseconds in the (always-release) extension build. + enc = tiktoken.get_encoding("cl100k_base") + pathological = " " * 2_000_000 + "x" + + with pytest.raises(ValueError): + enc.encode_ordinary(pathological) + + with pytest.raises(ValueError): + enc._encode_bytes(pathological.encode()) + + with pytest.raises(ValueError): + enc.encode_with_unstable(pathological) + + # ==================== # Roundtrip # ====================