mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42165 from BerriAI/refactor-tokenizer-rs
refactor(rust): split token counter backends
This commit is contained in:
yujonglee
GitHub
commit
65c6495616
38 changed files with 1205 additions and 560 deletions
7
.github/workflows/test-rust.yml
vendored
7
.github/workflows/test-rust.yml
vendored
|
|
@ -120,6 +120,13 @@ jobs:
|
|||
|
||||
- run: cargo test --workspace --doc --locked
|
||||
|
||||
- name: Test token counter feature combinations
|
||||
run: |
|
||||
for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do
|
||||
cargo test -p litellm-token-counter --locked --no-default-features --features "$features"
|
||||
cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}"
|
||||
done
|
||||
|
||||
rust-wheel:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
|
|
|||
84
litellm-rust/Cargo.lock
generated
84
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,25 @@ 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",
|
||||
"tokenizers",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-token-counter-fast"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"rustc-hash",
|
||||
|
|
@ -2226,6 +2275,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 +3859,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,8 @@ 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"] }
|
||||
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", "fast"]
|
||||
abi3 = ["pyo3/abi3-py310"]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
panic-test = []
|
||||
fast = ["litellm-token-counter/fast"]
|
||||
huggingface = ["litellm-token-counter/huggingface"]
|
||||
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,11 @@
|
|||
use std::{num::NonZero, sync::Arc, thread::available_parallelism};
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_host_python::{release_gil, run_async};
|
||||
#[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,
|
||||
};
|
||||
|
|
@ -28,17 +33,66 @@ 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 = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
|
||||
}
|
||||
#[cfg(all(not(feature = "fast"), feature = "huggingface"))]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
}
|
||||
#[cfg(not(any(feature = "fast", feature = "huggingface")))]
|
||||
{
|
||||
let _ = (py, tokenizer_json);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast or huggingface feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
#[cfg(feature = "fast")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
}
|
||||
#[cfg(not(feature = "fast"))]
|
||||
{
|
||||
let _ = (py, rank_file);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the fast feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
{
|
||||
Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
|
||||
}
|
||||
#[cfg(not(feature = "tiktoken"))]
|
||||
{
|
||||
let _ = (py, encoding);
|
||||
Err(RustBridgeDeclined::new_err(
|
||||
"tokenizer backend requires the tiktoken feature",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
|
|
@ -62,6 +116,7 @@ impl TokenCounter {
|
|||
}
|
||||
|
||||
impl TokenCounter {
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn load(
|
||||
py: Python<'_>,
|
||||
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
|
||||
|
|
@ -74,6 +129,7 @@ impl TokenCounter {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(1, NonZero::get)
|
||||
}
|
||||
|
|
@ -86,7 +142,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.workspace = true
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
240
litellm-rust/crates/token-counter-fast/src/tiktoken.rs
Normal file
240
litellm-rust/crates/token-counter-fast/src/tiktoken.rs
Normal file
|
|
@ -0,0 +1,240 @@
|
|||
//! 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}")))?;
|
||||
if rank == NO_RANK {
|
||||
return Err(Error::Ranks(format!("rank {rank} is reserved")));
|
||||
}
|
||||
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 reserved_merge_rank_is_rejected() {
|
||||
let bytes = (0..=u8::MAX)
|
||||
.map(|byte| format!("{} {byte}\n", STANDARD.encode([byte])))
|
||||
.collect::<String>();
|
||||
let rank_file = format!("{bytes}{} {NO_RANK}\n", STANDARD.encode(b"ab"));
|
||||
assert!(matches!(
|
||||
MergeRanks::parse(&rank_file),
|
||||
Err(Error::Ranks(_))
|
||||
));
|
||||
let valid_rank_file = format!("{bytes}{} {}\n", STANDARD.encode(b"ab"), NO_RANK - 1);
|
||||
let ranks = MergeRanks::parse(&valid_rank_file).unwrap();
|
||||
assert_eq!(ranks.count_piece(b"aab", &mut MergeScratch::default()), 2);
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
|
|
@ -2,8 +2,8 @@
|
|||
|
||||
Run from the repository root with the project environment, once per encoding:
|
||||
|
||||
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py cl100k_base
|
||||
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py o200k_base
|
||||
uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py cl100k_base
|
||||
uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py o200k_base
|
||||
|
||||
`<encoding>/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens`
|
||||
counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`,
|
||||
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.workspace = true
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
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),
|
||||
}
|
||||
23
litellm-rust/crates/token-counter-huggingface/src/lib.rs
Normal file
23
litellm-rust/crates/token-counter-huggingface/src/lib.rs
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod error;
|
||||
|
||||
pub use error::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
|
||||
5
litellm-rust/crates/token-counter-tiktoken/src/error.rs
Normal file
5
litellm-rust/crates/token-counter-tiktoken/src/error.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
use thiserror::Error as ThisError;
|
||||
|
||||
#[derive(Debug, ThisError)]
|
||||
#[error("unsupported tokenizer: {0}")]
|
||||
pub struct UnsupportedTokenizer(pub String);
|
||||
70
litellm-rust/crates/token-counter-tiktoken/src/lib.rs
Normal file
70
litellm-rust/crates/token-counter-tiktoken/src/lib.rs
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod error;
|
||||
|
||||
pub use error::UnsupportedTokenizer;
|
||||
|
||||
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 named_encodings_match_their_reference_counts() {
|
||||
let encodings = [
|
||||
("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", tiktoken_rs::r50k_base_singleton()),
|
||||
("gpt2", tiktoken_rs::r50k_base_singleton()),
|
||||
];
|
||||
let texts = [
|
||||
"",
|
||||
"Hello, how are you today?",
|
||||
"é e\u{301} 漢字 ع ३ 🙂 ＡfiⅣ",
|
||||
" def function():\n return 123456789\r\n",
|
||||
"<|endoftext|><|fim_prefix|><|start|>assistant<|message|>",
|
||||
];
|
||||
for (name, reference) in encodings {
|
||||
let counter = TiktokenTokenizer::from_name(name).unwrap();
|
||||
for text in texts {
|
||||
assert_eq!(
|
||||
counter.count_tokens(text),
|
||||
reference.encode_ordinary(text).len(),
|
||||
"{name}: {text:?}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_encoding_preserves_its_name() {
|
||||
let Err(UnsupportedTokenizer(name)) = TiktokenTokenizer::from_name("unknown-encoding")
|
||||
else {
|
||||
panic!("unknown encoding must be rejected");
|
||||
};
|
||||
assert_eq!(name, "unknown-encoding");
|
||||
}
|
||||
}
|
||||
|
|
@ -5,26 +5,34 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
default = ["fast", "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
|
||||
rand.workspace = true
|
||||
rstest.workspace = true
|
||||
tokenizers.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` uses 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`
|
||||
|
||||
All three backends are enabled by default. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. 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,22 +1,30 @@
|
|||
use rstest::rstest;
|
||||
use serde::Deserialize;
|
||||
|
||||
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
|
||||
#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
|
||||
use litellm_token_counter::TokenCounter;
|
||||
use litellm_token_counter::{CountableRequest, Error};
|
||||
|
||||
/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)`
|
||||
/// so this test also guards Python parity.
|
||||
fn counter() -> TokenCounter {
|
||||
let path = concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
|
||||
);
|
||||
let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo");
|
||||
TokenCounter::from_json(&json).expect("anthropic tokenizer loads")
|
||||
}
|
||||
#[cfg(any(feature = "fast", feature = "huggingface"))]
|
||||
mod json {
|
||||
use super::*;
|
||||
use litellm_token_counter::InputTokenCount;
|
||||
|
||||
const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#;
|
||||
/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)`
|
||||
/// so this test also guards Python parity.
|
||||
type JsonLoader = fn(&str) -> Result<TokenCounter, Error>;
|
||||
|
||||
const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[
|
||||
fn counter(load: JsonLoader) -> TokenCounter {
|
||||
let path = concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
|
||||
);
|
||||
let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo");
|
||||
load(&json).expect("anthropic tokenizer loads")
|
||||
}
|
||||
|
||||
const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#;
|
||||
|
||||
const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[
|
||||
{"role":"system","content":"You are a terse assistant."},
|
||||
{"role":"user","name":"alice","content":[
|
||||
{"type":"text","text":"Summarise this paragraph about ships and harbours."},
|
||||
|
|
@ -25,7 +33,7 @@ const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[
|
|||
{"type":"tool_reference","tool_name":"get_weather"}]},
|
||||
{"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#;
|
||||
|
||||
const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}],
|
||||
const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}],
|
||||
"tools":[
|
||||
{"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{
|
||||
"type":"object",
|
||||
|
|
@ -40,61 +48,188 @@ const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"
|
|||
{"type":"function","function":{"name":"noop"}}],
|
||||
"tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#;
|
||||
|
||||
const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5",
|
||||
const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5",
|
||||
"messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}],
|
||||
"tools":[{"name":"get_weather","description":"Get weather","input_schema":{
|
||||
"type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}],
|
||||
"tool_choice":"none"}"#;
|
||||
|
||||
const COMPLETIONS_PROMPT: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#;
|
||||
const COMPLETIONS_PROMPT: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#;
|
||||
|
||||
const COMPLETIONS_PROMPT_LIST: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#;
|
||||
const COMPLETIONS_PROMPT_LIST: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#;
|
||||
|
||||
const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[
|
||||
const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[
|
||||
{"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]},
|
||||
{"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#;
|
||||
|
||||
const EMBEDDINGS_TOKEN_IDS: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#;
|
||||
const EMBEDDINGS_TOKEN_IDS: &str =
|
||||
r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#;
|
||||
|
||||
const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
|
||||
const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
|
||||
"documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#;
|
||||
|
||||
/// Expected counts are pinned from
|
||||
/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`.
|
||||
#[rstest]
|
||||
#[case::text_only(SIMPLE, 14)]
|
||||
#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)]
|
||||
#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)]
|
||||
#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)]
|
||||
#[case::completions_prompt(COMPLETIONS_PROMPT, 7)]
|
||||
#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)]
|
||||
#[case::responses_input_items(RESPONSES_INPUT, 62)]
|
||||
#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)]
|
||||
#[case::rerank_query_and_documents(RERANK, 41)]
|
||||
fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) {
|
||||
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
|
||||
let count = counter().count_request(&request).expect("fixture counts");
|
||||
assert_eq!(
|
||||
count,
|
||||
InputTokenCount {
|
||||
model: Some("claude-sonnet-4-5".to_string()),
|
||||
input_tokens: expected,
|
||||
}
|
||||
);
|
||||
}
|
||||
fn assert_count_request_matches_python_token_counter(
|
||||
load: JsonLoader,
|
||||
body: &str,
|
||||
expected: usize,
|
||||
) {
|
||||
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
|
||||
let count = counter(load)
|
||||
.count_request(&request)
|
||||
.expect("fixture counts");
|
||||
assert_eq!(
|
||||
count,
|
||||
InputTokenCount {
|
||||
model: Some("claude-sonnet-4-5".to_string()),
|
||||
input_tokens: expected,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)]
|
||||
#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)]
|
||||
#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)]
|
||||
#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)]
|
||||
fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) {
|
||||
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
|
||||
let count = counter().count_request(&request).expect("fixture counts");
|
||||
assert_eq!(count.input_tokens, expected);
|
||||
fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) {
|
||||
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
|
||||
let count = counter(load)
|
||||
.count_request(&request)
|
||||
.expect("fixture counts");
|
||||
assert_eq!(count.input_tokens, expected);
|
||||
}
|
||||
|
||||
fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) {
|
||||
let request = CountableRequest::parse(body).expect("shape parses");
|
||||
assert!(matches!(
|
||||
counter(load).count_request(&request),
|
||||
Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems)
|
||||
));
|
||||
}
|
||||
|
||||
fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) {
|
||||
let counter = counter(load);
|
||||
let count = |body: &str| {
|
||||
counter
|
||||
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
|
||||
.expect("counts")
|
||||
.input_tokens
|
||||
};
|
||||
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
|
||||
assert_eq!(
|
||||
count(
|
||||
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#
|
||||
),
|
||||
base + 1
|
||||
);
|
||||
assert_eq!(
|
||||
count(
|
||||
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#
|
||||
),
|
||||
base
|
||||
);
|
||||
let with_tools = count(
|
||||
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#,
|
||||
);
|
||||
let with_tools_and_system = count(
|
||||
r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#,
|
||||
);
|
||||
assert_eq!(with_tools - with_tools_and_system, 4);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#),
|
||||
base
|
||||
);
|
||||
}
|
||||
|
||||
fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) {
|
||||
assert!(matches!(load("{}"), Err(Error::Load(_))));
|
||||
}
|
||||
|
||||
macro_rules! json_backend_tests {
|
||||
($loader:path) => {
|
||||
#[rstest]
|
||||
#[case::text_only(super::SIMPLE, 14)]
|
||||
#[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)]
|
||||
#[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)]
|
||||
#[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)]
|
||||
#[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)]
|
||||
#[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)]
|
||||
#[case::responses_input_items(super::RESPONSES_INPUT, 62)]
|
||||
#[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)]
|
||||
#[case::rerank_query_and_documents(super::RERANK, 41)]
|
||||
fn count_request_matches_python_token_counter(
|
||||
#[case] body: &str,
|
||||
#[case] expected: usize,
|
||||
) {
|
||||
super::assert_count_request_matches_python_token_counter($loader, body, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::null_messages_win_over_prompt(
|
||||
r#"{"model":"m","messages":null,"prompt":"ignored"}"#,
|
||||
3
|
||||
)]
|
||||
#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)]
|
||||
#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)]
|
||||
#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)]
|
||||
fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) {
|
||||
super::assert_key_presence_follows_python($loader, body, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
|
||||
#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]
|
||||
#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)]
|
||||
#[case::image_block(
|
||||
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"#
|
||||
)]
|
||||
#[case::tool_result_block(
|
||||
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"#
|
||||
)]
|
||||
#[case::array_without_items(
|
||||
br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"#
|
||||
)]
|
||||
fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) {
|
||||
super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_choice_and_system_discount_change_the_count() {
|
||||
super::assert_tool_choice_and_system_discount_change_the_count($loader);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoding_errors_preserve_the_backend_source() {
|
||||
use std::error::Error as _;
|
||||
|
||||
let tokenizer = tokenizers::Tokenizer::new(
|
||||
tokenizers::models::wordpiece::WordPiece::default(),
|
||||
);
|
||||
let expected = tokenizer.encode_fast("hello", true).unwrap_err();
|
||||
let counter = $loader(&tokenizer.to_string(false).unwrap()).unwrap();
|
||||
let request = CountableRequest::parse(br#"{"prompt":"hello"}"#).unwrap();
|
||||
let error = counter.count_request(&request).unwrap_err();
|
||||
assert!(matches!(error, Error::Encode(_)));
|
||||
assert_eq!(error.source().unwrap().to_string(), expected.to_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loading_a_bad_tokenizer_is_a_load_error() {
|
||||
super::assert_loading_a_bad_tokenizer_is_a_load_error($loader);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
mod fast_json {
|
||||
use super::*;
|
||||
|
||||
json_backend_tests!(TokenCounter::from_json_fast);
|
||||
}
|
||||
|
||||
#[cfg(feature = "huggingface")]
|
||||
mod huggingface_json {
|
||||
use super::*;
|
||||
|
||||
json_backend_tests!(TokenCounter::from_json);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -119,203 +254,204 @@ fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) {
|
|||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
|
||||
#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]
|
||||
#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)]
|
||||
#[case::image_block(
|
||||
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"#
|
||||
)]
|
||||
#[case::tool_result_block(
|
||||
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"#
|
||||
)]
|
||||
#[case::array_without_items(
|
||||
br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"#
|
||||
)]
|
||||
fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) {
|
||||
let request = CountableRequest::parse(body).expect("shape parses");
|
||||
#[cfg(any(feature = "fast", feature = "tiktoken"))]
|
||||
mod tiktoken {
|
||||
use super::*;
|
||||
use serde::Deserialize;
|
||||
|
||||
/// 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,
|
||||
source: TokenizerSource,
|
||||
load: fn(&str) -> Result<TokenCounter, Error>,
|
||||
model: &'static str,
|
||||
}
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
const CL100K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "cl100k",
|
||||
source: TokenizerSource::RankFile("9b5ad71b2ce5302211f9c61530b329a4922fc6a4"),
|
||||
load: TokenCounter::from_cl100k_ranks,
|
||||
model: "gpt-4",
|
||||
};
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
const O200K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "o200k",
|
||||
source: TokenizerSource::RankFile("fb374d419588a4632f3f557e76b4b70aebbca790"),
|
||||
load: TokenCounter::from_o200k_ranks,
|
||||
model: "gpt-4o",
|
||||
};
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
const TIKTOKEN_CL100K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "cl100k",
|
||||
source: TokenizerSource::Name("cl100k_base"),
|
||||
load: TokenCounter::from_tiktoken,
|
||||
model: "gpt-4",
|
||||
};
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
const TIKTOKEN_O200K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "o200k",
|
||||
source: TokenizerSource::Name("o200k_base"),
|
||||
load: TokenCounter::from_tiktoken,
|
||||
model: "gpt-4o",
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum TokenizerSource {
|
||||
#[cfg(feature = "fast")]
|
||||
RankFile(&'static str),
|
||||
#[cfg(feature = "tiktoken")]
|
||||
Name(&'static str),
|
||||
}
|
||||
|
||||
fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter {
|
||||
match encoding.source {
|
||||
#[cfg(feature = "tiktoken")]
|
||||
TokenizerSource::Name(name) => (encoding.load)(name).expect("encoding loads"),
|
||||
#[cfg(feature = "fast")]
|
||||
TokenizerSource::RankFile(file) => {
|
||||
let path = format!(
|
||||
"{}/../../../litellm/litellm_core_utils/tokenizers/{file}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
);
|
||||
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 token-counter-fast/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
|
||||
/// `token-counter-fast/tests/fixtures/generate.py`.
|
||||
#[rstest]
|
||||
#[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))]
|
||||
#[cfg_attr(feature = "fast", case::fast_o200k(O200K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_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]
|
||||
#[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))]
|
||||
#[cfg_attr(feature = "fast", case::fast_o200k(O200K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_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]
|
||||
#[cfg_attr(feature = "fast", case::fast_cl100k(CL100K))]
|
||||
#[cfg_attr(feature = "fast", case::fast_o200k(O200K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_cl100k(TIKTOKEN_CL100K))]
|
||||
#[cfg_attr(feature = "tiktoken", case::tiktoken_o200k(TIKTOKEN_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
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "fast")]
|
||||
#[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(_))));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "tiktoken")]
|
||||
#[test]
|
||||
fn unsupported_encoding_reaches_the_counter_caller() {
|
||||
assert!(matches!(
|
||||
counter().count_request(&request),
|
||||
Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems)
|
||||
TokenCounter::from_tiktoken("unknown-encoding"),
|
||||
Err(Error::UnsupportedTokenizer(name)) if name == "unknown-encoding"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_choice_and_system_discount_change_the_count() {
|
||||
let counter = counter();
|
||||
let count = |body: &str| {
|
||||
counter
|
||||
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
|
||||
.expect("counts")
|
||||
.input_tokens
|
||||
};
|
||||
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#),
|
||||
base + 1
|
||||
);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#),
|
||||
base
|
||||
);
|
||||
let with_tools = count(
|
||||
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#,
|
||||
);
|
||||
let with_tools_and_system = count(
|
||||
r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#,
|
||||
);
|
||||
assert_eq!(with_tools - with_tools_and_system, 4);
|
||||
assert_eq!(
|
||||
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#),
|
||||
base
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
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,
|
||||
}
|
||||
|
||||
const CL100K: TiktokenEncoding = TiktokenEncoding {
|
||||
fixtures: "cl100k",
|
||||
rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4",
|
||||
load: TokenCounter::from_cl100k_ranks,
|
||||
model: "gpt-4",
|
||||
};
|
||||
|
||||
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
|
||||
};
|
||||
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(_))));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -100,6 +100,8 @@ class TokenCounter:
|
|||
def from_cl100k_ranks(rank_file: str) -> TokenCounter: ...
|
||||
@staticmethod
|
||||
def from_o200k_ranks(rank_file: str) -> TokenCounter: ...
|
||||
@staticmethod
|
||||
def from_tiktoken(encoding: str) -> TokenCounter: ...
|
||||
def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...
|
||||
|
||||
def gil_stats() -> dict[str, int]: ...
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue