mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
refactor(rust): split token counter backends
This commit is contained in:
Yujong Lee
parent
0e61ed51c0
commit
84c26978a3
38 changed files with 944 additions and 501 deletions
83
litellm-rust/Cargo.lock
generated
83
litellm-rust/Cargo.lock
generated
|
|
@ -61,6 +61,12 @@ version = "1.0.14"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.104"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
|
||||
|
||||
[[package]]
|
||||
name = "arc-swap"
|
||||
version = "1.9.2"
|
||||
|
|
@ -598,6 +604,17 @@ dependencies = [
|
|||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bstr"
|
||||
version = "1.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"regex-automata",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.20.3"
|
||||
|
|
@ -1181,6 +1198,17 @@ dependencies = [
|
|||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fancy-regex"
|
||||
version = "0.17.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8"
|
||||
dependencies = [
|
||||
"bit-set",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fancy-regex"
|
||||
version = "0.19.2"
|
||||
|
|
@ -1926,6 +1954,12 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.186"
|
||||
|
|
@ -2083,7 +2117,7 @@ dependencies = [
|
|||
name = "litellm-core-utils"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"fancy-regex",
|
||||
"fancy-regex 0.19.2",
|
||||
"litellm-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -2212,10 +2246,24 @@ dependencies = [
|
|||
name = "litellm-token-counter"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"itoa",
|
||||
"litellm-token-counter-fast",
|
||||
"litellm-token-counter-huggingface",
|
||||
"litellm-token-counter-tiktoken",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-token-counter-fast"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"rustc-hash",
|
||||
|
|
@ -2226,6 +2274,22 @@ dependencies = [
|
|||
"unicode-normalization-alignments",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-token-counter-huggingface"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"thiserror 2.0.19",
|
||||
"tokenizers",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-token-counter-tiktoken"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"thiserror 2.0.19",
|
||||
"tiktoken-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-types"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3794,6 +3858,21 @@ dependencies = [
|
|||
"syn 3.0.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tiktoken-rs"
|
||||
version = "0.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"base64 0.22.1",
|
||||
"bstr",
|
||||
"fancy-regex 0.17.0",
|
||||
"lazy_static",
|
||||
"regex",
|
||||
"rustc-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.53"
|
||||
|
|
|
|||
|
|
@ -24,6 +24,9 @@ litellm-core-utils = { path = "crates/core-utils" }
|
|||
litellm-cache = { path = "crates/cache" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-token-counter = { path = "crates/token-counter" }
|
||||
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
|
||||
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
|
||||
litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
|
||||
litellm-host-python = { path = "crates/host-python" }
|
||||
|
||||
bytes = "1"
|
||||
|
|
@ -45,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std",
|
|||
sha2 = "0.10"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
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"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
|
|
|
|||
|
|
@ -10,10 +10,13 @@ name = "_native"
|
|||
crate-type = ["cdylib"]
|
||||
|
||||
[features]
|
||||
default = ["abi3"]
|
||||
default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"]
|
||||
abi3 = ["pyo3/abi3-py310"]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
panic-test = []
|
||||
token-counter-fast = ["litellm-token-counter/fast"]
|
||||
token-counter-huggingface = ["litellm-token-counter/huggingface"]
|
||||
token-counter-tiktoken = ["litellm-token-counter/tiktoken"]
|
||||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
|
|
@ -26,7 +29,7 @@ litellm-http.workspace = true
|
|||
litellm-llms.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
litellm-token-counter.workspace = true
|
||||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,6 +1,19 @@
|
|||
use std::{num::NonZero, sync::Arc, thread::available_parallelism};
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_host_python::{release_gil, run_async};
|
||||
#[cfg(any(
|
||||
feature = "token-counter-fast",
|
||||
feature = "token-counter-huggingface",
|
||||
feature = "token-counter-tiktoken"
|
||||
))]
|
||||
use std::{num::NonZero, thread::available_parallelism};
|
||||
|
||||
#[cfg(any(
|
||||
feature = "token-counter-fast",
|
||||
feature = "token-counter-huggingface",
|
||||
feature = "token-counter-tiktoken"
|
||||
))]
|
||||
use litellm_host_python::release_gil;
|
||||
use litellm_host_python::run_async;
|
||||
use litellm_token_counter::{
|
||||
CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
|
||||
};
|
||||
|
|
@ -28,17 +41,47 @@ pub(crate) struct TokenCounter {
|
|||
impl TokenCounter {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
#[cfg(feature = "token-counter-huggingface")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
}
|
||||
#[cfg(not(feature = "token-counter-huggingface"))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the token-counter-huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "token-counter-fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
|
||||
}
|
||||
#[cfg(not(feature = "token-counter-fast"))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the token-counter-fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "token-counter-tiktoken")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
|
||||
}
|
||||
#[cfg(not(feature = "token-counter-tiktoken"))]
|
||||
{
|
||||
let _ = (py, encoding);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the token-counter-tiktoken feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
|
|
@ -62,6 +105,11 @@ impl TokenCounter {
|
|||
}
|
||||
|
||||
impl TokenCounter {
|
||||
#[cfg(any(
|
||||
feature = "token-counter-fast",
|
||||
feature = "token-counter-huggingface",
|
||||
feature = "token-counter-tiktoken"
|
||||
))]
|
||||
fn load(
|
||||
py: Python<'_>,
|
||||
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
|
||||
|
|
@ -74,6 +122,11 @@ impl TokenCounter {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
feature = "token-counter-fast",
|
||||
feature = "token-counter-huggingface",
|
||||
feature = "token-counter-tiktoken"
|
||||
))]
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(1, NonZero::get)
|
||||
}
|
||||
|
|
@ -86,7 +139,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount
|
|||
fn token_count_error_to_pyerr(error: Error) -> PyErr {
|
||||
let message = error.to_string();
|
||||
match error {
|
||||
Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
|
||||
Error::Load(_)
|
||||
| Error::Ranks(_)
|
||||
| Error::UnicodeClasses
|
||||
| Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
|
||||
Error::RequestParse(_)
|
||||
| Error::MissingInput
|
||||
| Error::FloatText
|
||||
|
|
|
|||
19
litellm-rust/crates/token-counter-fast/Cargo.toml
Normal file
19
litellm-rust/crates/token-counter-fast/Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "litellm-token-counter-fast"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
rustc-hash = "2.1.3"
|
||||
thiserror.workspace = true
|
||||
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
|
||||
unicode-normalization-alignments = "0.1.12"
|
||||
|
||||
[dev-dependencies]
|
||||
rand.workspace = true
|
||||
rstest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
@ -473,13 +473,13 @@ mod tests {
|
|||
_ => unreachable!(),
|
||||
}
|
||||
assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none());
|
||||
let counter = crate::TokenCounter::from_json(
|
||||
let counter = crate::FastTokenizer::from_json(
|
||||
&anthropic_tokenizer.to_string(false).expect("serialize"),
|
||||
)
|
||||
.expect("load");
|
||||
for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", "<EOT> stop"] {
|
||||
assert_eq!(
|
||||
counter.count_text(text).expect("count"),
|
||||
counter.count_tokens(text).expect("count"),
|
||||
reference_count(&anthropic_tokenizer, text)
|
||||
);
|
||||
}
|
||||
|
|
@ -545,7 +545,7 @@ mod tests {
|
|||
.rstrip(rstrip)])
|
||||
.expect("add token");
|
||||
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
|
||||
let counter = crate::TokenCounter::from_json(
|
||||
let counter = crate::FastTokenizer::from_json(
|
||||
&anthropic_tokenizer.to_string(false).expect("serialize"),
|
||||
)
|
||||
.expect("load");
|
||||
|
|
@ -557,7 +557,7 @@ mod tests {
|
|||
] {
|
||||
assert_eq!(fast.count(&anthropic_tokenizer, text), None);
|
||||
assert_eq!(
|
||||
counter.count_text(text).expect("count"),
|
||||
counter.count_tokens(text).expect("count"),
|
||||
reference_count(&anthropic_tokenizer, text)
|
||||
);
|
||||
}
|
||||
|
|
@ -571,17 +571,17 @@ mod tests {
|
|||
assert_eq!(fast.count(&tokenizer, "hello"), None);
|
||||
assert!(tokenizer.encode_fast("hello", true).is_err());
|
||||
let counter =
|
||||
crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize"))
|
||||
crate::FastTokenizer::from_json(&tokenizer.to_string(false).expect("serialize"))
|
||||
.expect("load");
|
||||
assert!(matches!(
|
||||
counter.count_text("hello"),
|
||||
counter.count_tokens("hello"),
|
||||
Err(crate::Error::Encode(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) {
|
||||
let counter = crate::TokenCounter::from_json(
|
||||
let counter = crate::FastTokenizer::from_json(
|
||||
&anthropic_tokenizer.to_string(false).expect("serialize"),
|
||||
)
|
||||
.expect("load");
|
||||
|
|
@ -598,7 +598,7 @@ mod tests {
|
|||
scope.spawn(move || {
|
||||
for _ in 0..100 {
|
||||
for (text, count) in inputs.iter().zip(expected) {
|
||||
assert_eq!(counter.count_text(text).expect("count"), count);
|
||||
assert_eq!(counter.count_tokens(text).expect("count"), count);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
|
@ -614,10 +614,10 @@ mod tests {
|
|||
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
|
||||
assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1);
|
||||
assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None);
|
||||
let counter = crate::TokenCounter::from_json(
|
||||
let counter = crate::FastTokenizer::from_json(
|
||||
&anthropic_tokenizer.to_string(false).expect("serialize"),
|
||||
)
|
||||
.expect("load");
|
||||
assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1);
|
||||
assert_eq!(counter.count_tokens("ABCD EFGH").expect("count"), 1);
|
||||
}
|
||||
}
|
||||
13
litellm-rust/crates/token-counter-fast/src/error.rs
Normal file
13
litellm-rust/crates/token-counter-fast/src/error.rs
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
use thiserror::Error as ThisError;
|
||||
|
||||
#[derive(Debug, ThisError)]
|
||||
pub enum Error {
|
||||
#[error("failed to load tokenizer: {0}")]
|
||||
Load(#[source] tokenizers::Error),
|
||||
#[error("failed to load tokenizer: tiktoken rank file: {0}")]
|
||||
Ranks(String),
|
||||
#[error("failed to load tokenizer: Unicode character classes are unavailable")]
|
||||
UnicodeClasses,
|
||||
#[error("tokenization failed: {0}")]
|
||||
Encode(#[source] tokenizers::Error),
|
||||
}
|
||||
70
litellm-rust/crates/token-counter-fast/src/lib.rs
Normal file
70
litellm-rust/crates/token-counter-fast/src/lib.rs
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod byte_level;
|
||||
mod cl100k;
|
||||
mod error;
|
||||
mod o200k;
|
||||
mod scanner;
|
||||
mod tiktoken;
|
||||
mod unicode_classes;
|
||||
|
||||
use byte_level::ByteLevelCounter;
|
||||
use scanner::{SplitPattern, TiktokenCounter};
|
||||
|
||||
pub use error::Error;
|
||||
|
||||
enum Encoder {
|
||||
HuggingFace {
|
||||
tokenizer: Box<tokenizers::Tokenizer>,
|
||||
byte_level: Option<ByteLevelCounter>,
|
||||
},
|
||||
Tiktoken(TiktokenCounter),
|
||||
}
|
||||
|
||||
pub struct FastTokenizer(Encoder);
|
||||
|
||||
impl FastTokenizer {
|
||||
pub fn from_json(json: &str) -> Result<Self, Error> {
|
||||
let tokenizer = json.parse::<tokenizers::Tokenizer>().map_err(Error::Load)?;
|
||||
let byte_level = ByteLevelCounter::detect(&tokenizer);
|
||||
Ok(Self(Encoder::HuggingFace {
|
||||
tokenizer: Box::new(tokenizer),
|
||||
byte_level,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn from_cl100k_ranks(ranks: &str) -> Result<Self, Error> {
|
||||
Self::from_ranks(SplitPattern::Cl100k, ranks)
|
||||
}
|
||||
|
||||
pub fn from_o200k_ranks(ranks: &str) -> Result<Self, Error> {
|
||||
Self::from_ranks(SplitPattern::O200k, ranks)
|
||||
}
|
||||
|
||||
fn from_ranks(split: SplitPattern, ranks: &str) -> Result<Self, Error> {
|
||||
TiktokenCounter::from_ranks(split, ranks)
|
||||
.map(Encoder::Tiktoken)
|
||||
.map(Self)
|
||||
}
|
||||
|
||||
pub fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
match &self.0 {
|
||||
Encoder::Tiktoken(counter) => Ok(counter.count(text)),
|
||||
Encoder::HuggingFace {
|
||||
tokenizer,
|
||||
byte_level,
|
||||
} => {
|
||||
if let Some(count) = byte_level
|
||||
.as_ref()
|
||||
.and_then(|counter| counter.count(tokenizer, text))
|
||||
{
|
||||
return Ok(count);
|
||||
}
|
||||
tokenizer
|
||||
.encode_fast(text, true)
|
||||
.map(|encoding| encoding.len())
|
||||
.map_err(Error::Encode)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
222
litellm-rust/crates/token-counter-fast/src/tiktoken.rs
Normal file
222
litellm-rust/crates/token-counter-fast/src/tiktoken.rs
Normal file
|
|
@ -0,0 +1,222 @@
|
|||
//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and
|
||||
//! the merge loop that turns one regex piece into tokens. The merge order is
|
||||
//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is
|
||||
//! identical, but pairs are tracked in a heap so a long piece costs
|
||||
//! `O(n log n)` instead of tiktoken's `O(n^2)`.
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::BinaryHeap;
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use rustc_hash::FxHashMap;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
type Rank = u32;
|
||||
|
||||
const NO_RANK: Rank = Rank::MAX;
|
||||
const END: usize = usize::MAX;
|
||||
|
||||
pub(super) struct MergeRanks(FxHashMap<Box<[u8]>, Rank>);
|
||||
|
||||
impl MergeRanks {
|
||||
pub(super) fn parse(text: &str) -> Result<Self, Error> {
|
||||
let ranks = text
|
||||
.lines()
|
||||
.filter(|line| !line.is_empty())
|
||||
.map(parse_line)
|
||||
.collect::<Result<FxHashMap<_, _>, _>>()?;
|
||||
if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) {
|
||||
return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token")));
|
||||
}
|
||||
Ok(Self(ranks))
|
||||
}
|
||||
|
||||
fn rank(&self, bytes: &[u8]) -> Rank {
|
||||
self.0.get(bytes).copied().unwrap_or(NO_RANK)
|
||||
}
|
||||
|
||||
/// Token count of one regex piece, as `encode_ordinary` would produce.
|
||||
pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize {
|
||||
if piece.len() < 2 || self.0.contains_key(piece) {
|
||||
return 1;
|
||||
}
|
||||
scratch.reset(piece.len());
|
||||
for start in 0..piece.len() - 1 {
|
||||
scratch.set_rank(start, self.rank(&piece[start..start + 2]));
|
||||
}
|
||||
let mut parts = piece.len();
|
||||
while let Some(Reverse((rank, start))) = scratch.heap.pop() {
|
||||
if scratch.next[start] == END || scratch.rank[start] != rank {
|
||||
continue;
|
||||
}
|
||||
let merged = scratch.next[start];
|
||||
let after = scratch.next[merged];
|
||||
scratch.next[merged] = END;
|
||||
scratch.next[start] = after;
|
||||
parts -= 1;
|
||||
if after < piece.len() {
|
||||
scratch.prev[after] = start;
|
||||
scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)]));
|
||||
} else {
|
||||
scratch.rank[start] = NO_RANK;
|
||||
}
|
||||
let before = scratch.prev[start];
|
||||
if before != END {
|
||||
scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)]));
|
||||
}
|
||||
}
|
||||
parts
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> {
|
||||
let (token, rank) = line
|
||||
.split_once(' ')
|
||||
.ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?;
|
||||
let bytes = STANDARD
|
||||
.decode(token)
|
||||
.map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?;
|
||||
let rank = rank
|
||||
.parse()
|
||||
.map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?;
|
||||
Ok((bytes.into_boxed_slice(), rank))
|
||||
}
|
||||
|
||||
/// Buffers reused across the pieces of one text. Parts are addressed by the
|
||||
/// byte offset they start at, which also gives the leftmost-pair tie break.
|
||||
#[derive(Default)]
|
||||
pub(super) struct MergeScratch {
|
||||
next: Vec<usize>,
|
||||
prev: Vec<usize>,
|
||||
rank: Vec<Rank>,
|
||||
heap: BinaryHeap<Reverse<(Rank, usize)>>,
|
||||
}
|
||||
|
||||
impl MergeScratch {
|
||||
fn reset(&mut self, len: usize) {
|
||||
self.next.clear();
|
||||
self.next.extend(1..=len);
|
||||
self.prev.clear();
|
||||
self.prev.push(END);
|
||||
self.prev.extend(0..len - 1);
|
||||
self.rank.clear();
|
||||
self.rank.resize(len, NO_RANK);
|
||||
self.heap.clear();
|
||||
}
|
||||
|
||||
fn end(&self, start: usize) -> usize {
|
||||
self.next[start]
|
||||
}
|
||||
|
||||
fn set_rank(&mut self, start: usize, rank: Rank) {
|
||||
self.rank[start] = rank;
|
||||
if rank != NO_RANK {
|
||||
self.heap.push(Reverse((rank, start)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn ranks() -> MergeRanks {
|
||||
let path = concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
|
||||
);
|
||||
MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo"))
|
||||
.expect("rank file parses")
|
||||
}
|
||||
|
||||
/// tiktoken's `_byte_pair_merge`, transcribed, as the reference.
|
||||
fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize {
|
||||
if piece.len() < 2 || ranks.0.contains_key(piece) {
|
||||
return 1;
|
||||
}
|
||||
let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1)
|
||||
.map(|index| (index, ranks.rank(&piece[index..index + 2])))
|
||||
.chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)])
|
||||
.collect();
|
||||
let get_rank = |parts: &[(usize, Rank)], index: usize| {
|
||||
if index + 3 < parts.len() {
|
||||
ranks.rank(&piece[parts[index].0..parts[index + 3].0])
|
||||
} else {
|
||||
NO_RANK
|
||||
}
|
||||
};
|
||||
loop {
|
||||
let Some(index) = parts[..parts.len() - 1]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, (_, rank))| *rank != NO_RANK)
|
||||
.min_by_key(|(index, (_, rank))| (*rank, *index))
|
||||
.map(|(index, _)| index)
|
||||
else {
|
||||
return parts.len() - 1;
|
||||
};
|
||||
if index > 0 {
|
||||
parts[index - 1].1 = get_rank(&parts, index - 1);
|
||||
}
|
||||
parts[index].1 = get_rank(&parts, index);
|
||||
parts.remove(index + 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_byte_is_a_token() {
|
||||
let ranks = ranks();
|
||||
assert_eq!(ranks.0.len(), 100_256);
|
||||
assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heap_merge_matches_tiktokens_merge_loop() {
|
||||
let ranks = ranks();
|
||||
let mut scratch = MergeScratch::default();
|
||||
let mut rng = StdRng::seed_from_u64(99);
|
||||
let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123";
|
||||
for _ in 0..20_000 {
|
||||
let piece: Vec<u8> = (0..rng.gen_range(1..24))
|
||||
.map(|_| alphabet[rng.gen_range(0..alphabet.len())])
|
||||
.collect();
|
||||
assert_eq!(
|
||||
ranks.count_piece(&piece, &mut scratch),
|
||||
reference_count(&ranks, &piece),
|
||||
"piece {:?}",
|
||||
String::from_utf8_lossy(&piece)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_repeated_runs_cost_close_to_linear() {
|
||||
let ranks = ranks();
|
||||
let mut scratch = MergeScratch::default();
|
||||
let mut time = |len: usize| {
|
||||
let piece = vec![b' '; len];
|
||||
let started = std::time::Instant::now();
|
||||
assert!(ranks.count_piece(&piece, &mut scratch) > 0);
|
||||
started.elapsed()
|
||||
};
|
||||
let small = (0..3).map(|_| time(1 << 14)).min().unwrap();
|
||||
let large = time(1 << 18);
|
||||
assert!(
|
||||
large < small * 64,
|
||||
"{small:?} for 2^14 bytes, {large:?} for 2^18"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_rank_files_are_rejected() {
|
||||
assert!(MergeRanks::parse("IQ==").is_err());
|
||||
assert!(MergeRanks::parse("IQ== x").is_err());
|
||||
assert!(MergeRanks::parse("!!! 1").is_err());
|
||||
assert!(MergeRanks::parse("IQ== 1").is_err());
|
||||
}
|
||||
}
|
||||
10
litellm-rust/crates/token-counter-huggingface/Cargo.toml
Normal file
10
litellm-rust/crates/token-counter-huggingface/Cargo.toml
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
[package]
|
||||
name = "litellm-token-counter-huggingface"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
thiserror.workspace = true
|
||||
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
|
||||
29
litellm-rust/crates/token-counter-huggingface/src/lib.rs
Normal file
29
litellm-rust/crates/token-counter-huggingface/src/lib.rs
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
#[derive(Debug, ThisError)]
|
||||
pub enum Error {
|
||||
#[error("failed to load tokenizer: {0}")]
|
||||
Load(#[source] tokenizers::Error),
|
||||
#[error("tokenization failed: {0}")]
|
||||
Encode(#[source] tokenizers::Error),
|
||||
}
|
||||
|
||||
pub struct HuggingFaceTokenizer(Box<tokenizers::Tokenizer>);
|
||||
|
||||
impl HuggingFaceTokenizer {
|
||||
pub fn from_json(json: &str) -> Result<Self, Error> {
|
||||
json.parse::<tokenizers::Tokenizer>()
|
||||
.map(Box::new)
|
||||
.map(Self)
|
||||
.map_err(Error::Load)
|
||||
}
|
||||
|
||||
pub fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
self.0
|
||||
.encode_fast(text, true)
|
||||
.map(|encoding| encoding.len())
|
||||
.map_err(Error::Encode)
|
||||
}
|
||||
}
|
||||
10
litellm-rust/crates/token-counter-tiktoken/Cargo.toml
Normal file
10
litellm-rust/crates/token-counter-tiktoken/Cargo.toml
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
[package]
|
||||
name = "litellm-token-counter-tiktoken"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
thiserror.workspace = true
|
||||
tiktoken-rs.workspace = true
|
||||
54
litellm-rust/crates/token-counter-tiktoken/src/lib.rs
Normal file
54
litellm-rust/crates/token-counter-tiktoken/src/lib.rs
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
#[derive(Debug, ThisError)]
|
||||
#[error("unsupported tokenizer: {0}")]
|
||||
pub struct UnsupportedTokenizer(pub String);
|
||||
|
||||
pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE);
|
||||
|
||||
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(),
|
||||
_ => return Err(UnsupportedTokenizer(name.to_owned())),
|
||||
};
|
||||
Ok(Self(tokenizer))
|
||||
}
|
||||
|
||||
pub fn count_tokens(&self, text: &str) -> usize {
|
||||
self.0.count_ordinary(text)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn special_tokens_are_counted_as_ordinary_text() {
|
||||
let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap();
|
||||
assert!(counter.count_tokens("<|endoftext|>") > 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_python_tiktoken_encodings_are_available() {
|
||||
for name in [
|
||||
"cl100k_base",
|
||||
"o200k_base",
|
||||
"o200k_harmony",
|
||||
"p50k_base",
|
||||
"p50k_edit",
|
||||
"r50k_base",
|
||||
"gpt2",
|
||||
] {
|
||||
assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -5,16 +5,21 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
default = ["huggingface", "tiktoken"]
|
||||
fast = ["dep:litellm-token-counter-fast"]
|
||||
huggingface = ["dep:litellm-token-counter-huggingface"]
|
||||
tiktoken = ["dep:litellm-token-counter-tiktoken"]
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
indexmap = { version = "2.14.0", features = ["serde"] }
|
||||
itoa = "1.0"
|
||||
rustc-hash = "2.1.3"
|
||||
litellm-token-counter-fast = { workspace = true, optional = true }
|
||||
litellm-token-counter-huggingface = { workspace = true, optional = true }
|
||||
litellm-token-counter-tiktoken = { workspace = true, optional = true }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
|
||||
unicode-normalization-alignments = "0.1.12"
|
||||
|
||||
[dev-dependencies]
|
||||
criterion.workspace = true
|
||||
|
|
@ -24,7 +29,9 @@ rstest.workspace = true
|
|||
[[bench]]
|
||||
name = "token_counter"
|
||||
harness = false
|
||||
required-features = ["fast"]
|
||||
|
||||
[[bench]]
|
||||
name = "allocations"
|
||||
harness = false
|
||||
required-features = ["fast"]
|
||||
|
|
|
|||
23
litellm-rust/crates/token-counter/README.md
Normal file
23
litellm-rust/crates/token-counter/README.md
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
# Token counting
|
||||
|
||||
`Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface
|
||||
|
||||
The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation
|
||||
|
||||
The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation
|
||||
|
||||
The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2`
|
||||
|
||||
The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend
|
||||
|
||||
Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits
|
||||
|
||||
Run the feature matrix with:
|
||||
|
||||
```sh
|
||||
cargo test -p litellm-token-counter
|
||||
cargo test -p litellm-token-counter --no-default-features
|
||||
cargo test -p litellm-token-counter --no-default-features --features fast
|
||||
cargo test -p litellm-token-counter --no-default-features --features huggingface
|
||||
cargo test -p litellm-token-counter --no-default-features --features tiktoken
|
||||
```
|
||||
|
|
@ -87,7 +87,7 @@ fn main() {
|
|||
},
|
||||
);
|
||||
|
||||
let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads");
|
||||
let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("tokenizer loads");
|
||||
let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses");
|
||||
counter
|
||||
.count_request(&object)
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> {
|
|||
}
|
||||
|
||||
fn token_counter(c: &mut Criterion) {
|
||||
let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load");
|
||||
let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("token counter should load");
|
||||
let tokenizer = TOKENIZER_JSON
|
||||
.parse::<Tokenizer>()
|
||||
.expect("reference tokenizer should load");
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
use serde::Serialize;
|
||||
|
||||
use crate::Error;
|
||||
use crate::byte_level::ByteLevelCounter;
|
||||
use crate::python_json;
|
||||
use crate::scanner::{SplitPattern, TiktokenCounter};
|
||||
use crate::tools::format_function_definitions;
|
||||
use crate::types::{
|
||||
ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice,
|
||||
|
|
@ -24,73 +22,22 @@ pub struct InputTokenCount {
|
|||
pub input_tokens: usize,
|
||||
}
|
||||
|
||||
enum Encoder {
|
||||
HuggingFace {
|
||||
tokenizer: Box<tokenizers::Tokenizer>,
|
||||
byte_level: Option<ByteLevelCounter>,
|
||||
},
|
||||
Tiktoken(TiktokenCounter),
|
||||
}
|
||||
|
||||
/// A loaded tokenizer plus the message accounting Python applies on top of
|
||||
/// it. Encoding is CPU-bound and synchronous; hosts run it off their event
|
||||
/// loop.
|
||||
pub struct TokenCounter {
|
||||
encoder: Encoder,
|
||||
encoder: Box<dyn crate::Tokenizer>,
|
||||
}
|
||||
|
||||
impl TokenCounter {
|
||||
/// Load a HuggingFace `tokenizer.json` document. The host reads the file.
|
||||
pub fn from_json(tokenizer_json: &str) -> Result<Self, Error> {
|
||||
let tokenizer = tokenizer_json
|
||||
.parse::<tokenizers::Tokenizer>()
|
||||
.map_err(Error::Load)?;
|
||||
let byte_level = ByteLevelCounter::detect(&tokenizer);
|
||||
Ok(Self {
|
||||
encoder: Encoder::HuggingFace {
|
||||
tokenizer: Box::new(tokenizer),
|
||||
byte_level,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Load tiktoken's `cl100k_base` rank file (`base64(token) rank` lines).
|
||||
/// The host reads the file.
|
||||
pub fn from_cl100k_ranks(rank_file: &str) -> Result<Self, Error> {
|
||||
Self::from_tiktoken_ranks(SplitPattern::Cl100k, rank_file)
|
||||
}
|
||||
|
||||
/// Load tiktoken's `o200k_base` rank file (`base64(token) rank` lines).
|
||||
/// The host reads the file.
|
||||
pub fn from_o200k_ranks(rank_file: &str) -> Result<Self, Error> {
|
||||
Self::from_tiktoken_ranks(SplitPattern::O200k, rank_file)
|
||||
}
|
||||
|
||||
fn from_tiktoken_ranks(split: SplitPattern, rank_file: &str) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?),
|
||||
})
|
||||
pub fn new(tokenizer: impl crate::Tokenizer + 'static) -> Self {
|
||||
Self {
|
||||
encoder: Box::new(tokenizer),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn count_text(&self, text: &str) -> Result<usize, Error> {
|
||||
match &self.encoder {
|
||||
Encoder::Tiktoken(counter) => Ok(counter.count(text)),
|
||||
Encoder::HuggingFace {
|
||||
tokenizer,
|
||||
byte_level,
|
||||
} => {
|
||||
if let Some(count) = byte_level
|
||||
.as_ref()
|
||||
.and_then(|counter| counter.count(tokenizer, text))
|
||||
{
|
||||
return Ok(count);
|
||||
}
|
||||
tokenizer
|
||||
.encode_fast(text, true)
|
||||
.map(|encoding| encoding.len())
|
||||
.map_err(Error::Encode)
|
||||
}
|
||||
}
|
||||
self.encoder.count_tokens(text)
|
||||
}
|
||||
|
||||
/// Mirrors the host's key precedence: `messages`, then `prompt`, then
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ use thiserror::Error as ThisError;
|
|||
|
||||
#[derive(Debug, ThisError)]
|
||||
pub enum Error {
|
||||
#[error("unsupported tokenizer: {0}")]
|
||||
UnsupportedTokenizer(String),
|
||||
#[error("failed to load tokenizer: {0}")]
|
||||
Load(#[source] tokenizers::Error),
|
||||
Load(#[source] Box<dyn std::error::Error + Send + Sync>),
|
||||
#[error("failed to load tokenizer: tiktoken rank file: {0}")]
|
||||
Ranks(String),
|
||||
#[error("failed to load tokenizer: Unicode character classes are unavailable")]
|
||||
|
|
@ -29,7 +31,7 @@ pub enum Error {
|
|||
#[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")]
|
||||
JsonUtf8(#[source] FromUtf8Error),
|
||||
#[error("tokenization failed: {0}")]
|
||||
Encode(#[source] tokenizers::Error),
|
||||
Encode(#[source] Box<dyn std::error::Error + Send + Sync>),
|
||||
#[error("token counting task failed: {0}")]
|
||||
Task(String),
|
||||
}
|
||||
|
|
|
|||
41
litellm-rust/crates/token-counter/src/fast.rs
Normal file
41
litellm-rust/crates/token-counter/src/fast.rs
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
use litellm_token_counter_fast::Error as BackendError;
|
||||
pub use litellm_token_counter_fast::FastTokenizer;
|
||||
|
||||
use crate::{Error, TokenCounter, Tokenizer};
|
||||
|
||||
impl TokenCounter {
|
||||
pub fn from_json_fast(tokenizer_json: &str) -> Result<Self, Error> {
|
||||
FastTokenizer::from_json(tokenizer_json)
|
||||
.map(Self::new)
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
pub fn from_cl100k_ranks(rank_file: &str) -> Result<Self, Error> {
|
||||
FastTokenizer::from_cl100k_ranks(rank_file)
|
||||
.map(Self::new)
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
pub fn from_o200k_ranks(rank_file: &str) -> Result<Self, Error> {
|
||||
FastTokenizer::from_o200k_ranks(rank_file)
|
||||
.map(Self::new)
|
||||
.map_err(Error::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl Tokenizer for FastTokenizer {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
FastTokenizer::count_tokens(self, text).map_err(Error::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BackendError> for Error {
|
||||
fn from(error: BackendError) -> Self {
|
||||
match error {
|
||||
BackendError::Load(source) => Self::Load(source),
|
||||
BackendError::Ranks(message) => Self::Ranks(message),
|
||||
BackendError::UnicodeClasses => Self::UnicodeClasses,
|
||||
BackendError::Encode(source) => Self::Encode(source),
|
||||
}
|
||||
}
|
||||
}
|
||||
27
litellm-rust/crates/token-counter/src/huggingface.rs
Normal file
27
litellm-rust/crates/token-counter/src/huggingface.rs
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
use litellm_token_counter_huggingface::Error as BackendError;
|
||||
pub use litellm_token_counter_huggingface::HuggingFaceTokenizer;
|
||||
|
||||
use crate::{Error, TokenCounter, Tokenizer};
|
||||
|
||||
impl TokenCounter {
|
||||
pub fn from_json(tokenizer_json: &str) -> Result<Self, Error> {
|
||||
HuggingFaceTokenizer::from_json(tokenizer_json)
|
||||
.map(Self::new)
|
||||
.map_err(Error::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl Tokenizer for HuggingFaceTokenizer {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
HuggingFaceTokenizer::count_tokens(self, text).map_err(Error::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<BackendError> for Error {
|
||||
fn from(error: BackendError) -> Self {
|
||||
match error {
|
||||
BackendError::Load(source) => Self::Load(source),
|
||||
BackendError::Encode(source) => Self::Encode(source),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -4,18 +4,21 @@
|
|||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
mod byte_level;
|
||||
mod cl100k;
|
||||
mod counter;
|
||||
mod error;
|
||||
mod o200k;
|
||||
mod python_json;
|
||||
mod scanner;
|
||||
mod tiktoken;
|
||||
mod tokenizer;
|
||||
mod tools;
|
||||
mod types;
|
||||
mod unicode_classes;
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
pub mod fast;
|
||||
#[cfg(feature = "huggingface")]
|
||||
pub mod huggingface;
|
||||
#[cfg(feature = "tiktoken")]
|
||||
pub mod tiktoken;
|
||||
|
||||
pub use counter::{InputTokenCount, TokenCounter};
|
||||
pub use error::Error;
|
||||
pub use tokenizer::Tokenizer;
|
||||
pub use types::CountableRequest;
|
||||
|
|
|
|||
|
|
@ -1,222 +1,24 @@
|
|||
//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and
|
||||
//! the merge loop that turns one regex piece into tokens. The merge order is
|
||||
//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is
|
||||
//! identical, but pairs are tracked in a heap so a long piece costs
|
||||
//! `O(n log n)` instead of tiktoken's `O(n^2)`.
|
||||
pub use litellm_token_counter_tiktoken::TiktokenTokenizer;
|
||||
use litellm_token_counter_tiktoken::UnsupportedTokenizer;
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::BinaryHeap;
|
||||
use crate::{Error, TokenCounter, Tokenizer};
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use rustc_hash::FxHashMap;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
type Rank = u32;
|
||||
|
||||
const NO_RANK: Rank = Rank::MAX;
|
||||
const END: usize = usize::MAX;
|
||||
|
||||
pub(super) struct MergeRanks(FxHashMap<Box<[u8]>, Rank>);
|
||||
|
||||
impl MergeRanks {
|
||||
pub(super) fn parse(text: &str) -> Result<Self, Error> {
|
||||
let ranks = text
|
||||
.lines()
|
||||
.filter(|line| !line.is_empty())
|
||||
.map(parse_line)
|
||||
.collect::<Result<FxHashMap<_, _>, _>>()?;
|
||||
if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) {
|
||||
return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token")));
|
||||
}
|
||||
Ok(Self(ranks))
|
||||
}
|
||||
|
||||
fn rank(&self, bytes: &[u8]) -> Rank {
|
||||
self.0.get(bytes).copied().unwrap_or(NO_RANK)
|
||||
}
|
||||
|
||||
/// Token count of one regex piece, as `encode_ordinary` would produce.
|
||||
pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize {
|
||||
if piece.len() < 2 || self.0.contains_key(piece) {
|
||||
return 1;
|
||||
}
|
||||
scratch.reset(piece.len());
|
||||
for start in 0..piece.len() - 1 {
|
||||
scratch.set_rank(start, self.rank(&piece[start..start + 2]));
|
||||
}
|
||||
let mut parts = piece.len();
|
||||
while let Some(Reverse((rank, start))) = scratch.heap.pop() {
|
||||
if scratch.next[start] == END || scratch.rank[start] != rank {
|
||||
continue;
|
||||
}
|
||||
let merged = scratch.next[start];
|
||||
let after = scratch.next[merged];
|
||||
scratch.next[merged] = END;
|
||||
scratch.next[start] = after;
|
||||
parts -= 1;
|
||||
if after < piece.len() {
|
||||
scratch.prev[after] = start;
|
||||
scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)]));
|
||||
} else {
|
||||
scratch.rank[start] = NO_RANK;
|
||||
}
|
||||
let before = scratch.prev[start];
|
||||
if before != END {
|
||||
scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)]));
|
||||
}
|
||||
}
|
||||
parts
|
||||
impl TokenCounter {
|
||||
pub fn from_tiktoken(encoding: &str) -> Result<Self, Error> {
|
||||
TiktokenTokenizer::from_name(encoding)
|
||||
.map(Self::new)
|
||||
.map_err(Error::from)
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> {
|
||||
let (token, rank) = line
|
||||
.split_once(' ')
|
||||
.ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?;
|
||||
let bytes = STANDARD
|
||||
.decode(token)
|
||||
.map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?;
|
||||
let rank = rank
|
||||
.parse()
|
||||
.map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?;
|
||||
Ok((bytes.into_boxed_slice(), rank))
|
||||
}
|
||||
|
||||
/// Buffers reused across the pieces of one text. Parts are addressed by the
|
||||
/// byte offset they start at, which also gives the leftmost-pair tie break.
|
||||
#[derive(Default)]
|
||||
pub(super) struct MergeScratch {
|
||||
next: Vec<usize>,
|
||||
prev: Vec<usize>,
|
||||
rank: Vec<Rank>,
|
||||
heap: BinaryHeap<Reverse<(Rank, usize)>>,
|
||||
}
|
||||
|
||||
impl MergeScratch {
|
||||
fn reset(&mut self, len: usize) {
|
||||
self.next.clear();
|
||||
self.next.extend(1..=len);
|
||||
self.prev.clear();
|
||||
self.prev.push(END);
|
||||
self.prev.extend(0..len - 1);
|
||||
self.rank.clear();
|
||||
self.rank.resize(len, NO_RANK);
|
||||
self.heap.clear();
|
||||
}
|
||||
|
||||
fn end(&self, start: usize) -> usize {
|
||||
self.next[start]
|
||||
}
|
||||
|
||||
fn set_rank(&mut self, start: usize, rank: Rank) {
|
||||
self.rank[start] = rank;
|
||||
if rank != NO_RANK {
|
||||
self.heap.push(Reverse((rank, start)));
|
||||
}
|
||||
impl Tokenizer for TiktokenTokenizer {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
Ok(TiktokenTokenizer::count_tokens(self, text))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn ranks() -> MergeRanks {
|
||||
let path = concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
|
||||
);
|
||||
MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo"))
|
||||
.expect("rank file parses")
|
||||
}
|
||||
|
||||
/// tiktoken's `_byte_pair_merge`, transcribed, as the reference.
|
||||
fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize {
|
||||
if piece.len() < 2 || ranks.0.contains_key(piece) {
|
||||
return 1;
|
||||
}
|
||||
let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1)
|
||||
.map(|index| (index, ranks.rank(&piece[index..index + 2])))
|
||||
.chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)])
|
||||
.collect();
|
||||
let get_rank = |parts: &[(usize, Rank)], index: usize| {
|
||||
if index + 3 < parts.len() {
|
||||
ranks.rank(&piece[parts[index].0..parts[index + 3].0])
|
||||
} else {
|
||||
NO_RANK
|
||||
}
|
||||
};
|
||||
loop {
|
||||
let Some(index) = parts[..parts.len() - 1]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, (_, rank))| *rank != NO_RANK)
|
||||
.min_by_key(|(index, (_, rank))| (*rank, *index))
|
||||
.map(|(index, _)| index)
|
||||
else {
|
||||
return parts.len() - 1;
|
||||
};
|
||||
if index > 0 {
|
||||
parts[index - 1].1 = get_rank(&parts, index - 1);
|
||||
}
|
||||
parts[index].1 = get_rank(&parts, index);
|
||||
parts.remove(index + 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_byte_is_a_token() {
|
||||
let ranks = ranks();
|
||||
assert_eq!(ranks.0.len(), 100_256);
|
||||
assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heap_merge_matches_tiktokens_merge_loop() {
|
||||
let ranks = ranks();
|
||||
let mut scratch = MergeScratch::default();
|
||||
let mut rng = StdRng::seed_from_u64(99);
|
||||
let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123";
|
||||
for _ in 0..20_000 {
|
||||
let piece: Vec<u8> = (0..rng.gen_range(1..24))
|
||||
.map(|_| alphabet[rng.gen_range(0..alphabet.len())])
|
||||
.collect();
|
||||
assert_eq!(
|
||||
ranks.count_piece(&piece, &mut scratch),
|
||||
reference_count(&ranks, &piece),
|
||||
"piece {:?}",
|
||||
String::from_utf8_lossy(&piece)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_repeated_runs_cost_close_to_linear() {
|
||||
let ranks = ranks();
|
||||
let mut scratch = MergeScratch::default();
|
||||
let mut time = |len: usize| {
|
||||
let piece = vec![b' '; len];
|
||||
let started = std::time::Instant::now();
|
||||
assert!(ranks.count_piece(&piece, &mut scratch) > 0);
|
||||
started.elapsed()
|
||||
};
|
||||
let small = (0..3).map(|_| time(1 << 14)).min().unwrap();
|
||||
let large = time(1 << 18);
|
||||
assert!(
|
||||
large < small * 64,
|
||||
"{small:?} for 2^14 bytes, {large:?} for 2^18"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_rank_files_are_rejected() {
|
||||
assert!(MergeRanks::parse("IQ==").is_err());
|
||||
assert!(MergeRanks::parse("IQ== x").is_err());
|
||||
assert!(MergeRanks::parse("!!! 1").is_err());
|
||||
assert!(MergeRanks::parse("IQ== 1").is_err());
|
||||
impl From<UnsupportedTokenizer> for Error {
|
||||
fn from(error: UnsupportedTokenizer) -> Self {
|
||||
Self::UnsupportedTokenizer(error.0)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
31
litellm-rust/crates/token-counter/src/tokenizer.rs
Normal file
31
litellm-rust/crates/token-counter/src/tokenizer.rs
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
use crate::Error;
|
||||
|
||||
pub trait Tokenizer: Send + Sync {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{CountableRequest, TokenCounter};
|
||||
|
||||
struct Characters;
|
||||
|
||||
impl Tokenizer for Characters {
|
||||
fn count_tokens(&self, text: &str) -> Result<usize, Error> {
|
||||
Ok(text.chars().count())
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_accounting_works_with_an_injected_backend() {
|
||||
let counter = TokenCounter::new(Characters);
|
||||
let request =
|
||||
CountableRequest::parse(br#"{"messages":[{"role":"user","content":"hello"}]}"#)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
counter.count_request(&request).unwrap().input_tokens,
|
||||
3 + 4 + 5 + 3
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
#![cfg(feature = "huggingface")]
|
||||
|
||||
use rstest::rstest;
|
||||
use serde::Deserialize;
|
||||
|
||||
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
|
||||
|
||||
|
|
@ -176,146 +177,155 @@ fn loading_a_bad_tokenizer_is_a_load_error() {
|
|||
assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
|
||||
}
|
||||
|
||||
/// A tiktoken encoding: its fixture directory, the vendored rank file Python
|
||||
/// loads, the constructor, and the model `generate.py` counted the requests for.
|
||||
#[derive(Clone, Copy)]
|
||||
struct TiktokenEncoding {
|
||||
fixtures: &'static str,
|
||||
rank_file: &'static str,
|
||||
load: fn(&str) -> Result<TokenCounter, Error>,
|
||||
model: &'static str,
|
||||
}
|
||||
#[cfg(feature = "fast")]
|
||||
mod fast {
|
||||
use super::*;
|
||||
use serde::Deserialize;
|
||||
|
||||
const CL100K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "cl100k",
|
||||
rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4",
|
||||
load: TokenCounter::from_cl100k_ranks,
|
||||
model: "gpt-4",
|
||||
};
|
||||
/// A tiktoken encoding: its fixture directory, the vendored rank file Python
|
||||
/// loads, the constructor, and the model `generate.py` counted the requests for.
|
||||
#[derive(Clone, Copy)]
|
||||
struct TiktokenEncoding {
|
||||
fixtures: &'static str,
|
||||
rank_file: &'static str,
|
||||
load: fn(&str) -> Result<TokenCounter, Error>,
|
||||
model: &'static str,
|
||||
}
|
||||
|
||||
const O200K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "o200k",
|
||||
rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790",
|
||||
load: TokenCounter::from_o200k_ranks,
|
||||
model: "gpt-4o",
|
||||
};
|
||||
|
||||
fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter {
|
||||
let path = format!(
|
||||
"{}/../../../litellm/litellm_core_utils/tokenizers/{}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
encoding.rank_file
|
||||
);
|
||||
let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo");
|
||||
(encoding.load)(&ranks).expect("ranks load")
|
||||
}
|
||||
|
||||
fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String {
|
||||
let path = format!(
|
||||
"{}/tests/fixtures/{}/{name}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
encoding.fixtures
|
||||
);
|
||||
std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py")
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TextFixture {
|
||||
text: String,
|
||||
tokens: usize,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RequestFixture {
|
||||
body: String,
|
||||
input_tokens: usize,
|
||||
}
|
||||
|
||||
/// Reference counts come from `tiktoken.get_encoding(name)`; see
|
||||
/// `tests/fixtures/generate.py`.
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let fixtures: Vec<TextFixture> = tiktoken_fixture(encoding, "texts.jsonl")
|
||||
.lines()
|
||||
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
|
||||
.collect();
|
||||
assert!(fixtures.len() > 3000);
|
||||
let mismatches: Vec<_> = fixtures
|
||||
.iter()
|
||||
.filter_map(|fixture| {
|
||||
let count = counter.count_text(&fixture.text).expect("text counts");
|
||||
(count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count))
|
||||
})
|
||||
.collect();
|
||||
assert!(
|
||||
mismatches.is_empty(),
|
||||
"(text, tiktoken, rust): {mismatches:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Reference counts come from the proxy's admission counter
|
||||
/// (`_count_input_tokens(body, model)`), so this pins the shared message,
|
||||
/// tool and reply-priming accounting on the tiktoken paths as well.
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let fixtures: Vec<RequestFixture> = tiktoken_fixture(encoding, "requests.jsonl")
|
||||
.lines()
|
||||
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
|
||||
.collect();
|
||||
let counts: Vec<usize> = fixtures
|
||||
.iter()
|
||||
.map(|fixture| {
|
||||
let request = CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses");
|
||||
let count = counter.count_request(&request).expect("fixture counts");
|
||||
assert_eq!(count.model.as_deref(), Some(encoding.model));
|
||||
assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body);
|
||||
count.input_tokens
|
||||
})
|
||||
.collect();
|
||||
assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_shares_the_message_accounting_with_the_anthropic_path(
|
||||
#[case] encoding: TiktokenEncoding,
|
||||
) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let count = |body: &str| {
|
||||
counter
|
||||
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
|
||||
.expect("counts")
|
||||
.input_tokens
|
||||
const CL100K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "cl100k",
|
||||
rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4",
|
||||
load: TokenCounter::from_cl100k_ranks,
|
||||
model: "gpt-4",
|
||||
};
|
||||
let text = |text: &str| counter.count_text(text).expect("counts");
|
||||
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
|
||||
assert_eq!(base, 3 + text("user") + text("hi") + 3);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#),
|
||||
base + text("al") + 1
|
||||
);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#),
|
||||
base + 1
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty("")]
|
||||
#[case::not_base64("!!!! 0")]
|
||||
#[case::missing_rank("YQ==")]
|
||||
#[case::rank_not_a_number("YQ== x")]
|
||||
#[case::single_byte_tokens_missing("YWI= 0")]
|
||||
fn loading_a_bad_rank_file_is_a_load_error(
|
||||
#[case] rank_file: &str,
|
||||
#[values(CL100K, O200K)] encoding: TiktokenEncoding,
|
||||
) {
|
||||
assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_))));
|
||||
const O200K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "o200k",
|
||||
rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790",
|
||||
load: TokenCounter::from_o200k_ranks,
|
||||
model: "gpt-4o",
|
||||
};
|
||||
|
||||
fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter {
|
||||
let path = format!(
|
||||
"{}/../../../litellm/litellm_core_utils/tokenizers/{}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
encoding.rank_file
|
||||
);
|
||||
let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo");
|
||||
(encoding.load)(&ranks).expect("ranks load")
|
||||
}
|
||||
|
||||
fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String {
|
||||
let path = format!(
|
||||
"{}/../token-counter-fast/tests/fixtures/{}/{name}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
encoding.fixtures
|
||||
);
|
||||
std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py")
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TextFixture {
|
||||
text: String,
|
||||
tokens: usize,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RequestFixture {
|
||||
body: String,
|
||||
input_tokens: usize,
|
||||
}
|
||||
|
||||
/// Reference counts come from `tiktoken.get_encoding(name)`; see
|
||||
/// `tests/fixtures/generate.py`.
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let fixtures: Vec<TextFixture> = tiktoken_fixture(encoding, "texts.jsonl")
|
||||
.lines()
|
||||
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
|
||||
.collect();
|
||||
assert!(fixtures.len() > 3000);
|
||||
let mismatches: Vec<_> = fixtures
|
||||
.iter()
|
||||
.filter_map(|fixture| {
|
||||
let count = counter.count_text(&fixture.text).expect("text counts");
|
||||
(count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count))
|
||||
})
|
||||
.collect();
|
||||
assert!(
|
||||
mismatches.is_empty(),
|
||||
"(text, tiktoken, rust): {mismatches:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Reference counts come from the proxy's admission counter
|
||||
/// (`_count_input_tokens(body, model)`), so this pins the shared message,
|
||||
/// tool and reply-priming accounting on the tiktoken paths as well.
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let fixtures: Vec<RequestFixture> = tiktoken_fixture(encoding, "requests.jsonl")
|
||||
.lines()
|
||||
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
|
||||
.collect();
|
||||
let counts: Vec<usize> = fixtures
|
||||
.iter()
|
||||
.map(|fixture| {
|
||||
let request =
|
||||
CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses");
|
||||
let count = counter.count_request(&request).expect("fixture counts");
|
||||
assert_eq!(count.model.as_deref(), Some(encoding.model));
|
||||
assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body);
|
||||
count.input_tokens
|
||||
})
|
||||
.collect();
|
||||
assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::cl100k(CL100K)]
|
||||
#[case::o200k(O200K)]
|
||||
fn tiktoken_shares_the_message_accounting_with_the_anthropic_path(
|
||||
#[case] encoding: TiktokenEncoding,
|
||||
) {
|
||||
let counter = tiktoken_counter(encoding);
|
||||
let count = |body: &str| {
|
||||
counter
|
||||
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
|
||||
.expect("counts")
|
||||
.input_tokens
|
||||
};
|
||||
let text = |text: &str| counter.count_text(text).expect("counts");
|
||||
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
|
||||
assert_eq!(base, 3 + text("user") + text("hi") + 3);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#),
|
||||
base + text("al") + 1
|
||||
);
|
||||
assert_eq!(
|
||||
count(
|
||||
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#
|
||||
),
|
||||
base + 1
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty("")]
|
||||
#[case::not_base64("!!!! 0")]
|
||||
#[case::missing_rank("YQ==")]
|
||||
#[case::rank_not_a_number("YQ== x")]
|
||||
#[case::single_byte_tokens_missing("YWI= 0")]
|
||||
fn loading_a_bad_rank_file_is_a_load_error(
|
||||
#[case] rank_file: &str,
|
||||
#[values(CL100K, O200K)] encoding: TiktokenEncoding,
|
||||
) {
|
||||
assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_))));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -97,9 +97,9 @@ class ResponsesWebSocketConnection:
|
|||
class TokenCounter:
|
||||
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...
|
||||
@staticmethod
|
||||
def from_cl100k_ranks(rank_file: str) -> TokenCounter: ...
|
||||
def from_json_fast(tokenizer_json: str) -> TokenCounter: ...
|
||||
@staticmethod
|
||||
def from_o200k_ranks(rank_file: str) -> TokenCounter: ...
|
||||
def from_tiktoken(encoding: str) -> TokenCounter: ...
|
||||
def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...
|
||||
|
||||
def gil_stats() -> dict[str, int]: ...
|
||||
|
|
|
|||
|
|
@ -11,14 +11,21 @@ from pydantic import TypeAdapter
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.configuration import rust_enabled
|
||||
from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt
|
||||
from litellm.utils import claude_json_str, huggingface_tokenizer_kind
|
||||
|
||||
RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"]
|
||||
RustTokenizer = Literal[
|
||||
"anthropic",
|
||||
"cl100k_base",
|
||||
"o200k_base",
|
||||
"o200k_harmony",
|
||||
"p50k_base",
|
||||
"p50k_edit",
|
||||
"r50k_base",
|
||||
]
|
||||
|
||||
|
||||
class RustTokenCounter(Protocol):
|
||||
|
|
@ -30,10 +37,7 @@ class RustTokenCounterFactory(Protocol):
|
|||
def __call__(self, tokenizer_json: str) -> RustTokenCounter:
|
||||
raise NotImplementedError
|
||||
|
||||
def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter:
|
||||
raise NotImplementedError
|
||||
|
||||
def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter:
|
||||
def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
|
|
@ -63,9 +67,8 @@ def rust_tokenizer(model: str) -> RustTokenizer | None:
|
|||
"""The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count.
|
||||
|
||||
Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace
|
||||
downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust
|
||||
prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in
|
||||
Python."""
|
||||
downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every
|
||||
message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python."""
|
||||
if litellm.disable_token_counter is True:
|
||||
return None
|
||||
kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model)
|
||||
|
|
@ -73,24 +76,19 @@ def rust_tokenizer(model: str) -> RustTokenizer | None:
|
|||
return "anthropic"
|
||||
if kind is not None or uses_legacy_message_accounting(model):
|
||||
return None
|
||||
match openai_tokenizer_encoding(model).name:
|
||||
case "cl100k_base":
|
||||
return "cl100k_base"
|
||||
case "o200k_base":
|
||||
return "o200k_base"
|
||||
case _:
|
||||
return None
|
||||
encoding: Final = openai_tokenizer_encoding(model).name
|
||||
if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}:
|
||||
return cast(RustTokenizer, encoding)
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
@lru_cache(maxsize=8)
|
||||
def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter:
|
||||
match tokenizer:
|
||||
case "anthropic":
|
||||
return factory(claude_json_str)
|
||||
case "cl100k_base":
|
||||
return factory.from_cl100k_ranks(cl100k_base_rank_file())
|
||||
case "o200k_base":
|
||||
return factory.from_o200k_ranks(o200k_base_rank_file())
|
||||
case _:
|
||||
return factory.from_tiktoken(tokenizer)
|
||||
|
||||
|
||||
async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
|
||||
|
|
|
|||
|
|
@ -256,7 +256,7 @@ class _RecordingCounter:
|
|||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`."""
|
||||
"""Stands in for the native `TokenCounter` class."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
|
||||
|
|
@ -264,11 +264,8 @@ class _RecordingFactory:
|
|||
def __call__(self, tokenizer_json: str) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, "anthropic")
|
||||
|
||||
def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, "cl100k_base")
|
||||
|
||||
def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, "o200k_base")
|
||||
def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, encoding)
|
||||
|
||||
|
||||
class _DecliningCounter:
|
||||
|
|
@ -280,10 +277,7 @@ class _DecliningFactory:
|
|||
def __call__(self, tokenizer_json: str) -> _DecliningCounter:
|
||||
return _DecliningCounter()
|
||||
|
||||
def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter:
|
||||
return _DecliningCounter()
|
||||
|
||||
def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter:
|
||||
def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter:
|
||||
return _DecliningCounter()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ cases need the extension and are skipped when it is not built.
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -27,7 +26,6 @@ MODEL: Final = "claude-sonnet-4-5-20250929"
|
|||
CL100K_MODEL: Final = "gpt-4"
|
||||
O200K_MODEL: Final = "gpt-4o"
|
||||
TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base")
|
||||
RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998})
|
||||
BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode()
|
||||
|
||||
|
||||
|
|
@ -55,24 +53,20 @@ class _RecordingCounter:
|
|||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files."""
|
||||
"""Stands in for the native `TokenCounter` class."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counters: list[_RecordingCounter] = []
|
||||
self.rank_files: list[str] = []
|
||||
self.encodings: list[str] = []
|
||||
|
||||
def __call__(self, tokenizer_json: str) -> _RecordingCounter:
|
||||
counter = _RecordingCounter(tokenizer_json)
|
||||
self.counters.append(counter)
|
||||
return counter
|
||||
|
||||
def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
|
||||
self.rank_files.append(rank_file)
|
||||
return self("cl100k_base")
|
||||
|
||||
def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
|
||||
self.rank_files.append(rank_file)
|
||||
return self("o200k_base")
|
||||
def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter:
|
||||
self.encodings.append(encoding)
|
||||
return self(encoding)
|
||||
|
||||
|
||||
class _RaisingCounter:
|
||||
|
|
@ -92,10 +86,7 @@ class _RaisingFactory:
|
|||
def __call__(self, tokenizer_json: str) -> _RaisingCounter:
|
||||
return _RaisingCounter(self.error)
|
||||
|
||||
def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter:
|
||||
return _RaisingCounter(self.error)
|
||||
|
||||
def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter:
|
||||
def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter:
|
||||
return _RaisingCounter(self.error)
|
||||
|
||||
|
||||
|
|
@ -140,7 +131,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base"))
|
||||
async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None:
|
||||
async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
|
@ -149,9 +140,7 @@ async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokeni
|
|||
second: Final = await bridge.count_input_tokens(BODY, tokenizer)
|
||||
|
||||
assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42)
|
||||
assert len(factory.rank_files) == 1
|
||||
assert factory.rank_files[0].startswith("IQ== 0\n")
|
||||
assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer]
|
||||
assert factory.encodings == [tokenizer]
|
||||
assert factory.counters[0].tokenizer_json == tokenizer
|
||||
assert factory.counters[0].bodies == [BODY, BODY]
|
||||
|
||||
|
|
@ -235,13 +224,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected:
|
|||
("model", "python_encoding"),
|
||||
(("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")),
|
||||
)
|
||||
def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have(
|
||||
monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str
|
||||
def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust(
|
||||
monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model})
|
||||
|
||||
assert openai_tokenizer_encoding(model).name == python_encoding
|
||||
assert bridge.rust_tokenizer(model) is None
|
||||
assert bridge.rust_tokenizer(model) == python_encoding
|
||||
|
||||
|
||||
def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue