zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-06 23:51:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(chunker): drop dead atom-split path in mergeByTokenSizeFromJSON (#17873)

Browse Source
This commit is contained in:
bigfish-emily
2026-08-05 19:58:51 +08:00
committed by GitHub
parent d2183e610d
commit 2403988b6d
4 changed files with 33 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
internal/ingestion/component/chunker/token.go
View File

@@ -334,7 +334,7 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string
// Split-then-merge: split on delimiters, then greedily merge to
// chunk_token_size with optional overlap.
perItem := [][]schema.ChunkDoc{docs}
merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true, c.param.MergeStrategy())
merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, c.param.MergeStrategy())
return chunkOutputs(flatten(merged))
}
@@ -719,7 +719,7 @@ func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []s
// chunks across JSON items into one global token budget. Flatten the
// per-item structure into a single sequence first so the merge is
// global; non-text chunks still break the merge via their CKType.
attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false, c.param.MergeStrategy())
attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, c.param.MergeStrategy())
}
flat := flatten(attached)
@@ -956,15 +956,16 @@ func takeFromStart(text string, tokens int) string {
// mergeByTokenSizeFromJSON mirrors Python naive_merge's projected-total
// hard cap (rag/nlp/__init__.py after the strict chunk_token_num fix).
// Oversized text units are sub-split via splitOversizedUnit before merge;
// overlap is applied only when overlap+segment still fits the budget.
// Over-budget units are never atom-split: each one stands alone as its own
// chunk (Python naive_merge behavior, #17808 OVER_CAP contract). Overlap is
// applied only when overlap+segment still fits the budget.
//
// strategy selects the merge strategy (schema.MergeStrategy): MergeOverCap =
// OVER_CAP (Python's canonical default, a chunk may exceed the target by at
// most one incoming unit), MergeUnderCap = UNDER_CAP (never exceed the target;
// a projected overflow starts a fresh chunk). The TokenChunker threads its
// MergeStrategy() here.
func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool, strategy schema.MergeStrategy) [][]schema.ChunkDoc {
func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, strategy schema.MergeStrategy) [][]schema.ChunkDoc {
// overlappedPct is a [0,100] percentage. Clamp defensively because this
// helper is also exercised directly by tests.
if overlappedPct < 0 {
@@ -1059,30 +1060,10 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over
addTextChunk(ck)
continue
}
// Over-budget unit.
if !subSplitOversize {
// JSON path: Python keeps each over-budget item whole — it does
// not sub-split a single item, so emit it as one chunk.
addTextChunk(ck)
continue
}
// Text path: hard-cap atomic oversize units before merge, matching
// Python's _split_oversized_unit.
slog.Debug("TokenChunker: splitting oversized JSON unit via splitOversizedUnit",
"len", len(ck.Text), "tokens", tk, "chunk_token_size", chunkTokens)
for _, piece := range splitOversizedUnit(ck.Text, chunkTokens) {
if strings.TrimSpace(piece) == "" {
continue
}
cp := cloneChunkDoc(ck)
cp.Text = piece
cp.TKNums = intPtr(tokenizeStr(piece))
// Coordinates stay on the first piece only to avoid duplicating
// PDF bboxes across atom slices.
addTextChunk(cp)
ck.PDFPositions = nil
ck.Positions = nil
}
// Over-budget unit: keep it whole. Python's naive_merge never
// sub-splits a single item, so emit it as one chunk and let the
// model layer truncate it later.
addTextChunk(ck)
}
perItem[idx] = merged
}

20
internal/ingestion/component/chunker/token_batch1_test.go
View File

