refactor(chunker): converge delimiter_mode to {delimiter, one}, drop token_size (#17979)

Converge `TokenChunker.delimiter_mode` from three values (`token_size`,
`delimiter`, `one`) to two (`delimiter`, `one`). The unified `delimiter`
mode now carries the old `token_size` semantics: when no active
(backtick) delimiter is present, text/JSON chunks are merged up to
`chunk_token_size`; when a backtick delimiter is present, the text is
split by it and not merged. `one` continues to be handled by the
separate `OneChunker`.
This commit is contained in:
Jack
2026-08-07 16:11:42 +08:00
committed by GitHub
parent 4044d3bc5d
commit 1aa4e3c1f3
15 changed files with 167 additions and 83 deletions

View File

@@ -336,7 +336,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [],
"image_context_size": 0,
"outputs": {
@@ -576,7 +576,7 @@
"form": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
{
"value": "\n"

View File

@@ -219,7 +219,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [],
"image_context_size": 0,
"outputs": {
@@ -429,7 +429,7 @@
"data": {
"form": {
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
{
"value": "\n"

View File

@@ -192,7 +192,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [],
"image_context_size": 0,
"outputs": {
@@ -374,7 +374,7 @@
"form": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
{
"value": "\n"

View File

@@ -16,7 +16,7 @@
// SCOPE (honest) for token.go:
//
// - WHITELIST: delimiter_mode ∈ {"token_size","delimiter"} (the
// - WHITELIST: delimiter_mode ∈ {"delimiter"} (the
// single-chunk "one" behaviour moved to OneChunker in one.go).
// chunk_token_size > 0, overlapped_percent accepts a [0,1) fraction or a
// [0,90] percentage (normalized to [0,90] by normalizeOverlappedPercent,

View File

@@ -125,7 +125,7 @@ func TestMergeByTokenSize_TextPathStrictCap(t *testing.T) {
b.WriteString("\n\n")
}
comp, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": budget,
})
if err != nil {
@@ -159,7 +159,7 @@ func TestMergeByTokenSize_OversizedUnitStaysWhole(t *testing.T) {
// run is one unit that still exceeds the budget and must stay whole.
long := strings.TrimSpace(strings.Repeat("word ", 100))
comp, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": budget,
})
if err != nil {
@@ -196,7 +196,7 @@ func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) {
run := func(underCap bool) []map[string]any {
comp, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": budget,
"under_cap": underCap,
})
@@ -250,7 +250,7 @@ func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) {
b.WriteByte('\n')
}
comp, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": budget,
})
if err != nil {

View File

@@ -25,7 +25,7 @@ import (
// the body text still carries coordinate markers.
func TestTokenChunker_TextPath_StripsParserTags(t *testing.T) {
c, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": 1000,
"delimiters": []string{"\n"},
})

View File

@@ -145,7 +145,7 @@ func TestTokenChunker_DelimNeverStandaloneChunk(t *testing.T) {
// token-size merge and emit >=1 chunk.
func TestTokenChunker_InvokeTokenSize_FallbackToMerge(t *testing.T) {
c, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": 50,
"delimiters": []string{"`\n\n`"},
})
@@ -365,9 +365,9 @@ func TestTokenChunker_NewRejectsBadParam(t *testing.T) {
}{
{"bad delimiter_mode", map[string]any{"delimiter_mode": "nope"}},
{"one delimiter_mode (use OneChunker)", map[string]any{"delimiter_mode": "one"}},
{"zero chunk_token_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 0}},
{"negative chunk_token_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": -5}},
{"negative table_context_size", map[string]any{"delimiter_mode": "token_size", "chunk_token_size": 50, "table_context_size": -1}},
{"zero chunk_token_size", map[string]any{"delimiter_mode": "delimiter", "chunk_token_size": 0}},
{"negative chunk_token_size", map[string]any{"delimiter_mode": "delimiter", "chunk_token_size": -5}},
{"negative table_context_size", map[string]any{"delimiter_mode": "delimiter", "chunk_token_size": 50, "table_context_size": -1}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
@@ -380,14 +380,14 @@ func TestTokenChunker_NewRejectsBadParam(t *testing.T) {
// TestTokenChunker_NewAcceptsDefaults ensures the no-config
// constructor returns a usable component with a working default
// delimiter_mode = "token_size".
// delimiter_mode = "delimiter".
func TestTokenChunker_NewAcceptsDefaults(t *testing.T) {
c, err := NewTokenChunker(nil)
if err != nil {
t.Fatalf("NewTokenChunker(nil): %v", err)
}
if got := c.(*TokenChunkerComponent).param.DelimiterMode; got != "token_size" {
t.Errorf("default delimiter_mode = %q, want token_size", got)
if got := c.(*TokenChunkerComponent).param.DelimiterMode; got != "delimiter" {
t.Errorf("default delimiter_mode = %q, want delimiter", got)
}
}
@@ -439,7 +439,7 @@ func TestTokenChunker_NewAcceptsPythonOverlappedRange(t *testing.T) {
// percentages, including out-of-range inputs that Python clamps).
for _, pct := range []float64{0, 0.1, 0.5, 15, 30, 50, 90, 95, -5} {
conf := map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": 100,
"overlapped_percent": pct,
}
@@ -535,7 +535,7 @@ func TestTokenChunker_NormalizesOverlappedPercent(t *testing.T) {
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
c, err := NewTokenChunker(map[string]any{
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"chunk_token_size": 100,
"overlapped_percent": tc.in,
})
@@ -575,7 +575,7 @@ func TestTokenChunkerParam_ValidateOverlappedRange(t *testing.T) {
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
p := schema.TokenChunkerParam{
DelimiterMode: "token_size",
DelimiterMode: "delimiter",
ChunkTokenSize: 100,
OverlappedPercent: tc.in,
}

View File

@@ -173,7 +173,7 @@ type ChunkerOutputs struct {
type TokenChunkerParam struct {
// DelimiterMode selects the chunking strategy.
// Allowed values: "token_size", "delimiter".
// Allowed value: "delimiter".
// The single-chunk "one" behavior is provided by the separate
// OneChunker component.
DelimiterMode string `json:"delimiter_mode"`
@@ -246,7 +246,7 @@ func (p TokenChunkerParam) MergeStrategy() MergeStrategy {
// Defaults returns the Python default TokenChunkerParam.
func (TokenChunkerParam) Defaults() TokenChunkerParam {
return TokenChunkerParam{
DelimiterMode: "token_size",
DelimiterMode: "delimiter",
ChunkTokenSize: 512,
Delimiters: []string{"\n"},
OverlappedPercent: 0,
@@ -357,7 +357,7 @@ func (p *TokenChunkerParam) Validate() error {
p.OverlappedPercent = math.Round(v * 100)
}
switch p.DelimiterMode {
case "token_size", "delimiter":
case "delimiter":
default:
return errInvalidValue{Field: "delimiter_mode", Value: p.DelimiterMode}
}

View File

@@ -278,8 +278,8 @@ func TestChunkerOutputsJSONRoundTrip(t *testing.T) {
func TestTokenChunkerParamDefaults(t *testing.T) {
p := TokenChunkerParam{}.Defaults()
if p.DelimiterMode != "token_size" {
t.Errorf("default delimiter_mode = %q, want token_size", p.DelimiterMode)
if p.DelimiterMode != "delimiter" {
t.Errorf("default delimiter_mode = %q, want delimiter", p.DelimiterMode)
}
if p.ChunkTokenSize != 512 {
t.Errorf("default chunk_token_size = %d, want 512", p.ChunkTokenSize)

View File

@@ -83,7 +83,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
"\n",
"!",
@@ -265,7 +265,7 @@
"form": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
{
"value": "\n"

View File

@@ -161,7 +161,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
"\n",
"!",
@@ -391,7 +391,7 @@
"form": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
{
"value": "\n"

View File

@@ -157,7 +157,7 @@
"params": {
"children_delimiters": [],
"chunk_token_size": 512,
"delimiter_mode": "token_size",
"delimiter_mode": "delimiter",
"delimiters": [
"\n",
"!",

File diff suppressed because one or more lines are too long

View File

@@ -32,7 +32,7 @@ from rag.nlp import naive_merge
class TokenChunkerParam(ProcessParamBase):
def __init__(self):
super().__init__()
self.delimiter_mode = "token_size"
self.delimiter_mode = "delimiter"
self.chunk_token_size = 512
self.delimiters = ["\n"]
self.overlapped_percent = 0
@@ -41,7 +41,13 @@ class TokenChunkerParam(ProcessParamBase):
self.image_context_size = 0
def check(self):
self.check_valid_value(self.delimiter_mode, "Delimiter mode abnormal.", ["token_size", "delimiter", "one"])
# Backward-compat: "token_size" was removed but is behaviorally identical
# to "delimiter" at runtime (both route through the same code path), so
# accept and coerce it instead of rejecting legacy configs / pre-fix
# frontends. Only genuinely unknown values are rejected.
if self.delimiter_mode == "token_size":
self.delimiter_mode = "delimiter"
self.check_valid_value(self.delimiter_mode, "Delimiter mode abnormal.", ["delimiter", "one"])
if self.delimiters is None:
self.delimiters = []
elif isinstance(self.delimiters, str):
@@ -316,9 +322,7 @@ class TokenChunker(ProcessBase):
self.set_output("chunks", [{"text": payload}] if payload.strip() else [])
self.callback(1, "Done.")
return
if self._param.delimiter_mode == "delimiter":
cks = _split_text_by_pattern(payload, delimiter_pattern)
elif delimiter_pattern:
if delimiter_pattern:
cks = _split_text_by_pattern(payload, delimiter_pattern)
else:
cks = naive_merge(
@@ -358,8 +362,14 @@ class TokenChunker(ProcessBase):
self.callback(1, "Done.")
return
if self._param.delimiter_mode == "delimiter":
text_chunks = _build_json_chunks(json_result, "")
# Both branches start from per-item chunks (no pre-split by the
# delimiter pattern). The delimiter branch splits the buffered text
# stream while preserving per-segment PDF positions; the no-delimiter
# branch merges adjacent text items to chunk_token_size (the removed
# "token_size" behaviour, and a parity match with the Go JSON path).
text_chunks = _build_json_chunks(json_result, "")
if delimiter_pattern:
chunks = []
text_buffer = []
text_buffer_pos = []
@@ -372,9 +382,9 @@ class TokenChunker(ProcessBase):
# The delimiter is then applied to the combined text; a segment may
# span across item boundaries (the "\n" glue is not itself a
# delimiter), so each segment carries only the PDF positions of the
# buffered item(s) whose text contributed to it -- never the union of
# every item (which previously leaked page-N coordinates into
# page-M chunks and made all segments share one preview image).
# item(s) that contributed to it -- never the union of every item
# (which previously leaked page-N coordinates into page-M chunks and
# made all segments share one preview image).
parts = []
item_ranges = [] # (start, end) of each buffered item in combined_text
offset = 0
@@ -387,27 +397,24 @@ class TokenChunker(ProcessBase):
offset += 1
combined_text = "".join(parts[:-1]) # drop the trailing glue
if delimiter_pattern:
raw = re.split(r"(%s)" % delimiter_pattern, combined_text, flags=re.DOTALL)
segments = [] # (text, start, end) within combined_text
pos = 0
for i in range(0, len(raw), 2):
seg = raw[i]
seg_start = pos
seg_end = pos + len(seg)
if seg:
segments.append((seg, seg_start, seg_end))
pos = seg_end
if i + 1 < len(raw):
pos += len(raw[i + 1])
else:
segments = [(combined_text, 0, len(combined_text))]
raw = re.split(r"(%s)" % delimiter_pattern, combined_text, flags=re.DOTALL)
segments = [] # (text, start, end) within combined_text
pos = 0
for i in range(0, len(raw), 2):
seg = raw[i]
seg_start = pos
seg_end = pos + len(seg)
if seg:
segments.append((seg, seg_start, seg_end))
pos = seg_end
if i + 1 < len(raw):
pos += len(raw[i + 1])
for text, seg_start, seg_end in segments:
if not text.strip():
continue
seg_pos = []
for (istart, iend), item_pos in zip(item_ranges, text_buffer_pos):
for (istart, iend), item_pos in zip(item_ranges, text_buffer_pos, strict=True):
# A segment overlaps an item when their character ranges
# intersect; collect that item's coordinates.
if seg_start < iend and istart < seg_end:
@@ -436,22 +443,19 @@ class TokenChunker(ProcessBase):
if custom_pattern:
chunks = _split_chunk_docs_by_children(chunks, custom_pattern)
_attach_context_to_media_chunks(chunks, self._param.table_context_size, self._param.image_context_size)
await restore_pdf_text_previews(chunks, from_upstream, self._canvas)
self.set_output("chunks", _finalize_json_chunks(chunks))
self.callback(1, "Done.")
return
# Structured JSON input is normalized first, then optionally enriched with
# media context, and finally merged only when delimiter splitting is inactive.
chunks = _build_json_chunks(json_result, delimiter_pattern)
_attach_context_to_media_chunks(chunks, self._param.table_context_size, self._param.image_context_size)
if self._param.delimiter_mode == "token_size" and not delimiter_pattern:
chunks = _merge_text_chunks_by_token_size(chunks, self._param.chunk_token_size, overlapped_percent)
if custom_pattern:
chunks = _split_chunk_docs_by_children(chunks, custom_pattern)
else:
# No active delimiter: merge adjacent text items to chunk_token_size.
# This runs on the per-item chunks (NOT a single concatenated chunk),
# so the token cap is actually enforced -- matching the previous
# "token_size" mode and the Go JSON path. Media chunks break the merge.
# Media context is attached on the per-item chunks before merging, as
# the removed "token_size" branch did, to preserve context windows.
_attach_context_to_media_chunks(text_chunks, self._param.table_context_size, self._param.image_context_size)
chunks = _merge_text_chunks_by_token_size(text_chunks, self._param.chunk_token_size, overlapped_percent)
if custom_pattern:
chunks = _split_chunk_docs_by_children(chunks, custom_pattern)
await restore_pdf_text_previews(chunks, from_upstream, self._canvas)
cks = _finalize_json_chunks(chunks)
self.set_output("chunks", cks)
self.set_output("chunks", _finalize_json_chunks(chunks))
self.callback(1, "Done.")
return

View File

@@ -55,6 +55,19 @@ def _load_token_chunker_with_stubs():
def __init__(self):
pass
def check_valid_value(self, value, msg, allowed):
if value not in allowed:
raise ValueError(msg)
def check_positive_integer(self, value, msg):
pass
def check_decimal_float(self, value, msg):
pass
def check_nonnegative_number(self, value, msg):
pass
class ProcessBase:
def __init__(self, _pipeline, _id, param):
self._pipeline = _pipeline
@@ -277,6 +290,73 @@ def test_json_delimiter_mode_pdf_positions_retained():
assert chunks[0].get("pdf_positions") == [[1, 0, 10, 0, 5], [2, 0, 20, 0, 8]]
def test_token_size_mode_normalized_to_delimiter():
# Backward-compat: the removed "token_size" value must still be accepted by
# check() and coerced to "delimiter" (runtime behavior is identical), so
# legacy configs / pre-fix frontends don't get rejected. Unknown values are
# still rejected.
with _load_token_chunker_with_stubs() as token_chunker_module:
param = token_chunker_module.TokenChunkerParam()
param.delimiter_mode = "token_size"
param.check()
assert param.delimiter_mode == "delimiter"
bad = token_chunker_module.TokenChunkerParam()
bad.delimiter_mode = "nope"
try:
bad.check()
raise AssertionError("expected check() to reject unknown delimiter_mode")
except Exception:
pass
def test_json_no_delimiter_mode_merges_to_token_cap():
# Regression for #17979: with no active (backtick) delimiter, the JSON path
# must merge per-item text chunks up to chunk_token_size -- mirroring the old
# "token_size" mode and the Go JSON path. Concatenating every item into a
# single chunk before the merge would defeat the cap and emit one oversized
# chunk. The per-token stub (num_tokens_from_string -> 1) makes the cap easy
# to exceed: 12 one-token items under a cap of 5 must yield several chunks.
for _module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "chunk_token_size": 5}):
kwargs = {
"name": "token_chunker",
"output_format": "json",
"json_result": [{"text": f"item{i}"} for i in range(12)],
}
asyncio.run(chunker._invoke(**kwargs))
chunks = chunker._outputs["chunks"]
# 12 one-token items under a cap of 5 must NOT collapse into one chunk.
assert len(chunks) > 1, f"cap not enforced: 12 items collapsed to {len(chunks)} chunk(s)"
# Each merged chunk holds at most one overflow unit past the cap (<= 6
# items); the final output drops tk_nums, so count item markers instead.
for c in chunks:
assert c["text"].count("item") <= 6, f"chunk exceeds cap: {c['text'].count('item')} items"
# No text lost: all 12 items must survive, joined by the "\n" glue.
joined = "\n".join(c["text"] for c in chunks)
for i in range(12):
assert f"item{i}" in joined, f"item{i} dropped from output"
def test_json_no_delimiter_mode_media_breaks_merge():
# A non-text (media) chunk interleaved between text items must stay as its
# own chunk and reset the merge, so text before/after it are sized separately.
for _module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "chunk_token_size": 5}):
kwargs = {
"name": "token_chunker",
"output_format": "json",
"json_result": [{"text": f"t{i}", "doc_type_kwd": "text"} for i in range(6)]
+ [{"text": "IMG", "doc_type_kwd": "image", "img_id": "im1"}]
+ [{"text": f"u{i}", "doc_type_kwd": "text"} for i in range(12)],
}
asyncio.run(chunker._invoke(**kwargs))
chunks = chunker._outputs["chunks"]
doc_types = [c["doc_type_kwd"] for c in chunks]
# The media chunk is preserved and breaks the text merge.
assert doc_types.count("image") == 1, doc_types
# Several text chunks on each side of the media boundary.
assert doc_types.count("text") > 2, doc_types
def test_json_delimiter_mode_pdf_positions_per_segment_not_broadcast():
# Regression for #3 (PDF coordinate leak): when consecutive text items from
# different pages are buffered and then split by a custom delimiter, each
@@ -346,11 +426,11 @@ def test_json_delimiter_mode_consecutive_delimiter_keeps_boundary():
assert all("##" not in t for t in texts)
def test_text_delimiter_mode_token_size_zero_or_one_no_atom_split():
# token_size=0/1 must not atom-split delimiter segments into 1-token chunks;
# the delimiter path produces delimiter-boundary chunks regardless of cap.
for module, chunker in _build_json_chunker({"delimiter_mode": "token_size", "delimiters": ["`|`"]}):
# text path: delimiter_mode is token_size but a custom delimiter is
def test_text_delimiter_mode_zero_or_one_no_atom_split():
# chunk_token_size=0/1 must not atom-split delimiter segments into 1-token
# chunks; the delimiter path produces delimiter-boundary chunks regardless of cap.
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`|`"]}):
# text path: delimiter_mode is delimiter but a custom delimiter is
# present, so the delimiter branch (_split_text_by_pattern) is used.
kwargs = {
"name": "token_chunker",
@@ -362,4 +442,4 @@ def test_text_delimiter_mode_token_size_zero_or_one_no_atom_split():
asyncio.run(chunker._invoke(**kwargs))
chunks = chunker._outputs["chunks"]
texts = [c["text"] for c in chunks]
assert texts == ["aaa", "bbb", "ccc"], f"token_size={chunk_token_size} atom-split a delimiter segment: {texts}"
assert texts == ["aaa", "bbb", "ccc"], f"chunk_token_size={chunk_token_size} atom-split a delimiter segment: {texts}"