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:
Zhenjia ZHOU
2026-08-12 12:44:31 +08:00
committed by GitHub
parent 2483f57002
commit 6840153473
+56 -12
View File
@@ -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();