diff --git a/src/encoding.rs b/src/encoding.rs index 60257e7..2bb34f4 100644 --- a/src/encoding.rs +++ b/src/encoding.rs @@ -1,6 +1,6 @@ use crate::{ chat::{Author, Content, Message, ReasoningEffort, Role, SystemContent, TextContent}, - tiktoken::{CoreBPE, Rank}, + tiktoken::{CoreBPE, EncodeError, Rank}, }; use anyhow::Context as _; use std::{ @@ -32,6 +32,9 @@ pub(crate) enum RenderFormattingTokenError { token: FormattingToken, encoding: Vec, }, + + #[error(transparent)] + Encode(#[from] EncodeError), } /// These are formatting tokens that the renderer can use to generically @@ -339,7 +342,7 @@ impl HarmonyEncoding { let mapped = self .mapped_format_token(t) .ok_or(RenderFormattingTokenError::UnmappedToken(t))?; - let encoded = self.tokenizer.encode_with_special_tokens(mapped); + let encoded = self.tokenizer.encode_with_special_tokens(mapped)?; if encoded.len() != 1 { return Err(RenderFormattingTokenError::InvalidEncoding { token: t, @@ -367,7 +370,7 @@ impl HarmonyEncoding { T: AsRef, B: Extend, { - into.extend(self.tokenizer.encode_ordinary(text.as_ref())); + into.extend(self.tokenizer.encode_ordinary(text.as_ref())?); Ok(()) } diff --git a/src/py_module.rs b/src/py_module.rs index 345a887..cdac157 100644 --- a/src/py_module.rs +++ b/src/py_module.rs @@ -264,7 +264,12 @@ impl PyHarmonyEncoding { }; let allowed_set: std::collections::HashSet<&str> = allowed_vec.iter().map(|s| s.as_str()).collect(); - Ok(self.inner.tokenizer().encode(text, &allowed_set).0) + let (tokens, _) = self + .inner + .tokenizer() + .encode(text, &allowed_set) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(tokens) } /// Return the list of special tokens for this tokenizer. diff --git a/src/tests.rs b/src/tests.rs index 7aba934..1c36192 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -43,6 +43,7 @@ fn test_simple_convo() { load_test_data("../test-data/test_simple_convo.txt").as_str(), &encoding.tokenizer.special_tokens(), ) + .unwrap() .0; let convo = Conversation::from_messages([ Message::from_role_and_content( @@ -101,6 +102,7 @@ fn test_simple_convo_with_effort() { let expected_tokens = encoding .tokenizer .encode(expected_text.as_str(), &encoding.tokenizer.special_tokens()) + .unwrap() .0; let sys = SystemContent::new() .with_model_identity("You are ChatGPT, a large language model trained by OpenAI.") @@ -195,6 +197,7 @@ fn test_reasoning_system_message() { load_test_data("../test-data/test_reasoning_system_message.txt").as_str(), &encoding.tokenizer.special_tokens(), ) + .unwrap() .0; let convo = Conversation::from_messages([ Message::from_role_and_content( @@ -227,6 +230,7 @@ fn test_reasoning_system_message_no_instruction() { .as_str(), &encoding.tokenizer.special_tokens(), ) + .unwrap() .0; let convo = Conversation::from_messages([ Message::from_role_and_content( @@ -261,6 +265,7 @@ fn test_reasoning_system_message_with_dates() { .as_str(), &encoding.tokenizer.special_tokens(), ) + .unwrap() .0; let convo = Conversation::from_messages([ Message::from_role_and_content( @@ -548,6 +553,7 @@ fn test_tool_response_parsing() { let tokens = encoding .tokenizer .encode(&text_tokens, &encoding.tokenizer.special_tokens()) + .unwrap() .0; let expected_message = Message::from_author_and_content( @@ -575,6 +581,7 @@ fn test_encode_decode_roundtrip() { let tokens = encoding .tokenizer .encode(text, &std::collections::HashSet::new()) + .unwrap() .0; assert_eq!(encoding.tokenizer.decode_utf8(&tokens).unwrap(), text); } @@ -584,29 +591,61 @@ fn test_encode_allowed_special() { use std::collections::HashSet; let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let text = "hello world"; - let tokens = encoding.tokenizer.encode(text, &HashSet::new()).0; + let tokens = encoding.tokenizer.encode(text, &HashSet::new()).unwrap().0; assert_eq!(tokens, vec![24912, 2375]); // Allowed special token let mut allowed = HashSet::new(); allowed.insert("<|start|>"); - let tokens = encoding.tokenizer.encode("<|start|>", &allowed).0; + let tokens = encoding.tokenizer.encode("<|start|>", &allowed).unwrap().0; assert_eq!(tokens, vec![200006]); // Allowed special = all allowed = encoding.tokenizer.special_tokens(); // set of all special tokens - let tokens = encoding.tokenizer.encode("<|start|>", &allowed).0; + let tokens = encoding.tokenizer.encode("<|start|>", &allowed).unwrap().0; assert_eq!(tokens, vec![200006]); // Disallowed special (should error) - let result = encoding.tokenizer.encode("<|start|>", &HashSet::new()); + let result = encoding + .tokenizer + .encode("<|start|>", &HashSet::new()) + .unwrap(); assert!( result.0.is_empty() || result.0 != vec![200006], "Expected error or not special token for disallowed special token" ); // Disallowed special = empty (should not treat as special) - let tokens = encoding.tokenizer.encode("<|start|>", &HashSet::new()).0; + let tokens = encoding + .tokenizer + .encode("<|start|>", &HashSet::new()) + .unwrap() + .0; // This may not match the Python fallback, but should not be the special token assert_ne!(tokens, vec![200006]); } +#[test] +fn test_encode_long_single_char_run_returns_err() { + // A pathological long single-character run overflows the fancy_regex + // backtracking stack. This must be returned as a recoverable `Err` rather + // than panicking (which surfaces as an uncatchable PanicException in Python). + use std::collections::HashSet; + let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); + let long_run = "a".repeat(1_300_000); + let result = encoding.tokenizer.encode(&long_run, &HashSet::new()); + assert!( + result.is_err(), + "expected pathological long run to return an error, not panic or succeed" + ); + // The tokenizer must remain usable afterwards. + let tokens = encoding + .tokenizer + .encode("hello world", &HashSet::new()) + .unwrap() + .0; + assert_eq!( + encoding.tokenizer.decode_utf8(&tokens).unwrap(), + "hello world" + ); +} + #[test] fn test_is_special_token() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); @@ -630,6 +669,7 @@ fn test_streamable_parser() { let tokens = encoding .tokenizer .encode(&text, &encoding.tokenizer.special_tokens()) + .unwrap() .0; let mut parser = crate::encoding::StreamableParser::new(encoding.clone(), Some(Role::Assistant)).unwrap(); @@ -656,7 +696,10 @@ fn assert_tokens_eq(tokenizer: &CoreBPE, expected: &[Rank], actual: &[Rank]) { fn test_streamable_parser_tool_call_with_constrain_adjacent() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let text = "<|start|>assistant<|channel|>commentary to=functions.get_weather<|constrain|>json<|message|>{\"latitude\":48.8566,\"longitude\":2.3522}<|call|>"; - let tokens = encoding.tokenizer().encode_with_special_tokens(text); + let tokens = encoding + .tokenizer() + .encode_with_special_tokens(text) + .unwrap(); let mut parser = StreamableParser::new(encoding, None).unwrap(); for token in tokens { let _ = parser.process(token).unwrap(); @@ -678,7 +721,10 @@ fn test_streamable_parser_tool_call_with_constrain_adjacent() { fn test_missing_message_token_requires_non_strict_mode() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let malformed = "<|channel|>commentary Hello<|end|>"; - let tokens = encoding.tokenizer().encode_with_special_tokens(malformed); + let tokens = encoding + .tokenizer() + .encode_with_special_tokens(malformed) + .unwrap(); // Strict mode should continue to error on malformed headers. let strict_result = encoding @@ -708,7 +754,10 @@ fn test_missing_message_token_requires_non_strict_mode() { fn test_tool_call_with_constrain_marker_adjacent() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let text = "<|start|>assistant to=functions.get_weather<|channel|>commentary<|constrain|>json<|message|>{\"location\": \"Tokyo\"}<|end|>"; - let tokens = encoding.tokenizer().encode_with_special_tokens(text); + let tokens = encoding + .tokenizer() + .encode_with_special_tokens(text) + .unwrap(); let parsed = encoding .parse_messages_from_completion_tokens(tokens, None) .expect("expected to parse"); @@ -726,7 +775,10 @@ fn test_tool_call_with_constrain_marker_adjacent() { fn test_tool_call_with_channel_before_recipient_and_constrain_adjacent() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let text = "<|start|>assistant<|channel|>commentary to=functions.get_weather<|constrain|>json<|message|>{\"latitude\":48.8566,\"longitude\":2.3522}<|call|>"; - let tokens = encoding.tokenizer().encode_with_special_tokens(text); + let tokens = encoding + .tokenizer() + .encode_with_special_tokens(text) + .unwrap(); let parsed = encoding .parse_messages_from_completion_tokens(tokens, None) .expect("expected to parse"); @@ -759,7 +811,8 @@ fn test_streamable_parser_does_not_leak_bytes_between_messages() { tokens.extend( encoding .tokenizer() - .encode_with_special_tokens(first_prefix), + .encode_with_special_tokens(first_prefix) + .unwrap(), ); // Two invalid tokens to ensure we end the first message with incomplete UTF-8 bytes. tokens.push(9552); @@ -767,20 +820,28 @@ fn test_streamable_parser_does_not_leak_bytes_between_messages() { tokens.extend( encoding .tokenizer() - .encode_with_special_tokens(first_suffix), + .encode_with_special_tokens(first_suffix) + .unwrap(), ); // Second message should be clean and unaffected. tokens.extend( encoding .tokenizer() - .encode_with_special_tokens(second_prefix), + .encode_with_special_tokens(second_prefix) + .unwrap(), ); - tokens.extend(encoding.tokenizer().encode_with_special_tokens("Hi")); tokens.extend( encoding .tokenizer() - .encode_with_special_tokens(second_suffix), + .encode_with_special_tokens("Hi") + .unwrap(), + ); + tokens.extend( + encoding + .tokenizer() + .encode_with_special_tokens(second_suffix) + .unwrap(), ); let mut parser = StreamableParser::new(encoding, None).unwrap(); @@ -806,9 +867,15 @@ fn test_streamable_parser_flushes_partial_bytes_on_eos() { let mut tokens = encoding .tokenizer() - .encode_with_special_tokens("<|start|>assistant<|message|>"); + .encode_with_special_tokens("<|start|>assistant<|message|>") + .unwrap(); tokens.push(9552); - tokens.extend(encoding.tokenizer().encode_with_special_tokens("Hi")); + tokens.extend( + encoding + .tokenizer() + .encode_with_special_tokens("Hi") + .unwrap(), + ); let mut parser = StreamableParser::new(encoding.clone(), None).unwrap(); for token in tokens { @@ -833,12 +900,17 @@ fn test_streamable_parser_waits_for_multi_token_utf8_sequence() { let start_tokens = encoding .tokenizer() - .encode_with_special_tokens("<|start|>assistant<|message|>"); + .encode_with_special_tokens("<|start|>assistant<|message|>") + .unwrap(); for token in &start_tokens { parser.process(*token).unwrap(); } - let emoji_tokens = encoding.tokenizer().encode("💖", &HashSet::new()).0; + let emoji_tokens = encoding + .tokenizer() + .encode("💖", &HashSet::new()) + .unwrap() + .0; assert!( emoji_tokens.len() >= 2, "expected multi-token emoji encoding" @@ -852,7 +924,10 @@ fn test_streamable_parser_waits_for_multi_token_utf8_sequence() { assert_eq!(parser.last_content_delta().unwrap(), Some("💖".to_string())); assert_eq!(parser.current_content().unwrap(), "💖"); - let end_tokens = encoding.tokenizer().encode_with_special_tokens("<|end|>"); + let end_tokens = encoding + .tokenizer() + .encode_with_special_tokens("<|end|>") + .unwrap(); for token in end_tokens { parser.process(token).unwrap(); } @@ -868,9 +943,12 @@ fn test_parse_completion_with_invalid_content_token_errors_on_eos() { let encoding = load_harmony_encoding(HarmonyEncodingName::HarmonyGptOss).unwrap(); let mut parser = StreamableParser::new(encoding.clone(), None).unwrap(); - let start_tokens = encoding.tokenizer().encode_with_special_tokens( - "<|start|>assistant<|channel|>analysis<|message|>Practice invalid token handling.", - ); + let start_tokens = encoding + .tokenizer() + .encode_with_special_tokens( + "<|start|>assistant<|channel|>analysis<|message|>Practice invalid token handling.", + ) + .unwrap(); for token in &start_tokens { parser.process(*token).unwrap(); } diff --git a/src/tiktoken.rs b/src/tiktoken.rs index 9c5a290..0291c47 100644 --- a/src/tiktoken.rs +++ b/src/tiktoken.rs @@ -158,6 +158,27 @@ impl std::fmt::Display for DecodeError { impl std::error::Error for DecodeError {} +#[derive(Debug, Clone)] +pub struct EncodeError { + pub message: String, +} + +impl std::fmt::Display for EncodeError { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "Could not encode text: {}", self.message) + } +} + +impl std::error::Error for EncodeError {} + +impl From for EncodeError { + fn from(e: fancy_regex::Error) -> Self { + EncodeError { + message: e.to_string(), + } + } +} + const MAX_NUM_THREADS: usize = 128; #[derive(Clone)] @@ -218,22 +239,26 @@ impl CoreBPE { }) } - pub fn encode_ordinary(&self, text: &str) -> Vec { + pub fn encode_ordinary(&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 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(&self, text: &str, allowed_special: &HashSet<&str>) -> (Vec, usize) { + pub fn encode( + &self, + text: &str, + allowed_special: &HashSet<&str>, + ) -> Result<(Vec, usize), EncodeError> { let special_regex = self._get_tl_special_regex(); let regex = self._get_tl_regex(); let mut ret = vec![]; @@ -245,7 +270,7 @@ 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)?; match next_special { Some(m) => { if allowed_special.contains(&text[m.start()..m.end()]) { @@ -260,7 +285,7 @@ impl CoreBPE { // Okay, here we go, compare this logic to encode_ordinary for mat in regex.find_iter(&text[start..end]) { - let piece = mat.unwrap().as_str().as_bytes(); + let piece = mat?.as_str().as_bytes(); if let Some(token) = self.encoder.get(piece) { last_piece_token_len = 1; ret.push(*token); @@ -286,7 +311,7 @@ impl CoreBPE { // last_piece_token_len is how many tokens came from the last regex split. This is used // for determining unstable tokens, since you can't merge across (stable) regex splits - (ret, last_piece_token_len) + Ok((ret, last_piece_token_len)) } fn _increase_last_piece_token_len( @@ -332,12 +357,12 @@ impl CoreBPE { &self, text: &str, allowed_special: &HashSet<&str>, - ) -> (Vec, HashSet>) { - let (tokens, last_piece_token_len) = self.encode(text, allowed_special); + ) -> 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); @@ -353,7 +378,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 @@ -392,7 +417,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(s)?, // Technically, whether or not this arm is correct depends on whether there // would be a regex split before the UTF-8 truncation point. @@ -443,7 +468,7 @@ impl CoreBPE { } } - (tokens, completions) + Ok((tokens, completions)) } pub fn new( @@ -514,9 +539,9 @@ impl CoreBPE { .collect() } - pub fn encode_with_special_tokens(&self, text: &str) -> Vec { + pub fn encode_with_special_tokens(&self, text: &str) -> Result, EncodeError> { let allowed_special = self.special_tokens(); - self.encode(text, &allowed_special).0 + Ok(self.encode(text, &allowed_special)?.0) } pub fn is_special_token(&self, token: Rank) -> bool { diff --git a/src/wasm_module.rs b/src/wasm_module.rs index 0cbe281..a7daba7 100644 --- a/src/wasm_module.rs +++ b/src/wasm_module.rs @@ -214,7 +214,12 @@ impl JsHarmonyEncoding { }; let allowed_set: std::collections::HashSet<&str> = allowed_vec.iter().map(|s| s.as_str()).collect(); - Ok(self.inner.tokenizer().encode(text, &allowed_set).0) + let (tokens, _) = self + .inner + .tokenizer() + .encode(text, &allowed_set) + .map_err(|e| JsValue::from_str(&e.to_string()))?; + Ok(tokens) } #[wasm_bindgen(js_name = specialTokens)] diff --git a/tests/test_harmony.py b/tests/test_harmony.py index dbb9925..7f03e0b 100644 --- a/tests/test_harmony.py +++ b/tests/test_harmony.py @@ -1244,3 +1244,25 @@ def test_streamable_parser_tricky_utf8_decoding(): # Ensure if we're accumulating content deltas we still get the full utf-8 text assert "".join(content_deltas) == tricky_utf8_text + + +def test_encode_long_single_char_run_raises_catchable_error(): + """A pathological long single-character run makes the fancy_regex backtracking + engine exceed its stack limit. This must surface as a normal, catchable Python + exception rather than an uncatchable ``pyo3_runtime.PanicException`` that would + abort a host process. + """ + encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) + long_run = "a" * 1_300_000 + + # The exception must be catchable via ``except Exception`` (PanicException + # subclasses BaseException and would slip past such a handler). + with pytest.raises(Exception) as exc_info: + encoding.encode(long_run, allowed_special="all") + + assert isinstance(exc_info.value, HarmonyError) + assert type(exc_info.value).__name__ != "PanicException" + + # Ensure the tokenizer is still usable afterwards and ordinary text round-trips. + tokens = encoding.encode("hello world", allowed_special="all") + assert encoding.decode_utf8(tokens) == "hello world"