From 6840153473caa0d61e982215e16a8cf54b0b6cc7 Mon Sep 17 00:00:00 2001 From: Zhenjia ZHOU Date: Wed, 12 Aug 2026 12:44:31 +0800 Subject: [PATCH] fix(tokenizer): price CJK in the Rust fixed-ratio estimator (Python parity) (#2260) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 `ＡＰＩ` → 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 --- .../headroom-core/src/tokenizer/estimator.rs | 68 +++++++++++++++---- 1 file changed, 56 insertions(+), 12 deletions(-) diff --git a/crates/headroom-core/src/tokenizer/estimator.rs b/crates/headroom-core/src/tokenizer/estimator.rs index 5dde222d3..818077be0 100644 --- a/crates/headroom-core/src/tokenizer/estimator.rs +++ b/crates/headroom-core/src/tokenizer/estimator.rs @@ -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): ＡＰＩ -> 2, vs plain "API" -> 1 + assert_eq!(est.count_text("ＡＰＩ"), 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();