mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(rust): count tiktoken cl100k_base admission tokens in Rust (#40777)
* feat(rust): count tiktoken cl100k_base admission tokens in Rust The Rust admission token counter only had the Anthropic tokenizer, so every other model (OpenAI gpt-4 family, Azure, Gemini, Bedrock non-Claude, Mistral) tokenized with tiktoken on the Python inference worker. Add an exact cl100k_base counter to litellm-token-counter: the vendored rank file (base64 token / rank lines, the bytes Python's tiktoken uses) is parsed into a byte-level BPE model and the cl100k split pattern is a handwritten scanner over the shared Unicode classes, so no regex engine runs per request. Both tokenizers share the message, tool and reply-priming accounting. The PyO3 TokenCounter gains a from_cl100k_ranks constructor; Python reads the rank file and passes it in, the way claude_json_str already works. The bridge selects the counter through the same predicates litellm.token_counter uses (huggingface_tokenizer_kind, openai_tokenizer_encoding), declines o200k_base, downloaded HuggingFace and custom tokenizers to Python, and budget reservation counts once per distinct tokenizer a request names. The legacy gpt-3.5-turbo-0301 message accounting (4 per message, -1 per name) stays in Python: the selector declines it through the predicate token_counter itself uses. * feat(rust): count tiktoken o200k_base admission tokens in Rust (#40794) Add a handwritten o200k_base split scanner and TokenCounter::from_o200k_ranks next to the cl100k_base counter, sharing MergeRanks and the request accounting. The Python bridge selects it when openai_tokenizer_encoding names o200k_base, so gpt-4o, gpt-4.1, gpt-5, o1/o3/o4 and chatgpt-4o requests stop tokenizing on the Python worker under LITELLM_RUST=true Co-authored-by: yassin <yassin@berri.ai> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
d78861bb29
commit
359b7a8489
25 changed files with 10084 additions and 195 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -2002,11 +2002,13 @@ dependencies = [
|
|||
name = "litellm-token-counter"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"indexmap",
|
||||
"itoa",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"rustc-hash",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
|
|
|
|||
|
|
@ -30,12 +30,17 @@ struct TokenCounter {
|
|||
impl TokenCounter {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
.map_err(token_count_error_to_pyerr)?;
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
|
||||
})
|
||||
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
|
||||
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
|
||||
}
|
||||
|
||||
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
|
|
@ -58,6 +63,19 @@ impl TokenCounter {
|
|||
}
|
||||
}
|
||||
|
||||
impl TokenCounter {
|
||||
fn load(
|
||||
py: Python<'_>,
|
||||
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
|
||||
) -> PyResult<Self> {
|
||||
let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?;
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_parallelism() -> usize {
|
||||
available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get)
|
||||
}
|
||||
|
|
@ -70,7 +88,7 @@ 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(_) => PyValueError::new_err(message),
|
||||
Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
|
||||
Error::RequestParse(_)
|
||||
| Error::MissingInput
|
||||
| Error::FloatText
|
||||
|
|
|
|||
|
|
@ -6,8 +6,10 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
indexmap = { version = "2.14.0", features = ["serde"] }
|
||||
itoa = "1.0"
|
||||
rustc-hash = "2.1.3"
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ use tokenizers::pre_tokenizers::PreTokenizerWrapper;
|
|||
use tokenizers::{Model, Tokenizer};
|
||||
use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick};
|
||||
|
||||
use super::unicode_classes::UnicodeClasses;
|
||||
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
|
||||
|
||||
const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"];
|
||||
|
||||
|
|
@ -110,27 +110,6 @@ fn mapped_len(piece: &str) -> usize {
|
|||
.count()
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
enum Class {
|
||||
Letter,
|
||||
Number,
|
||||
Space,
|
||||
Other,
|
||||
}
|
||||
|
||||
fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
|
||||
match character {
|
||||
'A'..='Z' | 'a'..='z' => Class::Letter,
|
||||
'0'..='9' => Class::Number,
|
||||
'\t'..='\r' | ' ' => Class::Space,
|
||||
_ if character.is_ascii() => Class::Other,
|
||||
_ if unicode_classes.is_letter(character) => Class::Letter,
|
||||
_ if unicode_classes.is_number(character) => Class::Number,
|
||||
_ if unicode_classes.is_space(character) => Class::Space,
|
||||
_ => Class::Other,
|
||||
}
|
||||
}
|
||||
|
||||
/// The regex matches every character, so the pieces tile the text.
|
||||
fn pieces<'a>(
|
||||
text: &'a str,
|
||||
|
|
@ -169,12 +148,6 @@ fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize
|
|||
}
|
||||
}
|
||||
|
||||
fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
|
||||
text.char_indices()
|
||||
.find(|(_, character)| class(*character, unicode_classes) != run_class)
|
||||
.map_or(text.len(), |(index, _)| index)
|
||||
}
|
||||
|
||||
/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last
|
||||
/// character to start the next piece (` ?` on the following alternatives).
|
||||
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
|
|
|
|||
125
litellm-rust/crates/token-counter/src/cl100k.rs
Normal file
125
litellm-rust/crates/token-counter/src/cl100k.rs
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
//! Scanner for tiktoken's `cl100k_base` split regex,
|
||||
//! `'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s`.
|
||||
|
||||
use super::scanner::{contraction_len, digit_run_len, is_newline};
|
||||
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
|
||||
|
||||
/// The alternatives in regex order; the possessive quantifiers mean an
|
||||
/// alternative that starts matching and runs out of input fails as a whole.
|
||||
pub(super) fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
|
||||
if let Some(len) = contraction_len(text) {
|
||||
return len;
|
||||
}
|
||||
let first_class = class(first, unicode_classes);
|
||||
match first_class {
|
||||
Class::Letter => return run_len(text, Class::Letter, unicode_classes),
|
||||
Class::Number => return digit_run_len(text, unicode_classes),
|
||||
Class::Space | Class::Other => {}
|
||||
}
|
||||
let rest = &text[first.len_utf8()..];
|
||||
let second_class = rest
|
||||
.chars()
|
||||
.next()
|
||||
.map(|character| class(character, unicode_classes));
|
||||
if !is_newline(first) && second_class == Some(Class::Letter) {
|
||||
return first.len_utf8() + run_len(rest, Class::Letter, unicode_classes);
|
||||
}
|
||||
if first_class == Class::Other {
|
||||
return symbol_run_len(text, unicode_classes);
|
||||
}
|
||||
if first == ' ' && second_class == Some(Class::Other) {
|
||||
return 1 + symbol_run_len(rest, unicode_classes);
|
||||
}
|
||||
space_run_len(text, unicode_classes)
|
||||
}
|
||||
|
||||
/// `[^\s\p{L}\p{N}]++[\r\n]*+`
|
||||
fn symbol_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
let symbols = run_len(text, Class::Other, unicode_classes);
|
||||
symbols
|
||||
+ text[symbols..]
|
||||
.bytes()
|
||||
.take_while(|byte| matches!(byte, b'\r' | b'\n'))
|
||||
.count()
|
||||
}
|
||||
|
||||
/// `\s++$|\s*[\r\n]|\s+(?!\S)|\s`: whitespace to the end of the text is one
|
||||
/// piece; otherwise the piece ends at the last newline of the run, or leaves
|
||||
/// the run's last character for the next piece's optional leading space.
|
||||
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
let run = run_len(text, Class::Space, unicode_classes);
|
||||
if run == text.len() {
|
||||
return run;
|
||||
}
|
||||
if let Some(newline) = text[..run].rfind(['\r', '\n']) {
|
||||
return newline + 1;
|
||||
}
|
||||
let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
|
||||
match run - last {
|
||||
0 => run,
|
||||
shorter => shorter,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::scanner::pieces;
|
||||
|
||||
fn split(text: &str) -> Vec<&str> {
|
||||
pieces(
|
||||
text,
|
||||
piece_len,
|
||||
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
|
||||
)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("", &[])]
|
||||
#[case("Hello world", &["Hello", " world"])]
|
||||
#[case("don't I'LL you'Ve we'RE he'd I'm", &["don", "'t", " I", "'LL", " you", "'Ve", " we", "'RE", " he", "'d", " I", "'m"])]
|
||||
#[case("IT'SOK it'Dbe 'Sx 'Tx", &["IT", "'S", "OK", " it", "'D", "be", " '", "Sx", " '", "Tx"])]
|
||||
#[case("'Sx'Tx'Mx'LLx'VEx'REx'Dx", &["'S", "x", "'T", "x", "'M", "x", "'LL", "x", "'VE", "x", "'RE", "x", "'D", "x"])]
|
||||
#[case("'ſ 'lx", &["'ſ", " '", "lx"])]
|
||||
#[case("12345 6", &["123", "45", " ", "6"])]
|
||||
#[case("!abc !!abc", &["!abc", " !!", "abc"])]
|
||||
#[case(" !!!\r\n\r\nx", &[" !!!\r\n\r\n", "x"])]
|
||||
#[case("a b \n\n c", &["a", " ", " b", " \n\n", " ", " c"])]
|
||||
#[case("a\nb\r\nc\n\nd \n e", &["a", "\n", "b", "\r\n", "c", "\n\n", "d", " \n", " e"])]
|
||||
#[case("x \t\n \t y\n", &["x", " \t\n", " \t", " y", "\n"])]
|
||||
#[case("end ", &["end", " "])]
|
||||
#[case("\u{a0}abc\u{a0}!", &["\u{a0}abc", "\u{a0}", "!"])]
|
||||
#[case("<|endoftext|>", &["<|", "endoftext", "|>"])]
|
||||
#[case("e\u{301}a", &["e", "\u{301}a"])]
|
||||
#[case("日本語 ١٢٣٤", &["日本語", " ", "١٢٣", "٤"])]
|
||||
fn scanner_splits_like_the_regex(#[case] text: &str, #[case] expected: &[&str]) {
|
||||
assert_eq!(split(text), expected);
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct TextFixture {
|
||||
text: String,
|
||||
pieces: Vec<String>,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scanner_splits_the_fixture_corpus_like_tiktoken_regex() {
|
||||
let fixtures: Vec<TextFixture> = include_str!(concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/tests/fixtures/cl100k/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(|fixture| split(&fixture.text) != fixture.pieces)
|
||||
.map(|fixture| (&fixture.text, split(&fixture.text), &fixture.pieces))
|
||||
.collect();
|
||||
assert!(mismatches.is_empty(), "{mismatches:#?}");
|
||||
}
|
||||
}
|
||||
|
|
@ -3,6 +3,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,
|
||||
|
|
@ -23,12 +24,19 @@ pub struct InputTokenCount {
|
|||
pub input_tokens: usize,
|
||||
}
|
||||
|
||||
/// A loaded HuggingFace tokenizer plus the message accounting Python applies on
|
||||
/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their
|
||||
/// event loop.
|
||||
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 {
|
||||
tokenizer: tokenizers::Tokenizer,
|
||||
byte_level: Option<ByteLevelCounter>,
|
||||
encoder: Encoder,
|
||||
}
|
||||
|
||||
impl TokenCounter {
|
||||
|
|
@ -39,23 +47,50 @@ impl TokenCounter {
|
|||
.map_err(Error::Load)?;
|
||||
let byte_level = ByteLevelCounter::detect(&tokenizer);
|
||||
Ok(Self {
|
||||
tokenizer,
|
||||
byte_level,
|
||||
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 count_text(&self, text: &str) -> Result<usize, Error> {
|
||||
if let Some(count) = self
|
||||
.byte_level
|
||||
.as_ref()
|
||||
.and_then(|counter| counter.count(&self.tokenizer, text))
|
||||
{
|
||||
return Ok(count);
|
||||
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.tokenizer
|
||||
.encode_fast(text, true)
|
||||
.map(|encoding| encoding.len())
|
||||
.map_err(Error::Encode)
|
||||
}
|
||||
|
||||
/// Mirrors the host's key precedence: `messages`, then `prompt`, then
|
||||
|
|
|
|||
|
|
@ -6,6 +6,10 @@ use thiserror::Error as 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("unsupported by the rust token counter: request body could not be parsed: {0}")]
|
||||
RequestParse(#[source] serde_json::Error),
|
||||
#[error("unsupported by the rust token counter: request has no countable input")]
|
||||
|
|
|
|||
|
|
@ -5,9 +5,13 @@
|
|||
#![forbid(unsafe_code)]
|
||||
|
||||
mod byte_level;
|
||||
mod cl100k;
|
||||
mod counter;
|
||||
mod error;
|
||||
mod o200k;
|
||||
mod python_json;
|
||||
mod scanner;
|
||||
mod tiktoken;
|
||||
mod tools;
|
||||
mod types;
|
||||
mod unicode_classes;
|
||||
|
|
|
|||
217
litellm-rust/crates/token-counter/src/o200k.rs
Normal file
217
litellm-rust/crates/token-counter/src/o200k.rs
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
//! Scanner for tiktoken's `o200k_base` split regex,
|
||||
//! `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`.
|
||||
//! tiktoken runs it with a backtracking engine, so the letter alternatives
|
||||
//! below reproduce where the greedy quantifiers settle, not only what the
|
||||
//! classes say.
|
||||
|
||||
use super::scanner::{contraction_len, digit_run_len, is_newline};
|
||||
use super::unicode_classes::{Case, Class, UnicodeClasses, case, case_run_len, class, run_len};
|
||||
|
||||
/// The alternatives in regex order: a number is never a letter piece, a
|
||||
/// letter always is, and only whitespace and symbols reach the last three.
|
||||
pub(super) fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
|
||||
let first_class = class(first, unicode_classes);
|
||||
if first_class == Class::Number {
|
||||
return digit_run_len(text, unicode_classes);
|
||||
}
|
||||
if let Some(len) = letter_piece_len(text, first, first_class, unicode_classes) {
|
||||
return len;
|
||||
}
|
||||
if first_class == Class::Other {
|
||||
return symbol_run_len(text, unicode_classes);
|
||||
}
|
||||
let rest = &text[first.len_utf8()..];
|
||||
if first == ' '
|
||||
&& rest
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|character| class(character, unicode_classes) == Class::Other)
|
||||
{
|
||||
return 1 + symbol_run_len(rest, unicode_classes);
|
||||
}
|
||||
space_run_len(text, unicode_classes)
|
||||
}
|
||||
|
||||
/// The two letter alternatives, each first with then without the optional
|
||||
/// `[^\r\n\p{L}\p{N}]` prefix: the order the engine tries them in.
|
||||
fn letter_piece_len(
|
||||
text: &str,
|
||||
first: char,
|
||||
first_class: Class,
|
||||
unicode_classes: &UnicodeClasses,
|
||||
) -> Option<usize> {
|
||||
let prefix = (!is_newline(first) && matches!(first_class, Class::Space | Class::Other))
|
||||
.then(|| first.len_utf8());
|
||||
let after_prefix = |shape: fn(&str, &UnicodeClasses) -> Option<usize>| {
|
||||
prefix.and_then(|prefix| shape(&text[prefix..], unicode_classes).map(|len| prefix + len))
|
||||
};
|
||||
after_prefix(upper_then_lower_len)
|
||||
.or_else(|| upper_then_lower_len(text, unicode_classes))
|
||||
.or_else(|| after_prefix(upper_run_len))
|
||||
.or_else(|| upper_run_len(text, unicode_classes))
|
||||
}
|
||||
|
||||
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?`.
|
||||
/// The upper run is greedy; when no lower character follows it, the engine
|
||||
/// gives characters back until the one it just gave back is lower too, and
|
||||
/// that single character is the lower run.
|
||||
fn upper_then_lower_len(text: &str, unicode_classes: &UnicodeClasses) -> Option<usize> {
|
||||
let upper = case_run_len(text, Case::is_upper, unicode_classes);
|
||||
let lower = case_run_len(&text[upper..], Case::is_lower, unicode_classes);
|
||||
let letters = if lower > 0 {
|
||||
upper + lower
|
||||
} else {
|
||||
let (index, last_both) = text[..upper]
|
||||
.char_indices()
|
||||
.rev()
|
||||
.find(|(_, character)| case(*character, unicode_classes).is_lower())?;
|
||||
index + last_both.len_utf8()
|
||||
};
|
||||
Some(letters + contraction_len(&text[letters..]).unwrap_or(0))
|
||||
}
|
||||
|
||||
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?`
|
||||
fn upper_run_len(text: &str, unicode_classes: &UnicodeClasses) -> Option<usize> {
|
||||
let upper = case_run_len(text, Case::is_upper, unicode_classes);
|
||||
if upper == 0 {
|
||||
return None;
|
||||
}
|
||||
let letters = upper + case_run_len(&text[upper..], Case::is_lower, unicode_classes);
|
||||
Some(letters + contraction_len(&text[letters..]).unwrap_or(0))
|
||||
}
|
||||
|
||||
/// `[^\s\p{L}\p{N}]+[\r\n/]*`
|
||||
fn symbol_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
let symbols = run_len(text, Class::Other, unicode_classes);
|
||||
symbols
|
||||
+ text[symbols..]
|
||||
.bytes()
|
||||
.take_while(|byte| matches!(byte, b'\r' | b'\n' | b'/'))
|
||||
.count()
|
||||
}
|
||||
|
||||
/// `\s*[\r\n]+|\s+(?!\S)|\s+`: a run with a newline ends at its last newline,
|
||||
/// even at the end of the text; otherwise whitespace to the end of the text
|
||||
/// is one piece, or the run leaves its last character for the next piece's
|
||||
/// optional leading space.
|
||||
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
let run = run_len(text, Class::Space, unicode_classes);
|
||||
if let Some(newline) = text[..run].rfind(['\r', '\n']) {
|
||||
return newline + 1;
|
||||
}
|
||||
if run == text.len() {
|
||||
return run;
|
||||
}
|
||||
let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
|
||||
match run - last {
|
||||
0 => run,
|
||||
shorter => shorter,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use tokenizers::utils::SysRegex;
|
||||
|
||||
use super::*;
|
||||
use crate::scanner::pieces;
|
||||
|
||||
fn split(text: &str) -> Vec<&str> {
|
||||
pieces(
|
||||
text,
|
||||
piece_len,
|
||||
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
|
||||
)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("", &[])]
|
||||
#[case("Hello world", &["Hello", " world"])]
|
||||
#[case("camelCase PascalCase ABCdef ABCdeF ABC", &["camel", "Case", " Pascal", "Case", " ABCdef", " ABCde", "F", " ABC"])]
|
||||
#[case("日本ABC ABC日本 日本語abc abc日本語", &["日本", "ABC", " ABC日本", " 日本語abc", " abc日本語"])]
|
||||
#[case("\u{301}ABC \u{301}abc \u{301}\u{301}A A\u{301}\u{301} E\u{301}A aE\u{301}", &["\u{301}", "ABC", " \u{301}abc", " \u{301}\u{301}", "A", " A\u{301}\u{301}", " E\u{301}", "A", " a", "E\u{301}"])]
|
||||
#[case("ᵃbc ᵃBC Aᵃbc Aᵃ ᵃ' ᵃ's", &["ᵃbc", " ᵃ", "BC", " Aᵃbc", " Aᵃ", " ᵃ", "'", " ᵃ's"])]
|
||||
#[case("Džungla aDžB ADžB ADžb", &["Džungla", " a", "DžB", " ADžB", " ADžb"])]
|
||||
#[case("don'tx ABC's abc'S abc'ſ ABC'ſx IT'SOK it'Dbe", &["don't", "x", " ABC's", " abc'S", " abc'ſ", " ABC'ſ", "x", " IT'S", "OK", " it'D", "be"])]
|
||||
#[case("'sabc x's 's 'Sx'Tx 9'9 a'9 ' s", &["'sabc", " x's", " '", "s", " '", "Sx'T", "x", " ", "9", "'", "9", " a", "'", "9", " '", " s"])]
|
||||
#[case("!ABC !AbC !!abc !!\u{301}a \u{a0}\u{301}A", &["!ABC", " !", "Ab", "C", " !!", "abc", " !!\u{301}", "a", " ", "\u{a0}\u{301}", "A"])]
|
||||
#[case("!!/\n/x a/b !!\n/x /x //", &["!!/\n/", "x", " a", "/b", " !!\n/", "x", " ", " /", "x", " ", " //"])]
|
||||
#[case("12345 6 1abc abc1", &["123", "45", " ", "6", " ", "1", "abc", " abc", "1"])]
|
||||
#[case("x \n x \r\n \r\n y", &["x", " \n", " x", " \r\n \r\n", " y"])]
|
||||
#[case("x \n ", &["x", " \n", " "])]
|
||||
#[case("a b \n\n c", &["a", " ", " b", " \n\n", " ", " c"])]
|
||||
#[case("x\t\ty x\t\t", &["x", "\t", "\ty", " x", "\t\t"])]
|
||||
#[case("end ", &["end", " "])]
|
||||
#[case("\u{a0}abc\u{a0}!", &["\u{a0}abc", "\u{a0}", "!"])]
|
||||
#[case("<|endoftext|>", &["<|", "endoftext", "|>"])]
|
||||
#[case("İstanbul ΣΊΣΥΦΟΣ Ελληνικά Русский", &["İstanbul", " ΣΊΣΥΦΟΣ", " Ελληνικά", " Русский"])]
|
||||
#[case("日本語 ١٢٣٤", &["日本語", " ", "١٢٣", "٤"])]
|
||||
fn scanner_splits_like_the_regex(#[case] text: &str, #[case] expected: &[&str]) {
|
||||
assert_eq!(split(text), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_scalar_alone_is_one_piece() {
|
||||
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
|
||||
let text = character.to_string();
|
||||
assert_eq!(
|
||||
split(&text),
|
||||
[text.as_str()],
|
||||
"U+{:04X}",
|
||||
u32::from(character)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cases_match_oniguruma() {
|
||||
let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes");
|
||||
let upper = SysRegex::new(r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]").expect("regex");
|
||||
let lower = SysRegex::new(r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]").expect("regex");
|
||||
let whole =
|
||||
|regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len()));
|
||||
let mut text = String::new();
|
||||
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
|
||||
text.clear();
|
||||
text.push(character);
|
||||
let expected = match (whole(&upper, &text), whole(&lower, &text)) {
|
||||
(true, true) => Case::Both,
|
||||
(true, false) => Case::Upper,
|
||||
(false, true) => Case::Lower,
|
||||
(false, false) => Case::Neither,
|
||||
};
|
||||
assert_eq!(
|
||||
case(character, unicode_classes),
|
||||
expected,
|
||||
"U+{:04X}",
|
||||
u32::from(character)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct TextFixture {
|
||||
text: String,
|
||||
pieces: Vec<String>,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scanner_splits_the_fixture_corpus_like_tiktoken_regex() {
|
||||
let fixtures: Vec<TextFixture> = include_str!(concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/tests/fixtures/o200k/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(|fixture| split(&fixture.text) != fixture.pieces)
|
||||
.map(|fixture| (&fixture.text, split(&fixture.text), &fixture.pieces))
|
||||
.collect();
|
||||
assert!(mismatches.is_empty(), "{mismatches:#?}");
|
||||
}
|
||||
}
|
||||
107
litellm-rust/crates/token-counter/src/scanner.rs
Normal file
107
litellm-rust/crates/token-counter/src/scanner.rs
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
//! Exact token counting for tiktoken encodings. A hand-written scanner
|
||||
//! reproduces the piece boundaries of the encoding's split regex, and each
|
||||
//! piece is merged with the rank file. Special tokens are ordinary text, as
|
||||
//! with `encode(text, disallowed_special=())`.
|
||||
|
||||
use std::iter;
|
||||
|
||||
use super::tiktoken::{MergeRanks, MergeScratch};
|
||||
use super::unicode_classes::{Class, UnicodeClasses, class};
|
||||
use super::{cl100k, o200k};
|
||||
use crate::Error;
|
||||
|
||||
const MAX_DIGITS_PER_PIECE: usize = 3;
|
||||
|
||||
/// Byte length of the piece the split regex matches at the start of the
|
||||
/// text, given the text's first character.
|
||||
pub(super) type PieceLen = fn(&str, char, &UnicodeClasses) -> usize;
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(super) enum SplitPattern {
|
||||
Cl100k,
|
||||
O200k,
|
||||
}
|
||||
|
||||
impl SplitPattern {
|
||||
fn piece_len(self) -> PieceLen {
|
||||
match self {
|
||||
Self::Cl100k => cl100k::piece_len,
|
||||
Self::O200k => o200k::piece_len,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct TiktokenCounter {
|
||||
ranks: MergeRanks,
|
||||
piece_len: PieceLen,
|
||||
unicode_classes: &'static UnicodeClasses,
|
||||
}
|
||||
|
||||
impl TiktokenCounter {
|
||||
pub(super) fn from_ranks(split: SplitPattern, rank_file: &str) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
ranks: MergeRanks::parse(rank_file)?,
|
||||
piece_len: split.piece_len(),
|
||||
unicode_classes: UnicodeClasses::get().ok_or(Error::UnicodeClasses)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn count(&self, text: &str) -> usize {
|
||||
let mut scratch = MergeScratch::default();
|
||||
pieces(text, self.piece_len, self.unicode_classes)
|
||||
.map(|piece| self.ranks.count_piece(piece.as_bytes(), &mut scratch))
|
||||
.sum()
|
||||
}
|
||||
}
|
||||
|
||||
/// The regex matches every character, so the pieces tile the text.
|
||||
pub(super) fn pieces<'a>(
|
||||
text: &'a str,
|
||||
piece_len: PieceLen,
|
||||
unicode_classes: &'static UnicodeClasses,
|
||||
) -> impl Iterator<Item = &'a str> {
|
||||
iter::successors(
|
||||
split_piece(text, piece_len, unicode_classes),
|
||||
move |(_, rest)| split_piece(rest, piece_len, unicode_classes),
|
||||
)
|
||||
.map(|(piece, _)| piece)
|
||||
}
|
||||
|
||||
fn split_piece<'a>(
|
||||
text: &'a str,
|
||||
piece_len: PieceLen,
|
||||
unicode_classes: &UnicodeClasses,
|
||||
) -> Option<(&'a str, &'a str)> {
|
||||
let first = text.chars().next()?;
|
||||
Some(text.split_at(piece_len(text, first, unicode_classes)))
|
||||
}
|
||||
|
||||
/// `'(?i:s|t|re|ve|m|ll|d)`, the contraction both encodings spell out. Simple
|
||||
/// case folding also maps U+017F (long s) onto `s`.
|
||||
pub(super) fn contraction_len(text: &str) -> Option<usize> {
|
||||
let mut characters = text.chars();
|
||||
if characters.next()? != '\'' {
|
||||
return None;
|
||||
}
|
||||
let first = characters.next()?;
|
||||
let len = match first {
|
||||
's' | 'S' | '\u{17F}' | 'd' | 'D' | 'm' | 'M' | 't' | 'T' => first.len_utf8(),
|
||||
'l' | 'L' => matches!(characters.next(), Some('l' | 'L')).then_some(2)?,
|
||||
'v' | 'V' | 'r' | 'R' => matches!(characters.next(), Some('e' | 'E')).then_some(2)?,
|
||||
_ => return None,
|
||||
};
|
||||
Some(1 + len)
|
||||
}
|
||||
|
||||
pub(super) fn is_newline(character: char) -> bool {
|
||||
matches!(character, '\r' | '\n')
|
||||
}
|
||||
|
||||
/// `\p{N}{1,3}`
|
||||
pub(super) fn digit_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
|
||||
text.chars()
|
||||
.take(MAX_DIGITS_PER_PIECE)
|
||||
.take_while(|character| class(*character, unicode_classes) == Class::Number)
|
||||
.map(char::len_utf8)
|
||||
.sum()
|
||||
}
|
||||
215
litellm-rust/crates/token-counter/src/tiktoken.rs
Normal file
215
litellm-rust/crates/token-counter/src/tiktoken.rs
Normal file
|
|
@ -0,0 +1,215 @@
|
|||
//! 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_stay_cheap() {
|
||||
let ranks = ranks();
|
||||
let mut scratch = MergeScratch::default();
|
||||
let piece = vec![b' '; 1 << 20];
|
||||
let started = std::time::Instant::now();
|
||||
let count = ranks.count_piece(&piece, &mut scratch);
|
||||
assert!(count > 0);
|
||||
assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed());
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
|
|
@ -9,6 +9,8 @@ pub(super) struct UnicodeClasses {
|
|||
letters: Ranges,
|
||||
numbers: Ranges,
|
||||
spaces: Ranges,
|
||||
uppers: Ranges,
|
||||
lowers: Ranges,
|
||||
}
|
||||
|
||||
static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
|
||||
|
|
@ -19,6 +21,8 @@ static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
|
|||
letters: Ranges::load(r"\p{L}+", &scalars)?,
|
||||
numbers: Ranges::load(r"\p{N}+", &scalars)?,
|
||||
spaces: Ranges::load(r"\s+", &scalars)?,
|
||||
uppers: Ranges::load(r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+", &scalars)?,
|
||||
lowers: Ranges::load(r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+", &scalars)?,
|
||||
})
|
||||
});
|
||||
|
||||
|
|
@ -59,15 +63,102 @@ impl UnicodeClasses {
|
|||
CLASSES.as_ref()
|
||||
}
|
||||
|
||||
pub(super) fn is_letter(&self, character: char) -> bool {
|
||||
fn is_letter(&self, character: char) -> bool {
|
||||
self.letters.contains(character)
|
||||
}
|
||||
|
||||
pub(super) fn is_number(&self, character: char) -> bool {
|
||||
fn is_number(&self, character: char) -> bool {
|
||||
self.numbers.contains(character)
|
||||
}
|
||||
|
||||
pub(super) fn is_space(&self, character: char) -> bool {
|
||||
fn is_space(&self, character: char) -> bool {
|
||||
self.spaces.contains(character)
|
||||
}
|
||||
|
||||
fn is_upper(&self, character: char) -> bool {
|
||||
self.uppers.contains(character)
|
||||
}
|
||||
|
||||
fn is_lower(&self, character: char) -> bool {
|
||||
self.lowers.contains(character)
|
||||
}
|
||||
}
|
||||
|
||||
/// `\p{L}`, `\p{N}`, `\s` and everything else, the character classes the
|
||||
/// split regexes are written in.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(super) enum Class {
|
||||
Letter,
|
||||
Number,
|
||||
Space,
|
||||
Other,
|
||||
}
|
||||
|
||||
pub(super) fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
|
||||
match character {
|
||||
'A'..='Z' | 'a'..='z' => Class::Letter,
|
||||
'0'..='9' => Class::Number,
|
||||
'\t'..='\r' | ' ' => Class::Space,
|
||||
_ if character.is_ascii() => Class::Other,
|
||||
_ if unicode_classes.is_letter(character) => Class::Letter,
|
||||
_ if unicode_classes.is_number(character) => Class::Number,
|
||||
_ if unicode_classes.is_space(character) => Class::Space,
|
||||
_ => Class::Other,
|
||||
}
|
||||
}
|
||||
|
||||
/// Byte length of the leading run of `run_class` characters.
|
||||
pub(super) fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
|
||||
text.char_indices()
|
||||
.find(|(_, character)| class(*character, unicode_classes) != run_class)
|
||||
.map_or(text.len(), |(index, _)| index)
|
||||
}
|
||||
|
||||
/// Membership in the two letter classes of the o200k split regex,
|
||||
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]` and `[\p{Ll}\p{Lm}\p{Lo}\p{M}]`; `Lm`,
|
||||
/// `Lo` and `M` are in both.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(super) enum Case {
|
||||
Upper,
|
||||
Lower,
|
||||
Both,
|
||||
Neither,
|
||||
}
|
||||
|
||||
impl Case {
|
||||
pub(super) fn is_upper(self) -> bool {
|
||||
matches!(self, Case::Upper | Case::Both)
|
||||
}
|
||||
|
||||
pub(super) fn is_lower(self) -> bool {
|
||||
matches!(self, Case::Lower | Case::Both)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn case(character: char, unicode_classes: &UnicodeClasses) -> Case {
|
||||
match character {
|
||||
'A'..='Z' => Case::Upper,
|
||||
'a'..='z' => Case::Lower,
|
||||
_ if character.is_ascii() => Case::Neither,
|
||||
_ => match (
|
||||
unicode_classes.is_upper(character),
|
||||
unicode_classes.is_lower(character),
|
||||
) {
|
||||
(true, true) => Case::Both,
|
||||
(true, false) => Case::Upper,
|
||||
(false, true) => Case::Lower,
|
||||
(false, false) => Case::Neither,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Byte length of the leading run of characters whose case passes `in_class`.
|
||||
pub(super) fn case_run_len(
|
||||
text: &str,
|
||||
in_class: fn(Case) -> bool,
|
||||
unicode_classes: &UnicodeClasses,
|
||||
) -> usize {
|
||||
text.char_indices()
|
||||
.find(|(_, character)| !in_class(case(*character, unicode_classes)))
|
||||
.map_or(text.len(), |(index, _)| index)
|
||||
}
|
||||
|
|
|
|||
10
litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl
vendored
Normal file
10
litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
4053
litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl
vendored
Normal file
4053
litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
422
litellm-rust/crates/token-counter/tests/fixtures/generate.py
vendored
Normal file
422
litellm-rust/crates/token-counter/tests/fixtures/generate.py
vendored
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
"""Pin tiktoken reference counts for the Rust parity tests of one encoding.
|
||||
|
||||
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
|
||||
|
||||
`<encoding>/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens`
|
||||
counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`,
|
||||
the same call `litellm.token_counter` makes, and `pieces` the installed
|
||||
encoding's split pattern applied with the `regex` module tiktoken itself uses,
|
||||
so a scanner that splits differently fails even where BPE would count the same.
|
||||
`<encoding>/requests.jsonl` holds `{"body", "input_tokens"}` lines, `body` being
|
||||
the exact request bytes as a JSON string, counted with the proxy's admission
|
||||
counter (`_count_input_tokens(body, model)`) for a model Python counts with that
|
||||
encoding. Every message in the 50k-token body is shorter than the Python chunk
|
||||
size so the chunked Python count equals the exact whole-text tiktoken count the
|
||||
Rust counter produces.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import regex
|
||||
import tiktoken
|
||||
|
||||
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
|
||||
from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
|
||||
|
||||
HERE: Final = Path(__file__).resolve().parent
|
||||
MODELS: Final = {"cl100k_base": "gpt-4", "o200k_base": "gpt-4o"}
|
||||
ENCODING_NAME: Final = sys.argv[1]
|
||||
MODEL: Final = MODELS[ENCODING_NAME]
|
||||
ENCODING: Final = tiktoken.get_encoding(ENCODING_NAME)
|
||||
assert openai_tokenizer_encoding(MODEL).name == ENCODING_NAME
|
||||
SPLIT_PATTERN: Final = regex.compile(ENCODING._pat_str) # pyright: ignore[reportPrivateUsage] # tiktoken has no public accessor
|
||||
OUT: Final = HERE / ENCODING_NAME.removesuffix("_base")
|
||||
|
||||
# Mirrors ALPHABET in src/byte_level.rs, plus the pieces the tiktoken patterns treat differently.
|
||||
ALPHABET: Final = (
|
||||
"a",
|
||||
"Z",
|
||||
"e",
|
||||
"s",
|
||||
"t",
|
||||
"d",
|
||||
"m",
|
||||
"'",
|
||||
"'s",
|
||||
"'re",
|
||||
"'ll",
|
||||
"'S",
|
||||
"0",
|
||||
"9",
|
||||
" ",
|
||||
" ",
|
||||
"\t",
|
||||
"\n",
|
||||
"\r\n",
|
||||
"\x0b",
|
||||
".",
|
||||
",",
|
||||
"!",
|
||||
"-",
|
||||
"(",
|
||||
'"',
|
||||
"\xa0",
|
||||
"\x85",
|
||||
"\u2028",
|
||||
"\u3000",
|
||||
"\u200b",
|
||||
"\u200d",
|
||||
"é",
|
||||
"e\u0301",
|
||||
"ß",
|
||||
"漢",
|
||||
"字",
|
||||
"ع",
|
||||
"३",
|
||||
"½",
|
||||
"Ⅳ",
|
||||
"🙂",
|
||||
"👍🏽",
|
||||
"Ａ",
|
||||
"fi",
|
||||
"㍿",
|
||||
"㋿",
|
||||
"ꟲ",
|
||||
"𐞁",
|
||||
"a\u030a",
|
||||
"\u1e0b\u0323",
|
||||
"<",
|
||||
">",
|
||||
"EOT",
|
||||
"<EOT>",
|
||||
"<META_START>",
|
||||
"'D",
|
||||
"'M",
|
||||
"'T",
|
||||
"'VE",
|
||||
"'Re",
|
||||
"'ſ",
|
||||
"ſ",
|
||||
"12345678",
|
||||
"٣٤٥٦",
|
||||
"<|endoftext|>",
|
||||
"<|fim_prefix|>",
|
||||
"\r",
|
||||
"\r\n\r\n",
|
||||
" \n",
|
||||
"!!",
|
||||
"#$%",
|
||||
"\u00ad",
|
||||
"\u0301",
|
||||
"\U0001f600\U0001f3fd",
|
||||
"İ",
|
||||
"Dž",
|
||||
)
|
||||
|
||||
# The pieces the o200k case-shaped letter branch and slash-absorbing symbol branch split differently.
|
||||
CASE_ALPHABET: Final = ALPHABET + (
|
||||
"B",
|
||||
"Ab",
|
||||
"aB",
|
||||
"ABC",
|
||||
"ᵃ",
|
||||
"camelCase",
|
||||
"HTTPServer",
|
||||
"iOS",
|
||||
"Džungla",
|
||||
"/",
|
||||
"\n/",
|
||||
"/\r\n",
|
||||
" \n ",
|
||||
"a/b",
|
||||
)
|
||||
|
||||
CORPUS: Final = (
|
||||
"",
|
||||
"Hello, how are you today?",
|
||||
"I'm sure they're right, we'll see. WE'LL SEE, I'M SURE THEY'RE RIGHT, IT'S HERS AND IT'D BE 'D",
|
||||
"don't Don'T DON'T won'T i've I'VE i'Ve you'RE 'S 'T 'M 'D 'LL 'VE 'RE 'ſ 'x",
|
||||
"1234567890 123 12 1 0000000 ٣٤٥٦٧٨ ३४५६ 1,234,567.89 2026-09-11T18:00:00Z",
|
||||
"$abc %def &ghi @jkl _mno #pqr ~stu ^vwx |yz \\a /b :c ;d ?e !f (g )h [i ]j {k }l <m >n =o +p *q",
|
||||
"foo bar baz \t qux\t\tquux \n\nline\r\nline\r\n\r\n \n\t\r\n x ",
|
||||
"trailing spaces ",
|
||||
"trailing tabs\t\t",
|
||||
"trailing newline\n",
|
||||
"\n\n\n",
|
||||
"\r\n\r\n\r\n",
|
||||
" ",
|
||||
"😀😃😄 👍🏽 🇺🇸 👨👩👧👦 ✈️ ❤️🔥 ٭ ※ ⌘ ⏎",
|
||||
"漢字かな交じり文、東京都千代田区。日本語のテキストです。中文测试。한국어 텍스트",
|
||||
"مرحبا بالعالم، هذا نص عربي مع أرقام ١٢٣٤٥٦٧ و علامات ترقيم!",
|
||||
"Zürich, façade, naïve, Ærøskøbing, Ελληνικά, Русский текст, עברית, हिन्दी, ไทย",
|
||||
"e\u0301 a\u030a \u1e0b\u0323 \u0301\u0301 combining\u0308 marks\u0301!",
|
||||
"ΣΊΣΥΦΟΣ Džungla İstanbul file flow Ａｂｃ ㍿ ㋿ ꟲ 𐞁",
|
||||
"<|endoftext|> <|fim_prefix|>code<|fim_middle|>more<|fim_suffix|> <|endofprompt|> <|im_start|>",
|
||||
"<EOT> <META_START> <s> </s> [INST] [/INST] <<SYS>>",
|
||||
"def f(x):\n return {'a': x ** 2, \"b\": [1, 2, 3]} # comment\n\nprint(f(10))\n",
|
||||
'{"model":"gpt-4","messages":[{"role":"user","content":"hi\\n"}],"temperature":0.7}',
|
||||
"https://example.com/path?query=1&other=two#fragment user@example.com 192.168.0.1",
|
||||
"a" * 3000,
|
||||
" " * 3000,
|
||||
"." * 3000,
|
||||
"ab" * 1500,
|
||||
"\n" * 3000,
|
||||
"0" * 3000,
|
||||
"!" * 3000,
|
||||
"😀" * 1000,
|
||||
"漢" * 1000,
|
||||
"\u00a0abc\u00a0! \u2028x \u3000y \u200bz \u200d\u200d q",
|
||||
"x\u0085y \x0b\x0c z",
|
||||
"\x00\x01\x02 \x7f \ufffd",
|
||||
"tab\tseparated\tvalues\n1\t2\t3\n",
|
||||
"MiXeD cAsE wOrDs AND ACRONYMS like NASA, HTTP/2, gRPC, iOS, macOS",
|
||||
"snake_case_identifier camelCaseIdentifier PascalCaseIdentifier SCREAMING_SNAKE_CASE kebab-case",
|
||||
"x'sy x'ty x'rey x'vey x'my x'lly x'dy x'S x'T x'RE x'VE x'M x'LL x'D x'sS x'llL",
|
||||
"IT'SOK it'Dbe x'Sy x'Ty x'My x'Dy x'LLy x'VEy x'REy x'Ly x'Vy x'Ry 'Sx'Tx'Mx'LLx'VEx'REx'Dx",
|
||||
"'s't're've'm'll'd 'S'T'RE'VE'M'LL'D ''s '''s",
|
||||
"9'9 9's a'9 '9 ' 's' ' 's",
|
||||
"١٢٣٤ ½⅓¼ ⅣⅤ 𝟘𝟙𝟚𝟛𝟜𝟝𝟞𝟟𝟠𝟡 ①②③",
|
||||
"camelCase PascalCase ABCdef ABCdeF ABC aB Ab ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyzABC",
|
||||
"日本ABC ABC日本 日本語abc abc日本語 漢字Kanji kanji漢字 KANJI漢字kanji مرحباABC ABCمرحبا abcمرحبا",
|
||||
"\u0301ABC \u0301abc \u0301\u0301A A\u0301\u0301 E\u0301A aE\u0301 !!\u0301a \u00a0\u0301A x\u0308Y X\u0308y",
|
||||
"ᵃbc ᵃBC Aᵃbc Aᵃ ᵃ' ᵃ's Džungla aDžB ADžB ADžb DžDž Ljx İi ΣΊΣΥΦΟΣσ ΣσΣ",
|
||||
"don'tx ABC's abc'S abc'ſ ABC'ſx IT'SOK it'Dbe 'sabc x's 's 'Sx'Tx x’s X'LLx X'Ll",
|
||||
"!ABC !AbC !!abc #camelCase (ABCdef) \u00a0ABC\u00a0abc\u00a0Abc \tABC\tabc",
|
||||
"!!/\n/x a/b !!\n/x /x // path/to/file.rs http://x.y/z?a=b/c \\/\\/ //\r\n//\n",
|
||||
"x \n x \r\n \r\n y x \n a b \n\n c x\t\ty x\t\t end \n \n",
|
||||
"12345 6 1abc abc1 ABC123abc 123ABC ١٢٣٤٥abc",
|
||||
)
|
||||
|
||||
WORDS: Final = (
|
||||
"the",
|
||||
"quick",
|
||||
"brown",
|
||||
"fox",
|
||||
"jumps",
|
||||
"over",
|
||||
"lazy",
|
||||
"dog",
|
||||
"while",
|
||||
"counting",
|
||||
"tokens",
|
||||
"for",
|
||||
"budget",
|
||||
"reservation",
|
||||
"before",
|
||||
"admission",
|
||||
"on",
|
||||
"the",
|
||||
"gateway",
|
||||
"and",
|
||||
"every",
|
||||
"request",
|
||||
"body",
|
||||
"is",
|
||||
"scanned",
|
||||
"exactly",
|
||||
"once",
|
||||
"with",
|
||||
"a",
|
||||
"hand",
|
||||
"written",
|
||||
"piece",
|
||||
"scanner",
|
||||
"that",
|
||||
"mirrors",
|
||||
"tiktoken's",
|
||||
"regex",
|
||||
"boundaries",
|
||||
"It's",
|
||||
"faster",
|
||||
"because",
|
||||
"there's",
|
||||
"no",
|
||||
"backtracking",
|
||||
"engine",
|
||||
"involved",
|
||||
"so",
|
||||
"we'll",
|
||||
"keep",
|
||||
"it",
|
||||
"that",
|
||||
"way",
|
||||
"Zürich",
|
||||
"café",
|
||||
"naïve",
|
||||
"東京",
|
||||
"مرحبا",
|
||||
"🙂",
|
||||
"42",
|
||||
"1999",
|
||||
"3.14159",
|
||||
"$1,234.56",
|
||||
"100%",
|
||||
"user@example.com",
|
||||
"https://example.com/a/b?c=d",
|
||||
"C++",
|
||||
"F#",
|
||||
"node.js",
|
||||
"v1.2.3",
|
||||
"(parens)",
|
||||
"[brackets]",
|
||||
"{braces}",
|
||||
"<tags>",
|
||||
'"quotes"',
|
||||
"'single'",
|
||||
"don't",
|
||||
"WON'T",
|
||||
"I'M",
|
||||
"They'RE",
|
||||
)
|
||||
|
||||
|
||||
def random_text(rng: random.Random, alphabet: tuple[str, ...]) -> str:
|
||||
return "".join(rng.choice(alphabet) for _ in range(rng.randrange(0, 40)))
|
||||
|
||||
|
||||
def paragraph(rng: random.Random, words: int) -> str:
|
||||
return " ".join(rng.choice(WORDS) for _ in range(words))
|
||||
|
||||
|
||||
def short_paragraphs(rng: random.Random) -> Iterator[str]:
|
||||
while True:
|
||||
content = paragraph(rng, rng.randrange(60, 140))
|
||||
if len(content) < TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS:
|
||||
yield content
|
||||
|
||||
|
||||
def chat_body(rng: random.Random, target_tokens: int) -> dict[str, object]:
|
||||
candidates: Final = tuple(itertools.islice(short_paragraphs(rng), 2000))
|
||||
running: Final = tuple(itertools.accumulate(len(ENCODING.encode(content)) + 3 for content in candidates))
|
||||
turns: Final = next(index for index, total in enumerate(running) if total >= target_tokens) + 1
|
||||
contents: Final = candidates[: turns + (turns % 2)]
|
||||
return {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant. Answer precisely and cite sources."},
|
||||
*(
|
||||
{"role": "user" if index % 2 == 0 else "assistant", "content": content}
|
||||
for index, content in enumerate(contents)
|
||||
),
|
||||
{"role": "user", "content": "Summarise the conversation so far in three sentences."},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
TOOLS: Final = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {"type": "string", "description": "City name"},
|
||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||||
"days": {"type": "integer"},
|
||||
"tags": {"type": "array", "items": {"type": "string"}},
|
||||
"opts": {
|
||||
"type": "object",
|
||||
"properties": {"verbose": {"type": "boolean"}, "level": {"type": "integer", "enum": [1, 2]}},
|
||||
"required": ["verbose"],
|
||||
},
|
||||
"anything": {},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{"type": "function", "function": {"name": "noop"}},
|
||||
]
|
||||
|
||||
SMALL_REQUESTS: Final = (
|
||||
{"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]},
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a terse assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"name": "alice",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarise this paragraph about ships and harbours."},
|
||||
"plain string item",
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "Sure."}]},
|
||||
],
|
||||
},
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "weather?"}],
|
||||
"tools": TOOLS,
|
||||
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
|
||||
},
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "system", "content": "sys"}, {"role": "user", "content": "weather?"}],
|
||||
"tools": TOOLS,
|
||||
"tool_choice": "none",
|
||||
},
|
||||
{"model": MODEL, "prompt": "Write a haiku about ships."},
|
||||
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
|
||||
{
|
||||
"model": MODEL,
|
||||
"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",
|
||||
},
|
||||
{"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
|
||||
{
|
||||
"model": MODEL,
|
||||
"query": "best harbour",
|
||||
"documents": [
|
||||
"doc one",
|
||||
{"text": "doc two", "title": "T", "n": 3, "ok": True, "none": None, "tags": ["a", "b"]},
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rng: Final = random.Random(2026)
|
||||
case_rng: Final = random.Random(200_000)
|
||||
texts: Final = (
|
||||
tuple(CORPUS)
|
||||
+ tuple(random_text(rng, ALPHABET) for _ in range(3000))
|
||||
+ tuple(random_text(case_rng, CASE_ALPHABET) for _ in range(1000))
|
||||
)
|
||||
OUT.mkdir(exist_ok=True)
|
||||
with (OUT / "texts.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for text in texts:
|
||||
tokens = len(ENCODING.encode(text, disallowed_special=()))
|
||||
pieces = SPLIT_PATTERN.findall(text)
|
||||
handle.write(json.dumps({"text": text, "tokens": tokens, "pieces": pieces}, ensure_ascii=False) + "\n")
|
||||
bodies: Final = tuple(SMALL_REQUESTS) + (chat_body(rng, 50_000),)
|
||||
with (OUT / "requests.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for body in bodies:
|
||||
input_tokens = _count_input_tokens(dict(body), MODEL)
|
||||
assert input_tokens is not None
|
||||
handle.write(json.dumps({"body": json.dumps(body), "input_tokens": input_tokens}) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
10
litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl
vendored
Normal file
10
litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
4053
litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl
vendored
Normal file
4053
litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -1,4 +1,5 @@
|
|||
use rstest::rstest;
|
||||
use serde::Deserialize;
|
||||
|
||||
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
|
||||
|
||||
|
|
@ -174,3 +175,147 @@ fn tool_choice_and_system_discount_change_the_count() {
|
|||
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(_))));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -14,6 +15,20 @@ except (ImportError, AttributeError):
|
|||
|
||||
filename = pkg_resources.resource_filename(__name__, "litellm_core_utils/tokenizers")
|
||||
|
||||
CL100K_BASE_RANK_FILE: Final = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
|
||||
O200K_BASE_RANK_FILE: Final = "fb374d419588a4632f3f557e76b4b70aebbca790"
|
||||
|
||||
|
||||
def cl100k_base_rank_file() -> str:
|
||||
"""The vendored tiktoken `cl100k_base` rank file (`base64(token) rank` lines)."""
|
||||
return Path(filename, CL100K_BASE_RANK_FILE).read_text(encoding="ascii")
|
||||
|
||||
|
||||
def o200k_base_rank_file() -> str:
|
||||
"""The vendored tiktoken `o200k_base` rank file (`base64(token) rank` lines)."""
|
||||
return Path(filename, O200K_BASE_RANK_FILE).read_text(encoding="ascii")
|
||||
|
||||
|
||||
# Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory
|
||||
# unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR.
|
||||
# This keeps tiktoken fully offline-capable by default (see #1071).
|
||||
|
|
|
|||
|
|
@ -381,7 +381,7 @@ class _MessageCountParams:
|
|||
from litellm.utils import print_verbose
|
||||
|
||||
actual_model: Final = _fix_model_name(model)
|
||||
if actual_model == "gpt-3.5-turbo-0301":
|
||||
if uses_legacy_message_accounting(model):
|
||||
self.tokens_per_message = 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
|
||||
self.tokens_per_name = -1 # if there's a name, the role is omitted
|
||||
elif actual_model in litellm.open_ai_chat_completion_models or actual_model in litellm.azure_llms:
|
||||
|
|
@ -615,7 +615,7 @@ def _get_exact_count_function(
|
|||
) -> TokenCounterFunction:
|
||||
"""
|
||||
Get the function to count tokens based on the model and custom tokenizer."""
|
||||
from litellm.utils import _select_tokenizer, print_verbose
|
||||
from litellm.utils import _select_tokenizer
|
||||
|
||||
if model is not None or custom_tokenizer is not None:
|
||||
tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
|
||||
|
|
@ -627,15 +627,7 @@ def _get_exact_count_function(
|
|||
|
||||
return count_tokens
|
||||
elif tokenizer_json["type"] == "openai_tokenizer":
|
||||
model_to_use: Final = _fix_model_name(model)
|
||||
try:
|
||||
if "gpt-4o" in model_to_use:
|
||||
encoding = tiktoken.get_encoding("o200k_base")
|
||||
else:
|
||||
encoding = tiktoken.encoding_for_model(model_to_use)
|
||||
except KeyError:
|
||||
print_verbose("Warning: model not found. Using cl100k_base encoding.")
|
||||
encoding = tiktoken.get_encoding("cl100k_base")
|
||||
encoding: Final = openai_tokenizer_encoding(model)
|
||||
|
||||
def encode_length(text: str) -> int:
|
||||
return len(encoding.encode(text, disallowed_special=()))
|
||||
|
|
@ -651,6 +643,25 @@ def _get_exact_count_function(
|
|||
return _get_tiktoken_count_function(encode_length)
|
||||
|
||||
|
||||
def openai_tokenizer_encoding(model: str) -> tiktoken.Encoding:
|
||||
"""The tiktoken encoding `token_counter` uses for a model on the `openai_tokenizer` path."""
|
||||
from litellm.utils import print_verbose
|
||||
|
||||
model_to_use: Final = _fix_model_name(model)
|
||||
if "gpt-4o" in model_to_use:
|
||||
return tiktoken.get_encoding("o200k_base")
|
||||
try:
|
||||
return tiktoken.encoding_for_model(model_to_use)
|
||||
except KeyError:
|
||||
print_verbose("Warning: model not found. Using cl100k_base encoding.")
|
||||
return tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def uses_legacy_message_accounting(model: str) -> bool:
|
||||
"""Whether `token_counter` prices messages with the `gpt-3.5-turbo-0301` constants (4 per message, -1 per name)."""
|
||||
return _fix_model_name(model) == "gpt-3.5-turbo-0301"
|
||||
|
||||
|
||||
def _fix_model_name(model: str) -> str:
|
||||
"""We normalize some model names to others"""
|
||||
if model in litellm.azure_llms:
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.rust_bridge.token_counter import count_anthropic_input_tokens, uses_anthropic_tokenizer
|
||||
from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer
|
||||
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
|
|
@ -1365,8 +1365,9 @@ async def count_request_input_tokens(
|
|||
|
||||
Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so
|
||||
counting a large prompt inline stalls every other request on the worker.
|
||||
Models on the Anthropic tokenizer are counted from the raw body by the Rust
|
||||
bridge when it is enabled, which parses and tokenizes with the GIL released.
|
||||
Models whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base
|
||||
and o200k_base) are counted from the raw body by the bridge when it is enabled, once per
|
||||
distinct tokenizer, which parses and tokenizes with the GIL released.
|
||||
Everything it declines is counted in Python, large prompts in a worker
|
||||
thread. The counts are reused by both the max-cost and the input-cost
|
||||
estimate.
|
||||
|
|
@ -1374,23 +1375,31 @@ async def count_request_input_tokens(
|
|||
models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router)
|
||||
if not models:
|
||||
return MappingProxyType({})
|
||||
rust_count: Final = (
|
||||
await count_anthropic_input_tokens(raw_body)
|
||||
if raw_body is not None and any(uses_anthropic_tokenizer(model) for model in models)
|
||||
else None
|
||||
tokenizers: Final[Mapping[str, RustTokenizer | None]] = MappingProxyType(
|
||||
{model: rust_tokenizer(model) for model in models}
|
||||
)
|
||||
distinct_tokenizers: Final[tuple[RustTokenizer, ...]] = tuple(
|
||||
dict.fromkeys(tokenizer for tokenizer in tokenizers.values() if tokenizer is not None)
|
||||
)
|
||||
rust_counts_by_tokenizer: Final[Mapping[RustTokenizer, int]] = MappingProxyType(
|
||||
{
|
||||
tokenizer: count.input_tokens
|
||||
for tokenizer in distinct_tokenizers
|
||||
if raw_body is not None and (count := await count_input_tokens(raw_body, tokenizer)) is not None
|
||||
}
|
||||
)
|
||||
rust_counts: Final = MappingProxyType(
|
||||
{
|
||||
model: rust_count.input_tokens
|
||||
for model in models
|
||||
if rust_count is not None and uses_anthropic_tokenizer(model)
|
||||
model: rust_counts_by_tokenizer[tokenizer]
|
||||
for model, tokenizer in tokenizers.items()
|
||||
if tokenizer is not None and tokenizer in rust_counts_by_tokenizer
|
||||
}
|
||||
)
|
||||
python_models: Final = tuple(model for model in models if model not in rust_counts)
|
||||
if not python_models:
|
||||
return rust_counts
|
||||
python_counts: Final = (
|
||||
_count_input_tokens_for_models(request_body=request_body, models=python_models)
|
||||
MappingProxyType({})
|
||||
if not python_models
|
||||
else _count_input_tokens_for_models(request_body=request_body, models=python_models)
|
||||
if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS
|
||||
else await asyncio.to_thread(
|
||||
_count_input_tokens_for_models,
|
||||
|
|
@ -1398,6 +1407,7 @@ async def count_request_input_tokens(
|
|||
models=python_models,
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.debug("input token counts: rust=%s python=%s", dict(rust_counts), dict(python_counts))
|
||||
return MappingProxyType({**rust_counts, **python_counts})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,17 +5,20 @@ from __future__ import annotations
|
|||
from collections.abc import Awaitable
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
|
||||
from typing import Final, Literal, Protocol, cast # noqa: TID251 # native extension exposes untyped callables
|
||||
|
||||
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
|
||||
from litellm.utils import uses_anthropic_tokenizer as _python_uses_anthropic_tokenizer
|
||||
from litellm.utils import claude_json_str, huggingface_tokenizer_kind
|
||||
|
||||
RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"]
|
||||
|
||||
|
||||
class RustTokenCounter(Protocol):
|
||||
|
|
@ -27,6 +30,12 @@ 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:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InputTokenCount:
|
||||
|
|
@ -50,18 +59,41 @@ def _as_factory(value: object) -> RustTokenCounterFactory | None:
|
|||
TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory)
|
||||
|
||||
|
||||
def uses_anthropic_tokenizer(model: str) -> bool:
|
||||
if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True:
|
||||
return False
|
||||
return _python_uses_anthropic_tokenizer(model)
|
||||
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."""
|
||||
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)
|
||||
if kind == "anthropic":
|
||||
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
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def _anthropic_counter(factory: RustTokenCounterFactory) -> RustTokenCounter:
|
||||
return factory(claude_json_str)
|
||||
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())
|
||||
|
||||
|
||||
async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
|
||||
async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
|
||||
if not rust_enabled():
|
||||
return None
|
||||
factory: Final = TOKEN_COUNTER.load()
|
||||
|
|
@ -69,11 +101,14 @@ async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
|
|||
return None
|
||||
try:
|
||||
attempt: Final = await aattempt(
|
||||
native_call=lambda: _anthropic_counter(factory).acount_request(body),
|
||||
native_call=lambda: _counter(factory, tokenizer).acount_request(body),
|
||||
adapt=_INPUT_TOKEN_COUNT.validate_python,
|
||||
context=BridgeErrorContext(route="token_counter", provider="anthropic", model=""),
|
||||
context=BridgeErrorContext(route="token_counter", provider=tokenizer, model=""),
|
||||
)
|
||||
except (RuntimeError, ValueError) as error:
|
||||
verbose_logger.debug("Rust token counter failed, counting in Python: %s", error)
|
||||
verbose_logger.debug("Rust token counter (%s) failed, counting in Python: %s", tokenizer, error)
|
||||
return None
|
||||
return attempt.value if isinstance(attempt, RustHandled) else None
|
||||
if not isinstance(attempt, RustHandled):
|
||||
return None
|
||||
verbose_logger.debug("Rust token counter (%s) counted %d input tokens", tokenizer, attempt.value.input_tokens)
|
||||
return attempt.value
|
||||
|
|
|
|||
|
|
@ -2227,25 +2227,39 @@ def uses_anthropic_tokenizer(model: str) -> bool:
|
|||
return model in litellm.anthropic_models and "claude-3" not in model
|
||||
|
||||
|
||||
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
|
||||
HuggingFaceTokenizerKind = Literal["cohere", "anthropic", "llama2", "llama3"]
|
||||
|
||||
|
||||
def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
|
||||
"""Which HuggingFace tokenizer `token_counter` selects for a model; `None` means tiktoken."""
|
||||
if model in litellm.cohere_models and "command-r" in model:
|
||||
# cohere
|
||||
cohere_tokenizer: Final = Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer}
|
||||
# anthropic
|
||||
elif uses_anthropic_tokenizer(model):
|
||||
claude_tokenizer: Final = Tokenizer.from_str(claude_json_str)
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer}
|
||||
# llama2
|
||||
elif "llama-2" in model.lower() or "replicate" in model.lower():
|
||||
tokenizer = Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
|
||||
# llama3
|
||||
elif "llama-3" in model.lower():
|
||||
tokenizer = Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
|
||||
else:
|
||||
return "cohere"
|
||||
if uses_anthropic_tokenizer(model):
|
||||
return "anthropic"
|
||||
if "llama-2" in model.lower() or "replicate" in model.lower():
|
||||
return "llama2"
|
||||
if "llama-3" in model.lower():
|
||||
return "llama3"
|
||||
return None
|
||||
|
||||
|
||||
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
|
||||
kind: Final = huggingface_tokenizer_kind(model)
|
||||
if kind is None:
|
||||
return None
|
||||
return {"type": "huggingface_tokenizer", "tokenizer": _load_huggingface_tokenizer(kind)}
|
||||
|
||||
|
||||
def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> Tokenizer:
|
||||
match kind:
|
||||
case "cohere":
|
||||
return Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
|
||||
case "anthropic":
|
||||
return Tokenizer.from_str(claude_json_str)
|
||||
case "llama2":
|
||||
return Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
|
||||
case "llama3":
|
||||
return Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
|
||||
|
||||
|
||||
def encode(model="", text="", custom_tokenizer: dict | None = None):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -207,8 +210,13 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None:
|
|||
|
||||
|
||||
ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
CL100K_MODEL: Final = "gpt-4"
|
||||
O200K_MODEL: Final = "gpt-4o"
|
||||
RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
|
||||
RUST_INPUT_TOKENS: Final = 4_321
|
||||
RUST_INPUT_TOKENS_BY_TOKENIZER: Final = MappingProxyType(
|
||||
{"anthropic": RUST_INPUT_TOKENS, "cl100k_base": 1_234, "o200k_base": 2_345}
|
||||
)
|
||||
|
||||
|
||||
class _FakeDeclined(Exception):
|
||||
|
|
@ -225,33 +233,57 @@ class _FakeNative:
|
|||
|
||||
|
||||
class _RecordingCounter:
|
||||
bodies: Final[list[bytes]] = []
|
||||
"""Stands in for one native counter; records `(tokenizer, body)` on the shared factory."""
|
||||
|
||||
def __init__(self, tokenizer_json: str) -> None:
|
||||
pass
|
||||
def __init__(self, factory: _RecordingFactory, tokenizer: rust_token_counter.RustTokenizer) -> None:
|
||||
self.factory = factory
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
self.bodies.append(body)
|
||||
return {"model": ANTHROPIC_TOKENIZER_MODEL, "input_tokens": RUST_INPUT_TOKENS}
|
||||
self.factory.calls.append((self.tokenizer, body))
|
||||
return {"model": "", "input_tokens": RUST_INPUT_TOKENS_BY_TOKENIZER[self.tokenizer]}
|
||||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
|
||||
|
||||
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")
|
||||
|
||||
|
||||
class _DecliningCounter:
|
||||
def __init__(self, tokenizer_json: str) -> None:
|
||||
pass
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise _FakeDeclined("unsupported content block")
|
||||
|
||||
|
||||
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:
|
||||
return _DecliningCounter()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rust_counter(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
|
||||
rust_token_counter._anthropic_counter.cache_clear()
|
||||
rust_token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
_RecordingCounter.bodies.clear()
|
||||
yield
|
||||
rust_token_counter.TOKEN_COUNTER.reset()
|
||||
rust_token_counter._anthropic_counter.cache_clear()
|
||||
rust_token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
|
|
@ -270,8 +302,9 @@ def rust_counter(monkeypatch: pytest.MonkeyPatch):
|
|||
async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
|
||||
rust_counter: None, route: str, request_body: dict
|
||||
) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
raw_body: Final = json.dumps(request_body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
|
|
@ -279,53 +312,136 @@ async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
|
|||
)
|
||||
|
||||
assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS}
|
||||
assert _RecordingCounter.bodies == [raw_body]
|
||||
assert factory.calls == [("anthropic", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rust_decline_falls_back_to_python_count(rust_counter: None) -> None:
|
||||
@pytest.mark.parametrize("model", (CL100K_MODEL, "azure/gpt-35-turbo", "gemini/gemini-2.5-pro", "my-router-alias"))
|
||||
async def test_tiktoken_cl100k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_DecliningCounter)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"]}
|
||||
assert factory.calls == [("cl100k_base", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", (O200K_MODEL, "gpt-5", "o3", "gpt-4.1", "chatgpt-4o-latest"))
|
||||
async def test_tiktoken_o200k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"]}
|
||||
assert factory.calls == [("o200k_base", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_model_request_counts_once_per_tokenizer_and_python_for_the_rest(rust_counter: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
models: Final = (
|
||||
CL100K_MODEL,
|
||||
ANTHROPIC_TOKENIZER_MODEL,
|
||||
"gemini/gemini-2.5-pro",
|
||||
O200K_MODEL,
|
||||
"gpt-5",
|
||||
"replicate/meta/llama-2-70b-chat",
|
||||
)
|
||||
body: Final = {"model": list(models), "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
python_counts: Final = await count_request_input_tokens(
|
||||
request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None
|
||||
)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=RUST_COUNTED_BODY,
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert factory.calls == [("cl100k_base", raw_body), ("anthropic", raw_body), ("o200k_base", raw_body)]
|
||||
assert dict(counts) == {
|
||||
CL100K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
|
||||
"gemini/gemini-2.5-pro": RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
|
||||
ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS,
|
||||
O200K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
|
||||
"gpt-5": RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
|
||||
"replicate/meta/llama-2-70b-chat": python_counts["replicate/meta/llama-2-70b-chat"],
|
||||
}
|
||||
assert counts["replicate/meta/llama-2-70b-chat"] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", (ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL))
|
||||
async def test_rust_decline_falls_back_to_python_count(rust_counter: None, model: str) -> None:
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_DecliningFactory())
|
||||
body: Final = {**RUST_COUNTED_BODY, "model": model}
|
||||
python_counts: Final = await count_request_input_tokens(request_body=body, route="/v1/messages", llm_router=None)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body,
|
||||
route="/v1/messages",
|
||||
llm_router=None,
|
||||
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
|
||||
raw_body=json.dumps(body).encode(),
|
||||
)
|
||||
|
||||
assert dict(counts) == dict(python_counts)
|
||||
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
|
||||
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL], "messages": ANTHROPIC_MESSAGES}
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=RUST_COUNTED_BODY,
|
||||
route="/v1/messages",
|
||||
request_body=body,
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
|
||||
raw_body=json.dumps(body).encode(),
|
||||
)
|
||||
|
||||
assert _RecordingCounter.bodies == []
|
||||
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
|
||||
assert factory.calls == []
|
||||
assert set(counts) == {ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL}
|
||||
assert not set(counts.values()) & set(RUST_INPUT_TOKENS_BY_TOKENIZER.values())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_anthropic_tokenizer_models_stay_in_python(rust_counter: None) -> None:
|
||||
@pytest.mark.parametrize("model", ("replicate/meta/llama-2-70b-chat", "meta-llama/Llama-3-8b", "text-davinci-003"))
|
||||
async def test_models_without_a_rust_tokenizer_stay_in_python(
|
||||
rust_counter: None, monkeypatch: pytest.MonkeyPatch, model: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"text-davinci-003"}
|
||||
)
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
|
||||
body: Final = {"model": "gpt-4o", "messages": ANTHROPIC_MESSAGES}
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
python_counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None
|
||||
)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=json.dumps(body).encode()
|
||||
)
|
||||
|
||||
assert _RecordingCounter.bodies == []
|
||||
assert counts["gpt-4o"] != RUST_INPUT_TOKENS
|
||||
assert factory.calls == []
|
||||
assert dict(counts) == dict(python_counts)
|
||||
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
|
|
|||
|
|
@ -8,16 +8,26 @@ 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
|
||||
import tiktoken
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
import litellm
|
||||
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
|
||||
from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
|
||||
from litellm.rust_bridge import bindings, configuration
|
||||
from litellm.rust_bridge import token_counter as bridge
|
||||
from litellm.utils import claude_json_str
|
||||
|
||||
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()
|
||||
|
||||
|
||||
|
|
@ -44,69 +54,122 @@ class _RecordingCounter:
|
|||
return {"model": MODEL, "input_tokens": 42}
|
||||
|
||||
|
||||
class _DecliningCounter:
|
||||
def __init__(self, tokenizer_json: str) -> None:
|
||||
pass
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counters: list[_RecordingCounter] = []
|
||||
self.rank_files: 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")
|
||||
|
||||
|
||||
class _RaisingCounter:
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.error = error
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise _FakeDeclined("request has no messages")
|
||||
raise self.error
|
||||
|
||||
|
||||
class _FailingCounter:
|
||||
def __init__(self, tokenizer_json: str) -> None:
|
||||
pass
|
||||
class _RaisingFactory:
|
||||
"""Every counter it builds, for either tokenizer, raises `error` on count."""
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise RuntimeError("encode failed")
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.error = error
|
||||
|
||||
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:
|
||||
return _RaisingCounter(self.error)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
|
||||
bridge.TOKEN_COUNTER.reset()
|
||||
bridge._anthropic_counter.cache_clear()
|
||||
bridge._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
|
||||
yield
|
||||
bridge.TOKEN_COUNTER.reset()
|
||||
bridge._anthropic_counter.cache_clear()
|
||||
bridge._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_bridge_never_constructs_a_counter() -> None:
|
||||
constructed: list[str] = []
|
||||
|
||||
def factory(tokenizer_json: str) -> _RecordingCounter:
|
||||
constructed.append(tokenizer_json)
|
||||
return _RecordingCounter(tokenizer_json)
|
||||
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.RustTokenizer) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
assert await bridge.count_anthropic_input_tokens(BODY) is None
|
||||
assert constructed == []
|
||||
assert await bridge.count_input_tokens(BODY, tokenizer) is None
|
||||
assert factory.counters == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> None:
|
||||
counters: list[_RecordingCounter] = []
|
||||
|
||||
def factory(tokenizer_json: str) -> _RecordingCounter:
|
||||
counter = _RecordingCounter(tokenizer_json)
|
||||
counters.append(counter)
|
||||
return counter
|
||||
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
first: Final = await bridge.count_anthropic_input_tokens(BODY)
|
||||
second: Final = await bridge.count_anthropic_input_tokens(BODY)
|
||||
first: Final = await bridge.count_input_tokens(BODY, "anthropic")
|
||||
second: Final = await bridge.count_input_tokens(BODY, "anthropic")
|
||||
|
||||
assert first == bridge.InputTokenCount(model=MODEL, input_tokens=42)
|
||||
assert second == first
|
||||
assert len(counters) == 1
|
||||
assert counters[0].bodies == [BODY, BODY]
|
||||
assert json.loads(counters[0].tokenizer_json)["model"]["type"] == "BPE"
|
||||
assert len(factory.counters) == 1
|
||||
assert factory.counters[0].bodies == [BODY, BODY]
|
||||
assert json.loads(factory.counters[0].tokenizer_json)["model"]["type"] == "BPE"
|
||||
|
||||
|
||||
@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:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
first: Final = await bridge.count_input_tokens(BODY, tokenizer)
|
||||
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.counters[0].tokenizer_json == tokenizer
|
||||
assert factory.counters[0].bodies == [BODY, BODY]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
await bridge.count_input_tokens(BODY, "anthropic")
|
||||
await bridge.count_input_tokens(BODY, "cl100k_base")
|
||||
await bridge.count_input_tokens(BODY, "o200k_base")
|
||||
await bridge.count_input_tokens(BODY, "anthropic")
|
||||
await bridge.count_input_tokens(BODY, "o200k_base")
|
||||
|
||||
assert [counter.tokenizer_json for counter in factory.counters][1:] == ["cl100k_base", "o200k_base"]
|
||||
assert [len(counter.bodies) for counter in factory.counters] == [2, 1, 2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -114,38 +177,131 @@ async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch)
|
|||
litellm.rust(True)
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
|
||||
|
||||
assert await bridge.count_anthropic_input_tokens(BODY) is None
|
||||
assert [await bridge.count_input_tokens(BODY, tokenizer) for tokenizer in TOKENIZERS] == [None, None, None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_declined_request_falls_back() -> None:
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_DecliningCounter)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages")))
|
||||
|
||||
assert await bridge.count_anthropic_input_tokens(BODY) is None
|
||||
assert await bridge.count_input_tokens(BODY, tokenizer) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_failure_falls_back() -> None:
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_FailingCounter)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed")))
|
||||
|
||||
assert await bridge.count_anthropic_input_tokens(BODY) is None
|
||||
assert await bridge.count_input_tokens(BODY, tokenizer) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected"),
|
||||
((MODEL, True), ("claude-3-5-sonnet-20241022", False), ("gpt-4o", False), ("my-router-alias", False)),
|
||||
(
|
||||
(MODEL, "anthropic"),
|
||||
("claude-3-5-sonnet-20241022", "cl100k_base"),
|
||||
("gpt-4", "cl100k_base"),
|
||||
("gpt-4-turbo", "cl100k_base"),
|
||||
("gpt-3.5-turbo", "cl100k_base"),
|
||||
("azure/gpt-35-turbo", "cl100k_base"),
|
||||
("gemini/gemini-2.5-pro", "cl100k_base"),
|
||||
("mistral/mistral-large-latest", "cl100k_base"),
|
||||
("my-router-alias", "cl100k_base"),
|
||||
("azure/gpt-4o", "cl100k_base"),
|
||||
("command-r-plus", "cl100k_base"),
|
||||
("gpt-4o", "o200k_base"),
|
||||
("gpt-4o-mini", "o200k_base"),
|
||||
("gpt-4o-2024-08-06", "o200k_base"),
|
||||
("chatgpt-4o-latest", "o200k_base"),
|
||||
("gpt-4.1", "o200k_base"),
|
||||
("gpt-5", "o200k_base"),
|
||||
("gpt-5-mini", "o200k_base"),
|
||||
("o1", "o200k_base"),
|
||||
("o3", "o200k_base"),
|
||||
("o3-mini", "o200k_base"),
|
||||
("o4-mini", "o200k_base"),
|
||||
("replicate/meta/llama-2-70b-chat", None),
|
||||
("meta-llama/Llama-3-8b", None),
|
||||
),
|
||||
)
|
||||
def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bool) -> None:
|
||||
assert bridge.uses_anthropic_tokenizer(model) is expected
|
||||
def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bridge.RustTokenizer | None) -> None:
|
||||
assert bridge.rust_tokenizer(model) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter"))
|
||||
def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None:
|
||||
monkeypatch.setattr(litellm, flag, True)
|
||||
@pytest.mark.parametrize(
|
||||
("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
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model})
|
||||
|
||||
assert bridge.uses_anthropic_tokenizer(MODEL) is False
|
||||
assert openai_tokenizer_encoding(model).name == python_encoding
|
||||
assert bridge.rust_tokenizer(model) is None
|
||||
|
||||
|
||||
def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "cohere_models", litellm.cohere_models | {"command-r-plus"})
|
||||
|
||||
assert bridge.rust_tokenizer("command-r-plus") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("legacy_model", ("gpt-3.5-turbo-0301", "gpt-35-turbo-0301"))
|
||||
def test_rust_tokenizer_declines_legacy_message_accounting_python_prices_differently(
|
||||
monkeypatch: pytest.MonkeyPatch, legacy_model: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"gpt-3.5-turbo-0301"}
|
||||
)
|
||||
monkeypatch.setattr(litellm, "azure_llms", {**litellm.azure_llms, "gpt-35-turbo-0301": "azure"})
|
||||
messages: Final = [{"role": "user", "name": "bob", "content": "hello there"}]
|
||||
|
||||
assert litellm.token_counter(model=legacy_model, messages=messages) != litellm.token_counter(
|
||||
model=CL100K_MODEL, messages=messages
|
||||
)
|
||||
assert bridge.rust_tokenizer(legacy_model) is None
|
||||
assert bridge.rust_tokenizer(CL100K_MODEL) == "cl100k_base"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", (MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", "o3"))
|
||||
def test_rust_tokenizer_names_the_encoding_python_actually_counts_with(model: str) -> None:
|
||||
text: Final = (
|
||||
"Hello, world! camelCase ABCdef \u00e9\u00e8 12345 \u3053\u3093\u306b\u3061\u306f <|endoftext|>\r\n" * 9
|
||||
)
|
||||
python_count: Final = litellm.token_counter(model=model, text=text)
|
||||
cl100k_count: Final = len(tiktoken.get_encoding("cl100k_base").encode(text, disallowed_special=()))
|
||||
o200k_count: Final = len(tiktoken.get_encoding("o200k_base").encode(text, disallowed_special=()))
|
||||
assert cl100k_count != o200k_count
|
||||
match bridge.rust_tokenizer(model):
|
||||
case "cl100k_base":
|
||||
assert python_count == cl100k_count
|
||||
case "o200k_base":
|
||||
assert python_count == o200k_count
|
||||
case "anthropic":
|
||||
assert python_count == len(Tokenizer.from_str(claude_json_str).encode(text).ids)
|
||||
assert python_count not in {cl100k_count, o200k_count}
|
||||
case None:
|
||||
pytest.fail(f"{model} must have a Rust tokenizer")
|
||||
|
||||
|
||||
def test_disabled_hf_download_routes_anthropic_models_to_cl100k_like_python(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
|
||||
|
||||
assert bridge.rust_tokenizer(MODEL) == "cl100k_base"
|
||||
assert bridge.rust_tokenizer("meta-llama/Llama-3-8b") == "cl100k_base"
|
||||
assert bridge.rust_tokenizer(O200K_MODEL) == "o200k_base"
|
||||
|
||||
|
||||
def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_token_counter", True)
|
||||
|
||||
assert bridge.rust_tokenizer(MODEL) is None
|
||||
assert bridge.rust_tokenizer(CL100K_MODEL) is None
|
||||
assert bridge.rust_tokenizer(O200K_MODEL) is None
|
||||
|
||||
|
||||
PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
||||
|
|
@ -182,7 +338,16 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
},
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "x " * 20_000}],
|
||||
"messages": [{"role": "user", "content": "x " * 500}],
|
||||
},
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "I'VE got 1234567 things; it's \"fine\"...\r\n\r\n caf\u00e9 \u0645\u0631\u062d\u0628\u0627 \U0001f600 <|endoftext|>",
|
||||
}
|
||||
],
|
||||
},
|
||||
{"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20},
|
||||
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
|
||||
|
|
@ -190,7 +355,7 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
"model": MODEL,
|
||||
"instructions": "be terse",
|
||||
"input": [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": 'Summarise caf\u00e9 menus \u2014 "ok"?\n'}]},
|
||||
{"role": "assistant", "content": "Sure."},
|
||||
],
|
||||
},
|
||||
|
|
@ -202,27 +367,63 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
)
|
||||
|
||||
|
||||
PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
|
||||
(MODEL, "anthropic"),
|
||||
(CL100K_MODEL, "cl100k_base"),
|
||||
(O200K_MODEL, "o200k_base"),
|
||||
("gpt-5", "o200k_base"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("model", "tokenizer"), PARITY_MODELS)
|
||||
@pytest.mark.parametrize("request_body", PARITY_REQUESTS)
|
||||
async def test_native_count_matches_python_budget_counter(
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
body: Final = json.dumps(request_body).replace(MODEL, model)
|
||||
|
||||
rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode())
|
||||
python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL)
|
||||
rust_count: Final = await bridge.count_input_tokens(body.encode(), tokenizer)
|
||||
python_count: Final = _count_input_tokens(request_body=json.loads(body), model=model)
|
||||
|
||||
assert rust_count is not None
|
||||
assert rust_count.model == request_body.get("model")
|
||||
assert rust_count.model == json.loads(body).get("model")
|
||||
assert rust_count.input_tokens == python_count
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("model", "tokenizer"), ((CL100K_MODEL, "cl100k_base"), (O200K_MODEL, "o200k_base")))
|
||||
async def test_tiktoken_counts_long_text_exactly_where_python_chunks(
|
||||
monkeypatch: pytest.MonkeyPatch, model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
text: Final = "x " * 20_000
|
||||
body: Final = {"model": model, "messages": [{"role": "user", "content": text}]}
|
||||
encoding: Final = tiktoken.get_encoding(tokenizer)
|
||||
exact: Final = 3 + len(encoding.encode("user")) + len(encoding.encode(text)) + 3
|
||||
chunks: Final = -(-len(text) // TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS)
|
||||
|
||||
rust_count: Final = await bridge.count_input_tokens(json.dumps(body).encode(), tokenizer)
|
||||
python_count: Final = _count_input_tokens(request_body=body, model=model)
|
||||
|
||||
assert rust_count is not None
|
||||
assert rust_count.input_tokens == exact
|
||||
assert python_count is not None
|
||||
assert exact < python_count <= exact + chunks
|
||||
|
||||
|
||||
DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
||||
{
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}],
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}
|
||||
],
|
||||
},
|
||||
{"model": MODEL, "prompt": 1.5},
|
||||
{"model": MODEL, "documents": [{"score": 0.5}]},
|
||||
|
|
@ -231,12 +432,13 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
@pytest.mark.parametrize("request_body", DECLINED_REQUESTS)
|
||||
async def test_native_declines_shapes_python_prices_differently(
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
|
||||
assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None
|
||||
assert await bridge.count_input_tokens(json.dumps(request_body).encode(), tokenizer) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue