mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
feat(rust): add native tokenizer codec
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
f0ddb5045c
commit
b086dd17ae
15 changed files with 616 additions and 32 deletions
258
litellm-rust/Cargo.lock
generated
258
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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<Semaphore>,
|
||||
}
|
||||
|
||||
#[pyclass(frozen, name = "Tokenizer")]
|
||||
pub(crate) struct Tokenizer {
|
||||
inner: Arc<dyn TextCodec>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl TokenCounter {
|
||||
#[new]
|
||||
|
|
@ -115,6 +119,128 @@ impl TokenCounter {
|
|||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Tokenizer {
|
||||
#[staticmethod]
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[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<Self> {
|
||||
#[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<Self> {
|
||||
#[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<Vec<u32>> {
|
||||
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<u32>, skip_special_tokens: bool) -> PyResult<String> {
|
||||
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<usize> {
|
||||
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<T>(py: Python<'_>, load: impl FnOnce() -> Result<T, Error> + Send) -> PyResult<Self>
|
||||
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<String> {
|
||||
#[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),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,10 +14,63 @@ impl HuggingFaceTokenizer {
|
|||
.map_err(Error::Load)
|
||||
}
|
||||
|
||||
pub fn from_pretrained(
|
||||
identifier: &str,
|
||||
revision: &str,
|
||||
token: Option<&str>,
|
||||
) -> Result<Self, Error> {
|
||||
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<usize, Error> {
|
||||
self.0
|
||||
.encode_fast(text, true)
|
||||
.map(|encoding| encoding.len())
|
||||
.map_err(Error::Encode)
|
||||
}
|
||||
|
||||
pub fn encode(&self, text: &str) -> Result<Vec<u32>, 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<String, Error> {
|
||||
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("<SOS>hello<EOT>").unwrap();
|
||||
|
||||
assert!(tokenizer.decode(&ids, false).unwrap().contains("<SOS>"));
|
||||
assert_eq!(tokenizer.decode(&ids, true).unwrap(), "hello");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Self, UnsupportedTokenizer> {
|
||||
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<u32> {
|
||||
self.encoder.encode_ordinary(text)
|
||||
}
|
||||
|
||||
pub fn decode(&self, ids: &[u32]) -> Result<String, String> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ pub enum Error {
|
|||
UnsupportedTokenizer(String),
|
||||
#[error("failed to load tokenizer: {0}")]
|
||||
Load(#[source] Box<dyn std::error::Error + Send + Sync>),
|
||||
#[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<dyn std::error::Error + Send + Sync>),
|
||||
#[error("token decoding failed: {0}")]
|
||||
Decode(String),
|
||||
#[error("token counting task failed: {0}")]
|
||||
Task(String),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Self, Error> {
|
||||
|
|
@ -17,11 +17,27 @@ impl Tokenizer for HuggingFaceTokenizer {
|
|||
}
|
||||
}
|
||||
|
||||
impl TextCodec for HuggingFaceTokenizer {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, Error> {
|
||||
HuggingFaceTokenizer::encode(self, text).map_err(Error::from)
|
||||
}
|
||||
|
||||
fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String, Error> {
|
||||
HuggingFaceTokenizer::decode(self, ids, skip_special_tokens).map_err(Error::from)
|
||||
}
|
||||
|
||||
fn name(&self) -> &str {
|
||||
HuggingFaceTokenizer::name(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BackendError> 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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<Self, Error> {
|
||||
|
|
@ -17,6 +17,20 @@ impl Tokenizer for TiktokenTokenizer {
|
|||
}
|
||||
}
|
||||
|
||||
impl TextCodec for TiktokenTokenizer {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, Error> {
|
||||
Ok(TiktokenTokenizer::encode(self, text))
|
||||
}
|
||||
|
||||
fn decode(&self, ids: &[u32], _skip_special_tokens: bool) -> Result<String, Error> {
|
||||
TiktokenTokenizer::decode(self, ids).map_err(|error| Error::Decode(error.to_string()))
|
||||
}
|
||||
|
||||
fn name(&self) -> &str {
|
||||
TiktokenTokenizer::name(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<UnsupportedTokenizer> for Error {
|
||||
fn from(error: UnsupportedTokenizer) -> Self {
|
||||
Self::UnsupportedTokenizer(error.0)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,12 @@ pub trait Tokenizer: Send + Sync {
|
|||
fn count_tokens(&self, text: &str) -> Result<usize, Error>;
|
||||
}
|
||||
|
||||
pub trait TextCodec: Tokenizer {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, Error>;
|
||||
fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String, Error>;
|
||||
fn name(&self) -> &str;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
35
tests/test_litellm_rust/test_tokenizer.py
Normal file
35
tests/test_litellm_rust/test_tokenizer.py
Normal file
|
|
@ -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("<SOS>hello<EOT>")
|
||||
|
||||
assert "<SOS>" 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")
|
||||
Loading…
Add table
Reference in a new issue