From aa760d543568037cbc1338cf08b0d6078c25876a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 18:44:14 -0700 Subject: [PATCH] fix(tokenizer): reuse packaged vocabularies in the native wheel --- .github/scripts/verify_linux_native_wheel.py | 12 +- litellm-rust/Cargo.lock | 3 + .../crates/python-bridge/src/token_counter.rs | 3 +- .../crates/python-bridge/src/tokenizer.rs | 17 +- .../crates/token-counter-tiktoken/Cargo.toml | 3 + .../crates/token-counter-tiktoken/src/lib.rs | 10 + .../token-counter-tiktoken/src/ranks.rs | 235 ++++++++++++++++++ .../crates/token-counter/src/tiktoken.rs | 11 +- .../litellm_core_utils/test_tokenizer.py | 26 +- .../test_verify_linux_native_wheel.py | 17 +- 10 files changed, 327 insertions(+), 10 deletions(-) create mode 100644 litellm-rust/crates/token-counter-tiktoken/src/ranks.rs diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index d86ad7b15c6..91c7f6645be 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -134,7 +134,16 @@ def main( uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members) native_path: Final = wheel.parent / "native" / Path(native_member.filename).name native_path.parent.mkdir(parents=True, exist_ok=True) - native_path.write_bytes(archive.read(native_member)) + native_bytes: Final = archive.read(native_member) + native_path.write_bytes(native_bytes) + duplicated_vocabularies: Final = tuple( + member.filename + for member in wheel_members + if member.filename.startswith("litellm/litellm_core_utils/tokenizers/") + and re.fullmatch(r"[0-9a-f]{40}", PurePosixPath(member.filename).name) + and member.file_size > 0 + and archive.read(member) in native_bytes + ) wheel_metadata_tags_match: Final = ( len(wheel_metadata_tags) == len(expanded_filename_tags) @@ -223,6 +232,7 @@ def main( ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), ("Native extension does not exceed 35 MB", native_size_within_limit), + ("Tokenizer vocabularies are not duplicated in the native extension", not duplicated_vocabularies), ("Wheel contents are valid", not unexpected_members), ) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c886fd34711..cccc7329e8f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2288,6 +2288,9 @@ dependencies = [ name = "litellm-token-counter-tiktoken" version = "0.1.0" dependencies = [ + "base64 0.22.1", + "once_cell", + "rustc-hash", "thiserror 2.0.19", "tiktoken-rs", ] diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 9f0af284104..16bd9a73142 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -84,7 +84,8 @@ impl TokenCounter { fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { #[cfg(feature = "tiktoken")] { - Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) + let tokenizer = crate::tokenizer::load_tiktoken(py, encoding)?; + Self::load(py, || Ok(CoreTokenCounter::new(tokenizer))) } #[cfg(not(feature = "tiktoken"))] { diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index 92a21f7e0e1..4f8b270fcf2 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -30,6 +30,20 @@ use litellm_token_counter::huggingface::{ #[cfg(feature = "tiktoken")] use litellm_token_counter::tiktoken::TiktokenTokenizer; +#[cfg(feature = "tiktoken")] +pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + let resource: std::path::PathBuf = + PyModule::import(py, "litellm.litellm_core_utils.tokenizers")? + .getattr("__file__")? + .extract()?; + release_gil(py, || { + TiktokenTokenizer::from_cached_ranks(encoding, |file| { + std::fs::read_to_string(resource.with_file_name(file)) + }) + }) + .map_err(|error| token_count_error_to_pyerr(error.into())) +} + enum Codec { #[cfg(feature = "tiktoken")] Tiktoken(TiktokenTokenizer), @@ -59,8 +73,7 @@ impl Tokenizer { fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { #[cfg(feature = "tiktoken")] { - let tokenizer = release_gil(py, || TiktokenTokenizer::from_name(encoding)) - .map_err(|error| token_count_error_to_pyerr(error.into()))?; + let tokenizer = load_tiktoken(py, encoding)?; Ok(Self { inner: Arc::new(Codec::Tiktoken(tokenizer)), }) diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml index 494a9233e69..2fb3103e0c8 100644 --- a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -6,5 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true +once_cell = "1.21.3" +rustc-hash = "2.1.3" thiserror.workspace = true tiktoken-rs.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index c96591abf86..8f7404f1036 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -1,10 +1,12 @@ #![forbid(unsafe_code)] mod error; +mod ranks; use std::collections::HashSet; pub use error::UnsupportedTokenizer; +pub use ranks::LoadError; pub struct TiktokenTokenizer { encoder: &'static tiktoken_rs::CoreBPE, @@ -12,6 +14,14 @@ pub struct TiktokenTokenizer { } impl TiktokenTokenizer { + pub fn from_cached_ranks( + name: &str, + load: impl FnOnce(&str) -> std::io::Result, + ) -> Result { + let (encoder, name) = ranks::load(name, load)?; + Ok(Self { encoder, name }) + } + pub fn from_name(name: &str) -> Result { let (encoder, canonical_name) = match name { "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"), diff --git a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs new file mode 100644 index 00000000000..81739f0f945 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs @@ -0,0 +1,235 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use once_cell::sync::OnceCell; +use rustc_hash::FxHashMap; +use thiserror::Error; +use tiktoken_rs::{CoreBPE, O200K_BASE_PAT_STR, Rank}; + +use crate::UnsupportedTokenizer; + +const CL100K: &str = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"; +const O200K: &str = "fb374d419588a4632f3f557e76b4b70aebbca790"; +const P50K: &str = "ec7223a39ce59f226a68acc30dc1af2788490e15"; +const LEGACY_PATTERN: &str = + r"'(?:[sdmt]|ll|ve|re)| ?\p{L}++| ?\p{N}++| ?[^\s\p{L}\p{N}]++|\s++$|\s+(?!\S)|\s"; +const CL100K_PATTERN: &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"; + +static CL100K_ENCODER: OnceCell = OnceCell::new(); +static O200K_ENCODER: OnceCell = OnceCell::new(); +static HARMONY_ENCODER: OnceCell = OnceCell::new(); +static P50K_ENCODER: OnceCell = OnceCell::new(); +static EDIT_ENCODER: OnceCell = OnceCell::new(); +static R50K_ENCODER: OnceCell = OnceCell::new(); + +#[derive(Debug, Error)] +pub enum LoadError { + #[error(transparent)] + Unsupported(#[from] UnsupportedTokenizer), + #[error("failed to load tiktoken ranks: {0}")] + Ranks(String), +} + +pub(super) fn load( + name: &str, + load_file: impl FnOnce(&str) -> std::io::Result, +) -> Result<(&'static CoreBPE, &'static str), LoadError> { + let (name, file, cache) = match name { + "cl100k_base" => ("cl100k_base", CL100K, &CL100K_ENCODER), + "o200k_base" => ("o200k_base", O200K, &O200K_ENCODER), + "o200k_harmony" => ("o200k_harmony", O200K, &HARMONY_ENCODER), + "p50k_base" => ("p50k_base", P50K, &P50K_ENCODER), + "p50k_edit" => ("p50k_edit", P50K, &EDIT_ENCODER), + "r50k_base" | "gpt2" => ("r50k_base", P50K, &R50K_ENCODER), + _ => return Err(UnsupportedTokenizer(name.to_owned()).into()), + }; + let encoder = cache.get_or_try_init(|| { + let ranks = load_file(file).map_err(|error| LoadError::Ranks(error.to_string()))?; + build(name, &ranks) + })?; + Ok((encoder, name)) +} + +fn build(name: &str, ranks: &str) -> Result { + let parsed = ranks + .lines() + .map(parse_rank) + .collect::, _>>()?; + let encoder: FxHashMap<_, _> = parsed + .into_iter() + .filter(|(_, rank)| name != "r50k_base" || *rank < 50256) + .collect(); + if encoder + .values() + .collect::>() + .len() + != encoder.len() + || (0..=u8::MAX).any(|byte| !encoder.contains_key(&[byte][..])) + { + return Err(LoadError::Ranks("invalid vocabulary ranks".into())); + } + let (pattern, specials): (&str, &[(&str, Rank)]) = match name { + "cl100k_base" => ( + CL100K_PATTERN, + &[ + ("<|endoftext|>", 100257), + ("<|fim_prefix|>", 100258), + ("<|fim_middle|>", 100259), + ("<|fim_suffix|>", 100260), + ("<|endofprompt|>", 100276), + ], + ), + "o200k_base" => ( + O200K_BASE_PAT_STR, + &[("<|endoftext|>", 199999), ("<|endofprompt|>", 200018)], + ), + "o200k_harmony" => ( + O200K_BASE_PAT_STR, + &[ + ("<|startoftext|>", 199998), + ("<|endoftext|>", 199999), + ("<|reserved_200000|>", 200000), + ("<|reserved_200001|>", 200001), + ("<|return|>", 200002), + ("<|constrain|>", 200003), + ("<|reserved_200004|>", 200004), + ("<|channel|>", 200005), + ("<|start|>", 200006), + ("<|end|>", 200007), + ("<|message|>", 200008), + ("<|reserved_200009|>", 200009), + ("<|reserved_200010|>", 200010), + ("<|reserved_200011|>", 200011), + ("<|call|>", 200012), + ], + ), + "p50k_edit" => ( + LEGACY_PATTERN, + &[ + ("<|endoftext|>", 50256), + ("<|fim_prefix|>", 50281), + ("<|fim_middle|>", 50282), + ("<|fim_suffix|>", 50283), + ], + ), + _ => (LEGACY_PATTERN, &[("<|endoftext|>", 50256)]), + }; + let reserved = (200013..=201087) + .filter(|_| name == "o200k_harmony") + .map(|rank| (format!("<|reserved_{rank}|>"), rank)); + let special_tokens = specials + .iter() + .map(|(token, rank)| ((*token).to_owned(), *rank)) + .chain(reserved) + .collect(); + CoreBPE::new(encoder, special_tokens, pattern) + .map_err(|error| LoadError::Ranks(error.to_string())) +} + +fn parse_rank(line: &str) -> Result<(Vec, Rank), LoadError> { + let (token, rank) = line + .split_once(' ') + .ok_or_else(|| LoadError::Ranks("missing rank".into()))?; + let bytes = STANDARD + .decode(token) + .map_err(|error| LoadError::Ranks(error.to_string()))?; + let rank = rank + .parse() + .map_err(|error: std::num::ParseIntError| LoadError::Ranks(error.to_string()))?; + Ok((bytes, rank)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::TiktokenTokenizer; + + fn read_packaged_ranks(file: &str) -> std::io::Result { + std::fs::read_to_string( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../../litellm/litellm_core_utils/tokenizers") + .join(file), + ) + } + + #[test] + fn packaged_encodings_match_embedded_encodings_and_reuse_successful_loads() { + for name in [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ] { + if name != "gpt2" { + assert!( + TiktokenTokenizer::from_cached_ranks(name, |_| { + Err(std::io::Error::other("unreadable vocabulary")) + }) + .is_err() + ); + } + let loads = std::sync::atomic::AtomicUsize::new(0); + let barrier = std::sync::Barrier::new(4); + let encoders = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..4) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + TiktokenTokenizer::from_cached_ranks(name, |file| { + loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + read_packaged_ranks(file) + }) + .unwrap() + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + assert_eq!(loads.into_inner(), usize::from(name != "gpt2")); + let actual = &encoders[0]; + let expected = TiktokenTokenizer::from_name(name).unwrap(); + assert_eq!(actual.special_tokens(), expected.special_tokens()); + let specials: Vec<_> = expected.special_tokens().into_iter().collect(); + let special_text = specials.join(" "); + assert_eq!( + actual.encode_special(&special_text, &specials).unwrap(), + expected.encode_special(&special_text, &specials).unwrap() + ); + for text in [ + "", + "café 漢字 ع 🙂", + "a\r\nb\t ", + " hello 123456789", + &special_text, + ] { + let ids = expected.encode(text); + assert_eq!(actual.encode(text), ids, "{name}: {text:?}"); + assert_eq!(actual.count_tokens(text), ids.len(), "{name}: {text:?}"); + assert_eq!( + actual.decode_bytes(&ids).unwrap(), + expected.decode_bytes(&ids).unwrap() + ); + } + let cached = TiktokenTokenizer::from_cached_ranks(name, |_| { + panic!("reloaded cached vocabulary") + }) + .unwrap(); + assert_eq!(cached.encode("cached"), expected.encode("cached")); + } + } + + #[test] + fn malformed_ranks_return_errors_instead_of_panicking() { + for ranks in ["", "IQ==", "IQ== x", "!!! 1", "IQ== 1"] { + assert!(build("cl100k_base", ranks).is_err()); + } + let repeated_rank = (0..=u8::MAX) + .map(|byte| format!("{} 0\n", STANDARD.encode([byte]))) + .collect::(); + assert!(build("cl100k_base", &repeated_rank).is_err()); + } +} diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index c748e404845..51d6b23535f 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,4 +1,4 @@ -use litellm_token_counter_tiktoken::UnsupportedTokenizer; +use litellm_token_counter_tiktoken::{LoadError, UnsupportedTokenizer}; pub use litellm_token_counter_tiktoken::{TiktokenTokenizer, encoding_for_model}; use crate::{Error, TextCodec, TokenCounter, Tokenizer}; @@ -36,3 +36,12 @@ impl From for Error { Self::UnsupportedTokenizer(error.0) } } + +impl From for Error { + fn from(error: LoadError) -> Self { + match error { + LoadError::Unsupported(error) => error.into(), + LoadError::Ranks(message) => Self::Ranks(message), + } + } +} diff --git a/tests/test_litellm/litellm_core_utils/test_tokenizer.py b/tests/test_litellm/litellm_core_utils/test_tokenizer.py index caeb54fcd2a..31e136ed443 100644 --- a/tests/test_litellm/litellm_core_utils/test_tokenizer.py +++ b/tests/test_litellm/litellm_core_utils/test_tokenizer.py @@ -17,9 +17,11 @@ from litellm.utils import claude_json_str from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON -@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_base", "r50k_base", "o200k_harmony")) @pytest.mark.parametrize( - "text", ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff") + "name", ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "r50k_base", "gpt2", "o200k_harmony") +) +@pytest.mark.parametrize( + "text", ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64) ) def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: reference: Final = tiktoken.get_encoding(name) @@ -238,17 +240,26 @@ print("compatible") assert result.stdout.strip() == "compatible" -def test_installed_tokenization_does_not_import_python_tokenizer_packages() -> None: +def test_installed_tokenization_does_not_import_python_tokenizer_packages(tmp_path: Path) -> None: script: Final = """ import importlib.abc import sys sys.path.insert(0, sys.argv[1]) +def reject_network(event, args): + if event == "socket.connect": + raise AssertionError("tokenizer attempted a network connection") +sys.addaudithook(reject_network) class Block(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): if fullname.split(".")[0] in {"tiktoken", "tokenizers"}: raise AssertionError("unexpected runtime dependency: " + fullname) sys.meta_path.insert(0, Block()) import litellm +from litellm.litellm_core_utils.tokenizer import OpenAIEncoding +for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", "gpt2"): + encoding = OpenAIEncoding.from_tiktoken(name) + text = "offline café 漢字 🙂" + " " * 64 + assert encoding.decode(encoding.encode(text)) == text ids = litellm.encode(text="hello world") assert litellm.decode(tokens=ids) == "hello world" assert litellm.token_counter(model=None, text="hello world") == len(ids) @@ -261,10 +272,17 @@ print("compatible") capture_output=True, text=True, timeout=30, - env={**os.environ, "LITELLM_RUST": "0", "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + cwd=tmp_path, + env={ + **os.environ, + "LITELLM_RUST": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), + }, ) assert result.returncode == 0, result.stdout + result.stderr assert result.stdout.strip() == "compatible" + assert not (tmp_path / "unused-tokenizer-cache").exists() @pytest.mark.parametrize("is_pretokenized", (False, True)) diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index e449d4392d8..0eae041f535 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -56,10 +56,11 @@ def _write_wheel( metadata_tags: tuple[str, ...] | None = (_EXPECTED_TAG,), dist_info: str = _DIST_INFO, duplicate_wheel: bool = False, + native_bytes: bytes = b"synthetic native extension", ) -> Path: wheel: Final = tmp_path / f"litellm-1.100.0-{filename_tag}.whl" with zipfile.ZipFile(wheel, "w", compression=zipfile.ZIP_DEFLATED) as archive: - archive.writestr(_NATIVE_MEMBER, b"synthetic native extension") + archive.writestr(_NATIVE_MEMBER, native_bytes) archive.writestr( f"{dist_info}/METADATA", "Metadata-Version: 2.1\nName: litellm\nVersion: 1.100.0\n", @@ -195,3 +196,17 @@ def test_rejects_production_module_exposing_panic_hook(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) assert _run_verifier(wheel, exposes_panic=True) == 1 + + +@pytest.mark.parametrize("embedded", (False, True)) +def test_vocabulary_is_packaged_once(tmp_path: Path, embedded: bool) -> None: + ranks: Final = b"AA== 0\nAQ== 1\nAg== 2\n" + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + native_bytes=b"native engine" + (ranks if embedded else b""), + ) + with zipfile.ZipFile(wheel, "a") as archive: + archive.writestr("litellm/litellm_core_utils/tokenizers/" + "a" * 40, ranks) + + assert _run_verifier(wheel) == (1 if embedded else 0)