From 3e86227a144957cc0fd01165b7fe2aff9251587f Mon Sep 17 00:00:00 2001 From: Jack Date: Mon, 3 Aug 2026 17:48:39 +0800 Subject: [PATCH] fix(chunker): align splitOversizedUnitWith with Python running-sum flush (#17729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Align Go `splitOversizedUnitWith` with Python `rag/nlp._split_oversized_unit` so the whitespace-atom sub-split produces byte-identical chunk boundaries. ### Root cause of the divergence cl100k token counting is **not additive across whitespace joins** (`token(a)+token(b) != token(a+b)`). Go previously used the exact joined-string fit check `countFn(current+atom) > budget`, while Python accumulates a running sum `current_tokens + a_tokens > budget`. The two formulas disagree by one atom at the boundary, so Go and Python emitted the same chunk *count* but shifted *text*. ### Changes - `splitOversizedUnitWith` (`token.go`): replace the exact joined-string fit check with the running-sum check (mirroring Python's `current_tokens` accumulator), and after a flush keep the overflow whitespace atom (`current += atom`) instead of dropping it. - `token_strict_cap_test.go`: relax `TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit` to allow the same cl100k non-additive +1 overshoot Python exhibits (the invariant — an oversized unit is sub-split, not collapsed — is preserved). ### Test plan `bash build.sh --test ./internal/ingestion/component/chunker/...` — green. ## Note Test infrastructure for this change (golden parity harness, `split_oversized_test.go`, `testdata/parity/**`, `known_diffs.json`, `capture_golden.py`/`live_chunk.py`, and the `go-cmp` dependency promotion) is split into a separate, stacked PR #17735 so this PR stays minimal (production code only). This PR is **independent of #17712** (the offline BPE loader). It is based on `upstream/main` and contains only this change; no BPE-loader code is included. Co-authored-by: CodeBuddy --- internal/ingestion/component/chunker/token.go | 31 +++++++++---------- .../chunker/token_strict_cap_test.go | 11 +++++-- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 7d86b611be..0bd85079ed 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -397,14 +397,19 @@ func splitOversizedUnitWith(text string, chunkTokenNum int, countFn func(string) } var pieces []string current := "" + // Running sum of per-atom token counts for the current piece. Mirrors + // Python rag/nlp._split_oversized_unit's `current_tokens`. We flush when + // this running sum (not the exact count of the joined string) would + // exceed the budget, because cl100k token counting is not additive across + // whitespace joins: token(a)+token(b) can differ from token(a+b), so the + // joined-string fit check drifts one atom off Python's boundary. + currentTokens := 0 tokenCache := map[string]int{} atomTokens := func(atom string) int { // Whitespace-only atoms contribute 0 in isolation (mirrors Python // atom.isspace()), matching the packing heuristic used by - // rag/nlp._split_oversized_unit. Fit checks below still use an - // exact projected countFn(current+atom) so cl100k space-join - // effects cannot push a piece over the hard cap. + // rag/nlp._split_oversized_unit. if strings.TrimSpace(atom) == "" { return 0 } @@ -422,28 +427,20 @@ func splitOversizedUnitWith(text string, chunkTokenNum int, countFn func(string) if current != "" { pieces = append(pieces, current) current = "" + currentTokens = 0 } pieces = append(pieces, splitAtomByTokenBudget(atom, chunkTokenNum, countFn)...) continue } - // Exact projected-total check (not sum of atom counts): cl100k can - // count a joined "word word" differently than token(word)+token(word). - if current != "" && countFn(current+atom) > chunkTokenNum { + // Running-sum fit check, identical to Python's + // `current_tokens + a_tokens > chunk_token_num`. + if current != "" && currentTokens+aTokens > chunkTokenNum { pieces = append(pieces, current) current = "" - // Leading whitespace after a flush has no content value; drop it - // so the next piece does not start with a pure-space prefix that - // would never fit usefully on its own. - if strings.TrimSpace(atom) == "" { - continue - } - // If the atom alone still exceeds (pathological), carve it. - if atomTokens(atom) > chunkTokenNum { - pieces = append(pieces, splitAtomByTokenBudget(atom, chunkTokenNum, countFn)...) - continue - } + currentTokens = 0 } current += atom + currentTokens += aTokens } if current != "" { pieces = append(pieces, current) diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index c00774a552..5f70881d16 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -150,9 +150,16 @@ func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { if len(got[0]) < 2 { t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0])) } + // cl100k is not additive across whitespace joins: token(a)+token(b) can be + // one less than token(a+b), so the running-sum flush used by both Python's + // _split_oversized_unit and the aligned Go port can leave a piece exactly + // one token over the nominal budget. The invariant we defend here is that + // the oversized unit is sub-split (not collapsed into one chunk), not a + // byte-exact cap — matching the Python reference. + const slack = 1 for i, ck := range got[0] { - if n := tokenizeStr(ck.Text); n > budget { - t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + if n := tokenizeStr(ck.Text); n > budget+slack { + t.Errorf("chunk %d exceeds budget by more than cl100k slack: tokens=%d (cap=%d)", i, n, budget) } } }