fix(tokenizer): reuse packaged vocabularies in the native wheel

.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),
     )
 

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",
 ]

litellm-rust/crates/python-bridge/src/token_counter.rs

@@ -84,7 +84,8 @@ impl TokenCounter {
     fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
         #[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"))]
         {

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<TiktokenTokenizer> {
+    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<Self> {
         #[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)),
             })

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

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<String>,
+    ) -> Result<Self, LoadError> {
+        let (encoder, name) = ranks::load(name, load)?;
+        Ok(Self { encoder, name })
+    }
+
     pub fn from_name(name: &str) -> Result<Self, UnsupportedTokenizer> {
         let (encoder, canonical_name) = match name {
             "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"),

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<CoreBPE> = OnceCell::new();
+static O200K_ENCODER: OnceCell<CoreBPE> = OnceCell::new();
+static HARMONY_ENCODER: OnceCell<CoreBPE> = OnceCell::new();
+static P50K_ENCODER: OnceCell<CoreBPE> = OnceCell::new();
+static EDIT_ENCODER: OnceCell<CoreBPE> = OnceCell::new();
+static R50K_ENCODER: OnceCell<CoreBPE> = 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<String>,
+) -> 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<CoreBPE, LoadError> {
+    let parsed = ranks
+        .lines()
+        .map(parse_rank)
+        .collect::<Result<Vec<_>, _>>()?;
+    let encoder: FxHashMap<_, _> = parsed
+        .into_iter()
+        .filter(|(_, rank)| name != "r50k_base" || *rank < 50256)
+        .collect();
+    if encoder
+        .values()
+        .collect::<std::collections::HashSet<_>>()
+        .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<u8>, 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<String> {
+        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::<Vec<_>>()
+            });
+            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::<String>();
+        assert!(build("cl100k_base", &repeated_rank).is_err());
+    }
+}

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<UnsupportedTokenizer> for Error {
         Self::UnsupportedTokenizer(error.0)
     }
 }
+
+impl From<LoadError> for Error {
+    fn from(error: LoadError) -> Self {
+        match error {
+            LoadError::Unsupported(error) => error.into(),
+            LoadError::Ranks(message) => Self::Ranks(message),
+        }
+    }
+}

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

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)