@@ -91,7 +91,7 @@ func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) {
{Text: cText, DocType: "text", CKType: "text", TKNums: intPtr(cN)},
},
}
got := mergeByTokenSizeFromJSON(items, budget, 30.0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, budget, 30.0, schema.MergeOverCap)
merged := got[0]
if len(merged) != 2 {
t.Fatalf("want 2 chunks (overflow-closed + overlap chunk), got %d (a=%d b=%d c=%d budget=%d)", len(merged), aN, bN, cN, budget)
@@ -157,7 +157,7 @@ func TestMergeByTokenSizeFromJSON_NonTextBoundaryResetsPrevClosed(t *testing.T)
{Text: t4, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(t4))},
},
}
got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, budget, 0.0, schema.MergeOverCap)
merged := got[0]
// Expect: chunk0 (T1+T2, overflow-closed), N (non-text), chunk1 (T3+T4 merged).
if len(merged) != 3 {
@@ -213,7 +213,7 @@ func TestMergeByTokenSizeFromJSON_UnderCapNoOverflow(t *testing.T) {
{Text: cText, DocType: "text", CKType: "text", TKNums: intPtr(cN)},
},
}
got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, schema.MergeUnderCap)
got := mergeByTokenSizeFromJSON(items, budget, 0.0, schema.MergeUnderCap)
merged := got[0]
if len(merged) != 3 {
t.Fatalf("UNDER_CAP want 3 chunks (a, b, c separate), got %d", len(merged))
@@ -255,12 +255,12 @@ func clampOverlapFixture() [][]schema.ChunkDoc {
}
func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) {
at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true, schema.MergeOverCap)
at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, schema.MergeOverCap)
if at100 == nil || len(at100) == 0 {
t.Fatalf("overlappedPct=100: nil/empty result")
}
at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true, schema.MergeOverCap)
atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true, schema.MergeOverCap)
at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, schema.MergeOverCap)
atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, schema.MergeOverCap)
if !reflect.DeepEqual(at100, at150) {
t.Errorf("overlappedPct=150 should clamp to 100; output differs from 100")
}
@@ -268,12 +268,12 @@ func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) {
t.Errorf("overlappedPct=1e300 should clamp to 100; output differs from 100")
}
at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true, schema.MergeOverCap)
at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, schema.MergeOverCap)
if at0 == nil || len(at0) == 0 {
t.Fatalf("overlappedPct=0: nil/empty result")
}
atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true, schema.MergeOverCap)
atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true, schema.MergeOverCap)
atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, schema.MergeOverCap)
atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, schema.MergeOverCap)
if !reflect.DeepEqual(at0, atNeg) {
t.Errorf("overlappedPct=-5 should clamp to 0; output differs from 0")
}
@@ -294,7 +294,7 @@ func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) {
{Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)},
},
}
got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, 128, 0, schema.MergeOverCap)
merged := got[0]
if len(merged) != 1 {
t.Fatalf("want 1 merged chunk, got %d", len(merged))

6
internal/ingestion/component/chunker/token_pdfpos_test.go
View File

@@ -38,7 +38,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) {
{Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB},
},
}
got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, 128, 0, schema.MergeOverCap)
merged := got[0]
if len(merged) != 1 {
t.Fatalf("want 1 merged chunk, got %d", len(merged))
@@ -63,7 +63,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) {
{Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB},
},
}
got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, 128, 0, schema.MergeOverCap)
combined := string(got[0][0].Positions)
if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") {
t.Errorf("merged chunk dropped/omitted `positions`: %s", combined)
@@ -103,7 +103,7 @@ func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) {
{Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB},
},
}
got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON(items, 128, 0, schema.MergeOverCap)
m := got[0][0].ToMap()
raw, ok := m["positions"]
if !ok {

31
internal/ingestion/component/chunker/token_strict_cap_test.go
View File

@@ -108,7 +108,7 @@ func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) {
Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)),
})
}
got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, schema.MergeOverCap)
merged := got[0]
if len(merged) < 3 {
t.Fatalf("want >=3 chunks, got %d", len(merged))
@@ -135,7 +135,7 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) {
Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)),
})
}
got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true, schema.MergeOverCap)
got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, schema.MergeOverCap)
unit := tokenizeStr(sections[0].Text)
for i, ck := range got[0] {
// OVER_CAP allows one boundary overflow (prev + one unit). The JSON
@@ -150,31 +150,20 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) {
}
}
func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) {
// A single unit larger than the budget must be atom-split before merge.
func TestMergeByTokenSizeFromJSON_OversizedUnitStaysWhole(t *testing.T) {
// Per the #17808 contract an over-budget unit is never atom-split: it
// stands alone as one chunk and the model layer truncates it later.
const budget = 30
long := strings.TrimSpace(strings.Repeat("word ", 100))
items := [][]schema.ChunkDoc{{
{Text: long, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(long))},
}}
got := mergeByTokenSizeFromJSON(items, budget, 0, true, schema.MergeOverCap)
if len(got[0]) < 2 {
t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0]))
got := mergeByTokenSizeFromJSON(items, budget, 0, schema.MergeOverCap)
if len(got) != 1 || len(got[0]) != 1 {
t.Fatalf("over-budget unit must stay whole, got %d chunk(s)", 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 (each sub-split piece <= budget+1).
// OVER_CAP then merges at most two such pieces into one chunk before
// closing it, so the invariant we defend is that no chunk exceeds
// 2*(budget+1): the oversized unit is sub-split (not collapsed into one
// chunk) and at most one boundary overflow is allowed — matching the
// Python reference.
const slack = 2 * (budget + 1)
for i, ck := range got[0] {
if n := tokenizeStr(ck.Text); n > slack {
t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d (cap=%d)", i, n, budget)
}
if got[0][0].Text != long {
t.Errorf("over-budget chunk text changed: got %q", got[0][0].Text)
}
}