fix(tokenizer): price CJK in the Rust fixed-ratio estimator (Python parity) (#2260)
## Description The Rust `EstimatingCounter` priced every character at the Latin `chars_per_token` (default 4.0), but the Python `EstimatingTokenCounter` it explicitly mirrors already prices dense scripts (CJK / Kana / Hangul / full-width) at `CHARS_PER_TOKEN_CJK = 1.5` — so Rust under-counted CJK by ~2.5× and the two implementations diverged. #2080 fixed only the Python path; the Rust module doc still says "Mirrors …EstimatingTokenCounter" while it no longer did. This is the live count path for every provider-calibrated fixed-ratio counter (Anthropic 3.5, Google / Cohere 4.0, Moonshot 3.1), so CJK traffic was mis-budgeted (savings/estimates skewed). This counts dense-script codepoints — the same 8 `CJK_PATTERN` Unicode ranges Python uses — and prices them separately: `int(other / ratio + cjk / 1.5 + 0.5)`. Non-CJK output is byte-identical. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `crates/headroom-core/src/tokenizer/estimator.rs`: add `is_dense_script(c)` (8 ranges byte-mirroring Python `CJK_PATTERN`) and `CHARS_PER_TOKEN_CJK = 1.5`; `count_text` prices dense-script chars separately from Latin. - Reference tests for CJK / kana / full-width / mixed — values cross-checked against Python. ## Testing - [x] Unit tests pass (`cargo test`) - [x] Linting passes (`cargo clippy` / `cargo fmt`) - [x] New tests added - [x] Verified against Python (see Real Behavior Proof) ### Test Output ```text $ cargo test -p headroom-core --lib tokenizer test result: ok. 45 passed; 0 failed $ cargo clippy / fmt # clean ``` ## Real Behavior Proof - Environment: macOS (Darwin), Rust via cargo + Python in a uv venv, branch `feat/tokenizer-estimator-cjk` off `main`. - Exact command / steps: ran the same inputs through Python `EstimatingTokenCounter(4.0).count_text` and the Rust `EstimatingCounter::default().count_text`, comparing outputs. - Observed result: identical on every input — `数据库` → 2, `数据库连接失败` → 5 (was 2 under the old flat 7/4), `ひらが` → 2, full-width `API` → 2 vs plain `API` → 1, mixed `api数据` → 2. The existing non-CJK reference tests (`a`×40 → 10, Claude-3.5 densities, `héllo`/emoji char-count) are unchanged, confirming no ASCII regression. - Not tested: nothing further — parity is verified directly against the Python reference (same values on both sides). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation — N/A (internal estimator) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md — happy to add an entry if preferred. ## Additional Notes - Completes #2080 (which priced CJK in the Python fixed-ratio estimator) on the Rust side, restoring Rust↔Python parity for the density estimator. Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
@@ -1,12 +1,35 @@
|
||||
//! Character-density estimator. Used as a fallback for any tokenizer family
|
||||
//! we haven't wired in yet (Anthropic Claude, Google Gemini, Cohere, …).
|
||||
//!
|
||||
//! Mirrors `headroom.tokenizers.estimator.EstimatingTokenCounter`. The formula
|
||||
//! is `ceil(chars / chars_per_token)`. `chars` is *Unicode scalar count*, not
|
||||
//! byte length, to match Python's `len(text)` semantics on str.
|
||||
//! Mirrors `headroom.tokenizers.estimator.EstimatingTokenCounter`. Latin chars
|
||||
//! are priced at `chars_per_token`; dense scripts (CJK / Kana / Hangul / full-
|
||||
//! width) are priced separately at `CHARS_PER_TOKEN_CJK`, since they tokenize at
|
||||
//! ~1 token/char and the Latin ratio under-counts them 2-4x. `chars` is a
|
||||
//! *Unicode scalar count*, not byte length, to match Python's `len(text)`.
|
||||
|
||||
use super::{Backend, Tokenizer};
|
||||
|
||||
/// Chars-per-token for dense scripts. Byte-identical with Python
|
||||
/// `EstimatingTokenCounter.CHARS_PER_TOKEN_CJK`.
|
||||
const CHARS_PER_TOKEN_CJK: f64 = 1.5;
|
||||
|
||||
/// True for a "dense-script" codepoint (CJK ideographs + punctuation, Kana,
|
||||
/// Hangul, CJK compatibility, half/full-width forms, CJK Ext-A/B). Ranges kept
|
||||
/// byte-identical with Python `EstimatingTokenCounter.CJK_PATTERN`.
|
||||
fn is_dense_script(c: char) -> bool {
|
||||
matches!(
|
||||
c as u32,
|
||||
0x3000..=0x303F // CJK symbols and punctuation
|
||||
| 0x3040..=0x30FF // Hiragana + Katakana
|
||||
| 0x3400..=0x4DBF // CJK Unified Ideographs Ext A
|
||||
| 0x4E00..=0x9FFF // CJK Unified Ideographs
|
||||
| 0xAC00..=0xD7AF // Hangul syllables
|
||||
| 0xF900..=0xFAFF // CJK compatibility ideographs
|
||||
| 0xFF00..=0xFFEF // Half/full-width forms
|
||||
| 0x20000..=0x2A6DF // CJK Unified Ideographs Ext B
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct EstimatingCounter {
|
||||
chars_per_token: f64,
|
||||
@@ -42,15 +65,15 @@ impl Tokenizer for EstimatingCounter {
|
||||
if text.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
// Match Python `EstimatingTokenCounter.count_text`:
|
||||
// max(1, int(len(text) / chars_per_token + 0.5))
|
||||
// Python `int()` truncates toward zero; for non-negative inputs that's
|
||||
// identical to `as usize` saturating-cast semantics in Rust >= 1.45.
|
||||
// Adding 0.5 then truncating yields round-half-up. We previously used
|
||||
// ceil, which over-counted in the middle of the range (e.g. "aaaaa"
|
||||
// at 4.0 cpt returned 2 here vs 1 in Python).
|
||||
let chars = text.chars().count() as f64;
|
||||
let raw = (chars / self.chars_per_token + 0.5) as usize;
|
||||
// Match Python `EstimatingTokenCounter.count_text` (fixed-ratio path):
|
||||
// cjk = count_dense_script(text); other = len(text) - cjk
|
||||
// max(1, int(other / chars_per_token + cjk / CHARS_PER_TOKEN_CJK + 0.5))
|
||||
// Dense scripts tokenize at ~1 token/char, so the Latin `chars_per_token`
|
||||
// under-counts them; price them separately. `int()` truncates toward
|
||||
// zero (== `as usize` for non-negative); the `+ 0.5` gives round-half-up.
|
||||
let cjk = text.chars().filter(|&c| is_dense_script(c)).count();
|
||||
let other = (text.chars().count() - cjk) as f64;
|
||||
let raw = (other / self.chars_per_token + cjk as f64 / CHARS_PER_TOKEN_CJK + 0.5) as usize;
|
||||
raw.max(1)
|
||||
}
|
||||
|
||||
@@ -104,6 +127,27 @@ mod tests {
|
||||
assert_eq!(est.count_text("🦀🦀🦀🦀"), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dense_scripts_priced_at_cjk_ratio() {
|
||||
let est = EstimatingCounter::default(); // 4.0 for Latin
|
||||
// Pure CJK: cjk=3, other=0 -> 0/4 + 3/1.5 + 0.5 = 2.5 -> int -> 2
|
||||
assert_eq!(est.count_text("数据库"), 2);
|
||||
// 7 CJK -> 7/1.5 + 0.5 = 5.16 -> 5 (the old flat 7/4 -> 2 under-counted ~2.5x)
|
||||
assert_eq!(est.count_text("数据库连接失败"), 5);
|
||||
// Kana is dense: 3 hiragana -> 3/1.5 + 0.5 = 2.5 -> 2
|
||||
assert_eq!(est.count_text("ひらが"), 2);
|
||||
// Full-width Latin is dense (U+FF00-FFEF): API -> 2, vs plain "API" -> 1
|
||||
assert_eq!(est.count_text("API"), 2);
|
||||
assert_eq!(est.count_text("API"), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_ascii_and_cjk_prices_each_separately() {
|
||||
let est = EstimatingCounter::default();
|
||||
// "api数据": other=3, cjk=2 -> 3/4 + 2/1.5 + 0.5 = 0.75+1.33+0.5 = 2.58 -> 2
|
||||
assert_eq!(est.count_text("api数据"), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn min_is_one_for_non_empty_input() {
|
||||
let est = EstimatingCounter::default();
|
||||
|
||||
Reference in New Issue
Block a user