From b086dd17ae339c23ac32aae5102ac8d67d53b076 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 22:52:08 +0000 Subject: [PATCH] feat(rust): add native tokenizer codec Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 258 +++++++++++++++++- litellm-rust/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 4 +- .../crates/python-bridge/src/token_counter.rs | 133 ++++++++- .../token-counter-huggingface/src/error.rs | 4 + .../token-counter-huggingface/src/lib.rs | 53 ++++ .../crates/token-counter-tiktoken/src/lib.rs | 87 +++++- .../crates/token-counter/src/error.rs | 4 + .../crates/token-counter/src/huggingface.rs | 18 +- litellm-rust/crates/token-counter/src/lib.rs | 2 +- .../crates/token-counter/src/tiktoken.rs | 18 +- .../crates/token-counter/src/tokenizer.rs | 6 + litellm/rust_bridge/_native.pyi | 22 ++ tests/test_litellm_rust/test_tokenizer.py | 35 +++ 15 files changed, 616 insertions(+), 32 deletions(-) create mode 100644 tests/test_litellm_rust/test_tokenizer.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index af3a31ddbfa..19961d020a7 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -816,6 +816,19 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414" +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -1127,6 +1140,27 @@ dependencies = [ "ctutils", ] +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.61.2", +] + [[package]] name = "displaydoc" version = "0.2.6" @@ -1156,6 +1190,12 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "equivalent" version = "1.0.2" @@ -1513,6 +1553,26 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hf-hub" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" +dependencies = [ + "dirs", + "http 1.4.2", + "indicatif", + "libc", + "log", + "rand 0.9.5", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror 2.0.19", + "ureq", + "windows-sys 0.60.2", +] + [[package]] name = "hmac" version = "0.13.0" @@ -1672,7 +1732,7 @@ dependencies = [ "tokio", "tokio-rustls 0.26.4", "tower-service", - "webpki-roots", + "webpki-roots 1.0.9", ] [[package]] @@ -1854,6 +1914,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -1966,6 +2039,15 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libredox" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6480ccc157a1389bb2e4891b24751b0f798ba640d22386f23143fbcc89da195a" +dependencies = [ + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2182,7 +2264,7 @@ dependencies = [ "thiserror 2.0.19", "tokio", "veil", - "webpki-roots", + "webpki-roots 1.0.9", ] [[package]] @@ -2477,6 +2559,12 @@ dependencies = [ "autocfg", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "once_cell" version = "1.21.4" @@ -2517,6 +2605,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + [[package]] name = "outref" version = "0.5.2" @@ -3042,6 +3136,16 @@ dependencies = [ "bitflags", ] +[[package]] +name = "redox_users" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60dc65c0ff1a7ae1294b0c67b9f14baf70b644404010370171787bfac1038fc0" +dependencies = [ + "libredox", + "thiserror 2.0.19", +] + [[package]] name = "ref-cast" version = "1.0.27" @@ -3143,7 +3247,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots", + "webpki-roots 1.0.9", ] [[package]] @@ -3285,6 +3389,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -3682,6 +3787,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + [[package]] name = "spm_precompiled" version = "0.1.4" @@ -3952,6 +4068,7 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", + "hf-hub", "itertools 0.14.0", "log", "macro_rules_attribute", @@ -4285,6 +4402,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -4297,6 +4420,25 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "once_cell", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + [[package]] name = "url" version = "2.5.8" @@ -4524,6 +4666,15 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.9", +] + [[package]] name = "webpki-roots" version = "1.0.9" @@ -4629,7 +4780,25 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", ] [[package]] @@ -4647,14 +4816,31 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", ] [[package]] @@ -4663,48 +4849,96 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + [[package]] name = "winnow" version = "1.0.4" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 4250fa81d26..7ac088838b4 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -48,7 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" -tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +tokenizers = { version = "0.23.1", default-features = false, features = ["onig", "http", "rustls-tls"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a76b069935f..8579cbaea4b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "fast"] +default = ["abi3", "fast", "huggingface", "tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 46f98736aa1..d39b9545070 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -29,7 +29,7 @@ mod _native { #[pymodule_export] use crate::routes::responses::ResponsesWebSocketConnection; #[pymodule_export] - use crate::token_counter::TokenCounter; + use crate::token_counter::{TokenCounter, Tokenizer, tiktoken_encoding_for_model}; #[pymodule_export] use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; } @@ -65,6 +65,8 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", + "Tokenizer", + "tiktoken_encoding_for_model", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 244401e6696..b4079edc40b 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -3,11 +3,10 @@ use std::sync::Arc; #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use std::{num::NonZero, thread::available_parallelism}; -#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use litellm_host_python::release_gil; use litellm_host_python::run_async; use litellm_token_counter::{ - CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, + CountableRequest, Error, InputTokenCount, TextCodec, TokenCounter as CoreTokenCounter, }; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, @@ -29,6 +28,11 @@ pub(crate) struct TokenCounter { encode_slots: Arc, } +#[pyclass(frozen, name = "Tokenizer")] +pub(crate) struct Tokenizer { + inner: Arc, +} + #[pymethods] impl TokenCounter { #[new] @@ -115,6 +119,128 @@ impl TokenCounter { } } +#[pymethods] +impl Tokenizer { + #[staticmethod] + fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + #[cfg(feature = "tiktoken")] + { + let encoding = encoding.to_owned(); + Self::load(py, move || { + litellm_token_counter::tiktoken::TiktokenTokenizer::from_name(&encoding) + .map_err(Error::from) + }) + } + #[cfg(not(feature = "tiktoken"))] + { + let _ = (py, encoding); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the tiktoken feature", + )) + } + } + + #[staticmethod] + fn from_json(py: Python<'_>, tokenizer_json: &str) -> PyResult { + #[cfg(feature = "huggingface")] + { + let tokenizer_json = tokenizer_json.to_owned(); + Self::load(py, move || { + litellm_token_counter::huggingface::HuggingFaceTokenizer::from_json(&tokenizer_json) + .map_err(Error::from) + }) + } + #[cfg(not(feature = "huggingface"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the huggingface feature", + )) + } + } + + #[staticmethod] + #[pyo3(signature = (identifier, revision = "main", token = None))] + fn from_pretrained( + py: Python<'_>, + identifier: &str, + revision: &str, + token: Option<&str>, + ) -> PyResult { + #[cfg(feature = "huggingface")] + { + let identifier = identifier.to_owned(); + let revision = revision.to_owned(); + let token = token.map(str::to_owned); + Self::load(py, move || { + litellm_token_counter::huggingface::HuggingFaceTokenizer::from_pretrained( + &identifier, + &revision, + token.as_deref(), + ) + .map_err(Error::from) + }) + } + #[cfg(not(feature = "huggingface"))] + { + let _ = (py, identifier, revision, token); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the huggingface feature", + )) + } + } + + fn encode(&self, py: Python<'_>, text: &str) -> PyResult> { + let inner = Arc::clone(&self.inner); + let text = text.to_owned(); + release_gil(py, move || inner.encode(&text)).map_err(token_count_error_to_pyerr) + } + + #[pyo3(signature = (ids, skip_special_tokens = true))] + fn decode(&self, py: Python<'_>, ids: Vec, skip_special_tokens: bool) -> PyResult { + let inner = Arc::clone(&self.inner); + release_gil(py, move || inner.decode(&ids, skip_special_tokens)) + .map_err(token_count_error_to_pyerr) + } + + fn count(&self, py: Python<'_>, text: &str) -> PyResult { + let inner = Arc::clone(&self.inner); + let text = text.to_owned(); + release_gil(py, move || inner.count_tokens(&text)).map_err(token_count_error_to_pyerr) + } + + #[getter] + fn name(&self) -> &str { + self.inner.name() + } +} + +impl Tokenizer { + #[cfg(any(feature = "huggingface", feature = "tiktoken"))] + fn load(py: Python<'_>, load: impl FnOnce() -> Result + Send) -> PyResult + where + T: TextCodec + 'static, + { + let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?; + Ok(Self { + inner: Arc::new(inner), + }) + } +} + +#[pyfunction] +pub(crate) fn tiktoken_encoding_for_model(model: &str) -> Option { + #[cfg(feature = "tiktoken")] + { + litellm_token_counter::tiktoken::encoding_for_model(model).map(str::to_owned) + } + #[cfg(not(feature = "tiktoken"))] + { + let _ = model; + None + } +} + impl TokenCounter { #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] fn load( @@ -143,6 +269,7 @@ fn token_count_error_to_pyerr(error: Error) -> PyErr { let message = error.to_string(); match error { Error::Load(_) + | Error::Download(_) | Error::Ranks(_) | Error::UnicodeClasses | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), @@ -153,6 +280,6 @@ fn token_count_error_to_pyerr(error: Error) -> PyErr { | Error::ArrayItems | Error::JsonSerialization(_) | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), - Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + Error::Encode(_) | Error::Decode(_) | Error::Task(_) => PyRuntimeError::new_err(message), } } diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs index adc4551886f..e5da1305480 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/error.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -4,6 +4,10 @@ use thiserror::Error as ThisError; pub enum Error { #[error("failed to load tokenizer: {0}")] Load(#[source] tokenizers::Error), + #[error("failed to download tokenizer: {0}")] + Download(#[source] tokenizers::Error), #[error("tokenization failed: {0}")] Encode(#[source] tokenizers::Error), + #[error("token decoding failed: {0}")] + Decode(#[source] tokenizers::Error), } diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 8e05c2cca46..2d64e3582f2 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -14,10 +14,63 @@ impl HuggingFaceTokenizer { .map_err(Error::Load) } + pub fn from_pretrained( + identifier: &str, + revision: &str, + token: Option<&str>, + ) -> Result { + tokenizers::Tokenizer::from_pretrained( + identifier, + Some(tokenizers::FromPretrainedParameters { + revision: revision.to_owned(), + token: token.map(str::to_owned), + ..Default::default() + }), + ) + .map(Box::new) + .map(Self) + .map_err(Error::Download) + } + pub fn count_tokens(&self, text: &str) -> Result { self.0 .encode_fast(text, true) .map(|encoding| encoding.len()) .map_err(Error::Encode) } + + pub fn encode(&self, text: &str) -> Result, Error> { + self.0 + .encode_fast(text, true) + .map(|encoding| encoding.get_ids().to_vec()) + .map_err(Error::Encode) + } + + pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + self.0 + .decode(ids, skip_special_tokens) + .map_err(Error::Decode) + } + + pub fn name(&self) -> &str { + "huggingface" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn codecs_round_trip_and_skip_special_tokens() { + let json = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + )); + let tokenizer = HuggingFaceTokenizer::from_json(json).unwrap(); + let ids = tokenizer.encode("hello").unwrap(); + + assert!(tokenizer.decode(&ids, false).unwrap().contains("")); + assert_eq!(tokenizer.decode(&ids, true).unwrap(), "hello"); + } } diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index ecdb3946eee..2c33a0dc714 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -4,24 +4,56 @@ mod error; pub use error::UnsupportedTokenizer; -pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); +pub struct TiktokenTokenizer { + encoder: &'static tiktoken_rs::CoreBPE, + name: &'static str, +} impl TiktokenTokenizer { pub fn from_name(name: &str) -> Result { - let tokenizer = match name { - "cl100k_base" => tiktoken_rs::cl100k_base_singleton(), - "o200k_base" => tiktoken_rs::o200k_base_singleton(), - "o200k_harmony" => tiktoken_rs::o200k_harmony_singleton(), - "p50k_base" => tiktoken_rs::p50k_base_singleton(), - "p50k_edit" => tiktoken_rs::p50k_edit_singleton(), - "r50k_base" | "gpt2" => tiktoken_rs::r50k_base_singleton(), + let (encoder, canonical_name) = match name { + "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"), + "o200k_base" => (tiktoken_rs::o200k_base_singleton(), "o200k_base"), + "o200k_harmony" => (tiktoken_rs::o200k_harmony_singleton(), "o200k_harmony"), + "p50k_base" => (tiktoken_rs::p50k_base_singleton(), "p50k_base"), + "p50k_edit" => (tiktoken_rs::p50k_edit_singleton(), "p50k_edit"), + "r50k_base" => (tiktoken_rs::r50k_base_singleton(), "r50k_base"), + "gpt2" => (tiktoken_rs::r50k_base_singleton(), "r50k_base"), _ => return Err(UnsupportedTokenizer(name.to_owned())), }; - Ok(Self(tokenizer)) + Ok(Self { + encoder, + name: canonical_name, + }) } pub fn count_tokens(&self, text: &str) -> usize { - self.0.count_ordinary(text) + self.encoder.count_ordinary(text) + } + + pub fn encode(&self, text: &str) -> Vec { + self.encoder.encode_ordinary(text) + } + + pub fn decode(&self, ids: &[u32]) -> Result { + self.encoder.decode(ids).map_err(|error| error.to_string()) + } + + pub fn name(&self) -> &str { + self.name + } +} + +pub fn encoding_for_model(model: &str) -> Option<&'static str> { + match tiktoken_rs::tokenizer::get_tokenizer(model)? { + tiktoken_rs::tokenizer::Tokenizer::Cl100kBase => Some("cl100k_base"), + tiktoken_rs::tokenizer::Tokenizer::O200kBase => Some("o200k_base"), + tiktoken_rs::tokenizer::Tokenizer::O200kHarmony => Some("o200k_harmony"), + tiktoken_rs::tokenizer::Tokenizer::P50kBase => Some("p50k_base"), + tiktoken_rs::tokenizer::Tokenizer::P50kEdit => Some("p50k_edit"), + tiktoken_rs::tokenizer::Tokenizer::R50kBase | tiktoken_rs::tokenizer::Tokenizer::Gpt2 => { + Some("r50k_base") + } } } @@ -66,5 +98,40 @@ mod tests { panic!("unknown encoding must be rejected"); }; assert_eq!(name, "unknown-encoding"); + assert_eq!( + TiktokenTokenizer::from_name("gpt2").unwrap().name(), + "r50k_base" + ); + } + + #[test] + fn codecs_round_trip_named_encodings() { + let encodings = [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ]; + let texts = ["hello world", "café 漢字 مرحبا 🙂", "line one\nline two"]; + for name in encodings { + let tokenizer = TiktokenTokenizer::from_name(name).unwrap(); + for text in texts { + assert_eq!( + tokenizer.decode(&tokenizer.encode(text)).unwrap(), + text, + "{name}: {text:?}", + ); + } + } + } + + #[test] + fn encoding_for_model_maps_known_models() { + assert_eq!(encoding_for_model("gpt-4o"), Some("o200k_base")); + assert_eq!(encoding_for_model("text-davinci-003"), Some("p50k_base")); + assert_eq!(encoding_for_model("unknown-model"), None); } } diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index b05ce007e46..dc1458d5602 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -8,6 +8,8 @@ pub enum Error { UnsupportedTokenizer(String), #[error("failed to load tokenizer: {0}")] Load(#[source] Box), + #[error("failed to download tokenizer: {0}")] + Download(String), #[error("failed to load tokenizer: tiktoken rank file: {0}")] Ranks(String), #[error("failed to load tokenizer: Unicode character classes are unavailable")] @@ -32,6 +34,8 @@ pub enum Error { JsonUtf8(#[source] FromUtf8Error), #[error("tokenization failed: {0}")] Encode(#[source] Box), + #[error("token decoding failed: {0}")] + Decode(String), #[error("token counting task failed: {0}")] Task(String), } diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs index fb7683b373e..1c0780bcf95 100644 --- a/litellm-rust/crates/token-counter/src/huggingface.rs +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -1,7 +1,7 @@ use litellm_token_counter_huggingface::Error as BackendError; pub use litellm_token_counter_huggingface::HuggingFaceTokenizer; -use crate::{Error, TokenCounter, Tokenizer}; +use crate::{Error, TextCodec, TokenCounter, Tokenizer}; impl TokenCounter { pub fn from_json(tokenizer_json: &str) -> Result { @@ -17,11 +17,27 @@ impl Tokenizer for HuggingFaceTokenizer { } } +impl TextCodec for HuggingFaceTokenizer { + fn encode(&self, text: &str) -> Result, Error> { + HuggingFaceTokenizer::encode(self, text).map_err(Error::from) + } + + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + HuggingFaceTokenizer::decode(self, ids, skip_special_tokens).map_err(Error::from) + } + + fn name(&self) -> &str { + HuggingFaceTokenizer::name(self) + } +} + impl From for Error { fn from(error: BackendError) -> Self { match error { BackendError::Load(source) => Self::Load(source), + BackendError::Download(source) => Self::Download(source.to_string()), BackendError::Encode(source) => Self::Encode(source), + BackendError::Decode(source) => Self::Decode(source.to_string()), } } } diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index 446c91049de..84bf35ca2ba 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -20,5 +20,5 @@ pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; pub use error::Error; -pub use tokenizer::Tokenizer; +pub use tokenizer::{TextCodec, Tokenizer}; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index 07c1c9f5b73..c748e404845 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,7 +1,7 @@ -pub use litellm_token_counter_tiktoken::TiktokenTokenizer; use litellm_token_counter_tiktoken::UnsupportedTokenizer; +pub use litellm_token_counter_tiktoken::{TiktokenTokenizer, encoding_for_model}; -use crate::{Error, TokenCounter, Tokenizer}; +use crate::{Error, TextCodec, TokenCounter, Tokenizer}; impl TokenCounter { pub fn from_tiktoken(encoding: &str) -> Result { @@ -17,6 +17,20 @@ impl Tokenizer for TiktokenTokenizer { } } +impl TextCodec for TiktokenTokenizer { + fn encode(&self, text: &str) -> Result, Error> { + Ok(TiktokenTokenizer::encode(self, text)) + } + + fn decode(&self, ids: &[u32], _skip_special_tokens: bool) -> Result { + TiktokenTokenizer::decode(self, ids).map_err(|error| Error::Decode(error.to_string())) + } + + fn name(&self) -> &str { + TiktokenTokenizer::name(self) + } +} + impl From for Error { fn from(error: UnsupportedTokenizer) -> Self { Self::UnsupportedTokenizer(error.0) diff --git a/litellm-rust/crates/token-counter/src/tokenizer.rs b/litellm-rust/crates/token-counter/src/tokenizer.rs index 88c29c672a7..146ac5b4d0d 100644 --- a/litellm-rust/crates/token-counter/src/tokenizer.rs +++ b/litellm-rust/crates/token-counter/src/tokenizer.rs @@ -4,6 +4,12 @@ pub trait Tokenizer: Send + Sync { fn count_tokens(&self, text: &str) -> Result; } +pub trait TextCodec: Tokenizer { + fn encode(&self, text: &str) -> Result, Error>; + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result; + fn name(&self) -> &str; +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..9d79800208d 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -104,6 +104,26 @@ class TokenCounter: def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... +@final +class Tokenizer: + @staticmethod + def from_tiktoken(encoding: str) -> Tokenizer: ... + @staticmethod + def from_json(tokenizer_json: str) -> Tokenizer: ... + @staticmethod + def from_pretrained( + identifier: str, + revision: str = "main", + token: str | None = None, + ) -> Tokenizer: ... + def encode(self, text: str) -> list[int]: ... + def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str: ... + def count(self, text: str) -> int: ... + @property + def name(self) -> str: ... + +def tiktoken_encoding_for_model(model: str) -> str | None: ... + def gil_stats() -> dict[str, int]: ... def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... @@ -115,6 +135,7 @@ __all__ = [ "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", + "Tokenizer", "achat_completions", "amessages", "aocr", @@ -126,5 +147,6 @@ __all__ = [ "ocr", "process_state_started", "reserve_process_for_forking", + "tiktoken_encoding_for_model", "transcription", ] diff --git a/tests/test_litellm_rust/test_tokenizer.py b/tests/test_litellm_rust/test_tokenizer.py new file mode 100644 index 00000000000..19640ff440f --- /dev/null +++ b/tests/test_litellm_rust/test_tokenizer.py @@ -0,0 +1,35 @@ +from typing import Final + +import pytest + +from litellm.rust_bridge import _native +from litellm.utils import claude_json_str + +pytestmark = pytest.mark.requires_rust_extension + + +def test_tiktoken_codec_round_trips_and_counts() -> None: + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + encoded: Final = tokenizer.encode("hello world") + + assert tokenizer.name == "cl100k_base" + assert tokenizer.count("hello world") == len(encoded) + assert tokenizer.decode(encoded) == "hello world" + + +def test_huggingface_codec_skips_special_tokens() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + encoded: Final = tokenizer.encode("hello") + + assert "" in tokenizer.decode(encoded, skip_special_tokens=False) + assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" + + +def test_tiktoken_encoding_for_model() -> None: + assert _native.tiktoken_encoding_for_model("gpt-4o") == "o200k_base" + assert _native.tiktoken_encoding_for_model("unknown-model") is None + + +def test_unknown_tiktoken_encoding_raises_value_error() -> None: + with pytest.raises(ValueError, match="unsupported tokenizer"): + _native.Tokenizer.from_tiktoken("unknown-encoding")