mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b1373d2456
commit
f6d332f577
17 changed files with 5776 additions and 192 deletions
|
|
@ -38,6 +38,11 @@ impl TokenCounter {
|
|||
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>> {
|
||||
let counter = Arc::clone(&self.inner);
|
||||
let encode_slots = Arc::clone(&self.encode_slots);
|
||||
|
|
|
|||
|
|
@ -1,61 +1,14 @@
|
|||
//! Exact token counting for tiktoken's `cl100k_base`. A scanner reproduces the
|
||||
//! piece boundaries of the encoding's 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`,
|
||||
//! and each piece is merged with the rank file. Special tokens are ordinary
|
||||
//! text, as with `encode(text, disallowed_special=())`.
|
||||
//! 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 std::iter;
|
||||
|
||||
use super::tiktoken::{MergeRanks, MergeScratch};
|
||||
use super::scanner::{contraction_len, digit_run_len, is_newline};
|
||||
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
|
||||
use crate::Error;
|
||||
|
||||
const MAX_DIGITS_PER_PIECE: usize = 3;
|
||||
|
||||
pub(super) struct Cl100kCounter {
|
||||
ranks: MergeRanks,
|
||||
unicode_classes: &'static UnicodeClasses,
|
||||
}
|
||||
|
||||
impl Cl100kCounter {
|
||||
pub(super) fn from_ranks(rank_file: &str) -> Result<Self, Error> {
|
||||
Ok(Self {
|
||||
ranks: MergeRanks::parse(rank_file)?,
|
||||
unicode_classes: UnicodeClasses::get().ok_or(Error::UnicodeClasses)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn count(&self, text: &str) -> usize {
|
||||
let mut scratch = MergeScratch::default();
|
||||
pieces(text, 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.
|
||||
fn pieces<'a>(
|
||||
text: &'a str,
|
||||
unicode_classes: &'static UnicodeClasses,
|
||||
) -> impl Iterator<Item = &'a str> {
|
||||
iter::successors(split_piece(text, unicode_classes), move |(_, rest)| {
|
||||
split_piece(rest, unicode_classes)
|
||||
})
|
||||
.map(|(piece, _)| piece)
|
||||
}
|
||||
|
||||
fn split_piece<'a>(text: &'a str, unicode_classes: &UnicodeClasses) -> Option<(&'a str, &'a str)> {
|
||||
let first = text.chars().next()?;
|
||||
Some(text.split_at(piece_len(text, first, unicode_classes)))
|
||||
}
|
||||
|
||||
/// The alternatives in regex order; the possessive quantifiers mean an
|
||||
/// alternative that starts matching and runs out of input fails as a whole.
|
||||
fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
|
||||
if first == '\''
|
||||
&& let Some(len) = contraction_len(&text[1..])
|
||||
{
|
||||
return 1 + len;
|
||||
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 {
|
||||
|
|
@ -80,32 +33,6 @@ fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize
|
|||
space_run_len(text, unicode_classes)
|
||||
}
|
||||
|
||||
/// `(?i:[sdmt]|ll|ve|re)` after the apostrophe. Simple case folding also maps
|
||||
/// U+017F (long s) onto `s`.
|
||||
fn contraction_len(rest: &str) -> Option<usize> {
|
||||
let mut characters = rest.chars();
|
||||
let first = characters.next()?;
|
||||
match first {
|
||||
's' | 'S' | '\u{17F}' | 'd' | 'D' | 'm' | 'M' | 't' | 'T' => Some(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),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_newline(character: char) -> bool {
|
||||
matches!(character, '\r' | '\n')
|
||||
}
|
||||
|
||||
/// `\p{N}{1,3}+`
|
||||
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()
|
||||
}
|
||||
|
||||
/// `[^\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);
|
||||
|
|
@ -139,10 +66,12 @@ 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()
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@ use serde::Serialize;
|
|||
|
||||
use crate::Error;
|
||||
use crate::byte_level::ByteLevelCounter;
|
||||
use crate::cl100k::Cl100kCounter;
|
||||
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,
|
||||
|
|
@ -29,7 +29,7 @@ enum Encoder {
|
|||
tokenizer: Box<tokenizers::Tokenizer>,
|
||||
byte_level: Option<ByteLevelCounter>,
|
||||
},
|
||||
Cl100k(Cl100kCounter),
|
||||
Tiktoken(TiktokenCounter),
|
||||
}
|
||||
|
||||
/// A loaded tokenizer plus the message accounting Python applies on top of
|
||||
|
|
@ -57,14 +57,24 @@ impl TokenCounter {
|
|||
/// 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::Cl100k(Cl100kCounter::from_ranks(rank_file)?),
|
||||
encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn count_text(&self, text: &str) -> Result<usize, Error> {
|
||||
match &self.encoder {
|
||||
Encoder::Cl100k(counter) => Ok(counter.count(text)),
|
||||
Encoder::Tiktoken(counter) => Ok(counter.count(text)),
|
||||
Encoder::HuggingFace {
|
||||
tokenizer,
|
||||
byte_level,
|
||||
|
|
|
|||
|
|
@ -8,7 +8,9 @@ mod byte_level;
|
|||
mod cl100k;
|
||||
mod counter;
|
||||
mod error;
|
||||
mod o200k;
|
||||
mod python_json;
|
||||
mod scanner;
|
||||
mod tiktoken;
|
||||
mod tools;
|
||||
mod types;
|
||||
|
|
|
|||
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()
|
||||
}
|
||||
|
|
@ -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)?,
|
||||
})
|
||||
});
|
||||
|
||||
|
|
@ -70,6 +74,14 @@ impl UnicodeClasses {
|
|||
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
|
||||
|
|
@ -101,3 +113,52 @@ pub(super) fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeCla
|
|||
.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)
|
||||
}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,24 +1,27 @@
|
|||
"""Pin tiktoken cl100k_base reference counts for the Rust parity tests.
|
||||
"""Pin tiktoken reference counts for the Rust parity tests of one encoding.
|
||||
|
||||
Run from the repository root with the project environment:
|
||||
Run from the repository root with the project environment, once per encoding:
|
||||
|
||||
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/cl100k/generate.py
|
||||
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
|
||||
|
||||
`texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens` counted with
|
||||
`tiktoken.get_encoding("cl100k_base").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. `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, "gpt-4")`). 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.
|
||||
`<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
|
||||
|
|
@ -27,13 +30,19 @@ 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
|
||||
ENCODING: Final = tiktoken.get_encoding("cl100k_base")
|
||||
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 cl100k pattern treats differently.
|
||||
# Mirrors ALPHABET in src/byte_level.rs, plus the pieces the tiktoken patterns treat differently.
|
||||
ALPHABET: Final = (
|
||||
"a",
|
||||
"Z",
|
||||
|
|
@ -114,6 +123,24 @@ ALPHABET: Final = (
|
|||
"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?",
|
||||
|
|
@ -159,6 +186,15 @@ CORPUS: Final = (
|
|||
"'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 = (
|
||||
|
|
@ -244,8 +280,8 @@ WORDS: Final = (
|
|||
)
|
||||
|
||||
|
||||
def random_text(rng: random.Random) -> str:
|
||||
return "".join(rng.choice(ALPHABET) for _ in range(rng.randrange(0, 40)))
|
||||
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:
|
||||
|
|
@ -265,7 +301,7 @@ def chat_body(rng: random.Random, target_tokens: int) -> dict[str, object]:
|
|||
turns: Final = next(index for index, total in enumerate(running) if total >= target_tokens) + 1
|
||||
contents: Final = candidates[: turns + (turns % 2)]
|
||||
return {
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant. Answer precisely and cite sources."},
|
||||
*(
|
||||
|
|
@ -305,9 +341,9 @@ TOOLS: Final = [
|
|||
]
|
||||
|
||||
SMALL_REQUESTS: Final = (
|
||||
{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello, how are you today?"}]},
|
||||
{"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]},
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a terse assistant."},
|
||||
{
|
||||
|
|
@ -322,21 +358,21 @@ SMALL_REQUESTS: Final = (
|
|||
],
|
||||
},
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "weather?"}],
|
||||
"tools": TOOLS,
|
||||
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
|
||||
},
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "system", "content": "sys"}, {"role": "user", "content": "weather?"}],
|
||||
"tools": TOOLS,
|
||||
"tool_choice": "none",
|
||||
},
|
||||
{"model": "gpt-4", "prompt": "Write a haiku about ships."},
|
||||
{"model": "gpt-4", "prompt": ["first prompt", "second prompt"]},
|
||||
{"model": MODEL, "prompt": "Write a haiku about ships."},
|
||||
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"input": [
|
||||
{
|
||||
"role": "user",
|
||||
|
|
@ -348,9 +384,9 @@ SMALL_REQUESTS: Final = (
|
|||
],
|
||||
"instructions": "be terse",
|
||||
},
|
||||
{"model": "gpt-4", "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
|
||||
{"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"model": MODEL,
|
||||
"query": "best harbour",
|
||||
"documents": [
|
||||
"doc one",
|
||||
|
|
@ -362,16 +398,22 @@ SMALL_REQUESTS: Final = (
|
|||
|
||||
def main() -> None:
|
||||
rng: Final = random.Random(2026)
|
||||
texts: Final = tuple(CORPUS) + tuple(random_text(rng) for _ in range(3000))
|
||||
with (HERE / "texts.jsonl").open("w", encoding="utf-8") as handle:
|
||||
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 (HERE / "requests.jsonl").open("w", encoding="utf-8") as handle:
|
||||
with (OUT / "requests.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for body in bodies:
|
||||
input_tokens = _count_input_tokens(dict(body), "gpt-4")
|
||||
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")
|
||||
|
||||
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
|
|
@ -176,21 +176,47 @@ fn loading_a_bad_tokenizer_is_a_load_error() {
|
|||
assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
|
||||
}
|
||||
|
||||
fn cl100k_counter() -> TokenCounter {
|
||||
let path = concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
|
||||
);
|
||||
let ranks = std::fs::read_to_string(path).expect("cl100k rank file is in the repo");
|
||||
TokenCounter::from_cl100k_ranks(&ranks).expect("cl100k ranks 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,
|
||||
}
|
||||
|
||||
fn cl100k_fixture(name: &str) -> String {
|
||||
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!(
|
||||
"{}/tests/fixtures/cl100k/{name}",
|
||||
env!("CARGO_MANIFEST_DIR")
|
||||
"{}/../../../litellm/litellm_core_utils/tokenizers/{}",
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
encoding.rank_file
|
||||
);
|
||||
std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/cl100k/generate.py")
|
||||
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)]
|
||||
|
|
@ -205,12 +231,14 @@ struct RequestFixture {
|
|||
input_tokens: usize,
|
||||
}
|
||||
|
||||
/// Reference counts come from `tiktoken.get_encoding("cl100k_base")`; see
|
||||
/// `tests/fixtures/cl100k/generate.py`.
|
||||
#[test]
|
||||
fn cl100k_text_counts_match_tiktoken() {
|
||||
let counter = cl100k_counter();
|
||||
let fixtures: Vec<TextFixture> = cl100k_fixture("texts.jsonl")
|
||||
/// 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();
|
||||
|
|
@ -229,12 +257,14 @@ fn cl100k_text_counts_match_tiktoken() {
|
|||
}
|
||||
|
||||
/// Reference counts come from the proxy's admission counter
|
||||
/// (`_count_input_tokens(body, "gpt-4")`), so this pins the shared message,
|
||||
/// tool and reply-priming accounting on the cl100k path as well.
|
||||
#[test]
|
||||
fn cl100k_request_counts_match_python_admission_counter() {
|
||||
let counter = cl100k_counter();
|
||||
let fixtures: Vec<RequestFixture> = cl100k_fixture("requests.jsonl")
|
||||
/// (`_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();
|
||||
|
|
@ -243,7 +273,7 @@ fn cl100k_request_counts_match_python_admission_counter() {
|
|||
.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("gpt-4"));
|
||||
assert_eq!(count.model.as_deref(), Some(encoding.model));
|
||||
assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body);
|
||||
count.input_tokens
|
||||
})
|
||||
|
|
@ -251,9 +281,13 @@ fn cl100k_request_counts_match_python_admission_counter() {
|
|||
assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cl100k_shares_the_message_accounting_with_the_anthropic_path() {
|
||||
let counter = cl100k_counter();
|
||||
#[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"))
|
||||
|
|
@ -279,9 +313,9 @@ fn cl100k_shares_the_message_accounting_with_the_anthropic_path() {
|
|||
#[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) {
|
||||
assert!(matches!(
|
||||
TokenCounter::from_cl100k_ranks(rank_file),
|
||||
Err(Error::Ranks(_))
|
||||
));
|
||||
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(_))));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ 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:
|
||||
|
|
@ -23,6 +24,11 @@ def cl100k_base_rank_file() -> str:
|
|||
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).
|
||||
|
|
|
|||
|
|
@ -1365,8 +1365,8 @@ 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 whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base)
|
||||
are counted from the raw body by the bridge when it is enabled, once per
|
||||
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
|
||||
|
|
|
|||
|
|
@ -11,14 +11,14 @@ from pydantic import TypeAdapter
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file
|
||||
from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.configuration import rust_enabled
|
||||
from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt
|
||||
from litellm.utils import claude_json_str, huggingface_tokenizer_kind
|
||||
|
||||
RustTokenizer = Literal["anthropic", "cl100k_base"]
|
||||
RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"]
|
||||
|
||||
|
||||
class RustTokenCounter(Protocol):
|
||||
|
|
@ -33,6 +33,9 @@ class RustTokenCounterFactory(Protocol):
|
|||
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:
|
||||
|
|
@ -60,8 +63,9 @@ 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 only `cl100k_base` does. Rust prices every message with
|
||||
the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python."""
|
||||
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)
|
||||
|
|
@ -69,7 +73,13 @@ def rust_tokenizer(model: str) -> RustTokenizer | None:
|
|||
return "anthropic"
|
||||
if kind is not None or uses_legacy_message_accounting(model):
|
||||
return None
|
||||
return "cl100k_base" if openai_tokenizer_encoding(model).name == "cl100k_base" else 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)
|
||||
|
|
@ -79,6 +89,8 @@ def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> Rust
|
|||
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_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
|
||||
|
|
|
|||
|
|
@ -211,9 +211,12 @@ 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})
|
||||
RUST_INPUT_TOKENS_BY_TOKENIZER: Final = MappingProxyType(
|
||||
{"anthropic": RUST_INPUT_TOKENS, "cl100k_base": 1_234, "o200k_base": 2_345}
|
||||
)
|
||||
|
||||
|
||||
class _FakeDeclined(Exception):
|
||||
|
|
@ -242,7 +245,7 @@ class _RecordingCounter:
|
|||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_cl100k_ranks`."""
|
||||
"""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]] = []
|
||||
|
|
@ -253,6 +256,9 @@ class _RecordingFactory:
|
|||
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:
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
|
|
@ -266,6 +272,9 @@ class _DecliningFactory:
|
|||
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):
|
||||
|
|
@ -323,12 +332,36 @@ async def test_tiktoken_cl100k_models_are_counted_by_rust(rust_counter: None, mo
|
|||
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", "gpt-4o")
|
||||
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(
|
||||
|
|
@ -339,18 +372,20 @@ async def test_multi_model_request_counts_once_per_tokenizer_and_python_for_the_
|
|||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert factory.calls == [("cl100k_base", raw_body), ("anthropic", 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,
|
||||
"gpt-4o": python_counts["gpt-4o"],
|
||||
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["gpt-4o"] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
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))
|
||||
@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())
|
||||
|
|
@ -373,7 +408,7 @@ 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(factory)
|
||||
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL], "messages": ANTHROPIC_MESSAGES}
|
||||
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL], "messages": ANTHROPIC_MESSAGES}
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body,
|
||||
|
|
@ -383,13 +418,18 @@ async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None
|
|||
)
|
||||
|
||||
assert factory.calls == []
|
||||
assert set(counts) == {ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL}
|
||||
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
|
||||
@pytest.mark.parametrize("model", ("gpt-4o", "o3", "replicate/meta/llama-2-70b-chat"))
|
||||
async def test_models_without_a_rust_tokenizer_stay_in_python(rust_counter: None, model: str) -> 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(factory)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ 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
|
||||
|
|
@ -16,6 +17,7 @@ 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
|
||||
|
|
@ -23,6 +25,9 @@ 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()
|
||||
|
||||
|
||||
|
|
@ -50,7 +55,7 @@ class _RecordingCounter:
|
|||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_cl100k_ranks` for ranks."""
|
||||
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counters: list[_RecordingCounter] = []
|
||||
|
|
@ -65,6 +70,10 @@ class _RecordingFactory:
|
|||
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:
|
||||
|
|
@ -86,6 +95,9 @@ class _RaisingFactory:
|
|||
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):
|
||||
|
|
@ -100,7 +112,7 @@ def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.RustTokenizer) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
|
|
@ -127,18 +139,20 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cl100k_counter_is_built_from_the_vendored_rank_file_once() -> None:
|
||||
@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, "cl100k_base")
|
||||
second: Final = await bridge.count_input_tokens(BODY, "cl100k_base")
|
||||
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") == 100_256
|
||||
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]
|
||||
|
||||
|
||||
|
|
@ -150,10 +164,12 @@ async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
|
|||
|
||||
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 == "cl100k_base" for counter in factory.counters] == [False, True]
|
||||
assert [len(counter.bodies) for counter in factory.counters] == [2, 1]
|
||||
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
|
||||
|
|
@ -161,12 +177,11 @@ 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_input_tokens(BODY, "anthropic") is None
|
||||
assert await bridge.count_input_tokens(BODY, "cl100k_base") is None
|
||||
assert [await bridge.count_input_tokens(BODY, tokenizer) for tokenizer in TOKENIZERS] == [None, None, None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages")))
|
||||
|
|
@ -175,7 +190,7 @@ async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> N
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed")))
|
||||
|
|
@ -197,11 +212,17 @@ async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> No
|
|||
("my-router-alias", "cl100k_base"),
|
||||
("azure/gpt-4o", "cl100k_base"),
|
||||
("command-r-plus", "cl100k_base"),
|
||||
("gpt-4o", None),
|
||||
("gpt-4o-mini", None),
|
||||
("gpt-4.1", None),
|
||||
("gpt-5", None),
|
||||
("o3", None),
|
||||
("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),
|
||||
),
|
||||
|
|
@ -210,6 +231,19 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected:
|
|||
assert bridge.rust_tokenizer(model) == expected
|
||||
|
||||
|
||||
@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 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"})
|
||||
|
||||
|
|
@ -233,18 +267,25 @@ def test_rust_tokenizer_declines_legacy_message_accounting_python_prices_differe
|
|||
assert bridge.rust_tokenizer(CL100K_MODEL) == "cl100k_base"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", (MODEL, "gpt-4", "gpt-4o"))
|
||||
@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! \u00e9\u00e8 12345 \u3053\u3093\u306b\u3061\u306f <|endoftext|>\r\n" * 10
|
||||
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) != cl100k_count
|
||||
assert python_count == len(Tokenizer.from_str(claude_json_str).encode(text).ids)
|
||||
assert python_count not in {cl100k_count, o200k_count}
|
||||
case None:
|
||||
assert python_count != cl100k_count
|
||||
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:
|
||||
|
|
@ -252,7 +293,7 @@ def test_disabled_hf_download_routes_anthropic_models_to_cl100k_like_python(monk
|
|||
|
||||
assert bridge.rust_tokenizer(MODEL) == "cl100k_base"
|
||||
assert bridge.rust_tokenizer("meta-llama/Llama-3-8b") == "cl100k_base"
|
||||
assert bridge.rust_tokenizer("gpt-4o") is None
|
||||
assert bridge.rust_tokenizer(O200K_MODEL) == "o200k_base"
|
||||
|
||||
|
||||
def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -260,6 +301,7 @@ def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyP
|
|||
|
||||
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], ...]] = (
|
||||
|
|
@ -328,6 +370,8 @@ 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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -351,19 +395,22 @@ async def test_native_count_matches_python_budget_counter(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cl100k_counts_long_text_exactly_where_python_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@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": CL100K_MODEL, "messages": [{"role": "user", "content": text}]}
|
||||
encoding: Final = tiktoken.get_encoding("cl100k_base")
|
||||
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(), "cl100k_base")
|
||||
python_count: Final = _count_input_tokens(request_body=body, model=CL100K_MODEL)
|
||||
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
|
||||
|
|
@ -385,7 +432,7 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
|
||||
@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], tokenizer: bridge.RustTokenizer
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue