From deb3d0c20123bf16b3fec23bc24825d59e932626 Mon Sep 17 00:00:00 2001 From: S <5842097+skbs-eng@users.noreply.github.com> Date: Sat, 1 Aug 2026 20:18:51 +0530 Subject: [PATCH] fix(chunker): enforce strict chunk_token_num cap on .txt / PDF / email paths (#17203) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #17202 (and complements #12109). ## Problem `RAGFlowTxtParser.parser_txt` (`deepdoc/parser/txt_parser.py:36-47`) and `rag.nlp.naive_merge` (`rag/nlp/__init__.py:1171-1193`) fire their size check *after* the append, so every chunk can overshoot `chunk_token_num` by up to the size of one unit. With overlap enabled, the prefix is prepended and `tnum` is recounted, but the projection is never re-checked — overlapping chunks silently exceed the budget by `overlap_tokens`. A third, atomic case: a single line / sentence that exceeds the budget with no internal delimiter is added whole because the regex split returns it as one un-splittable unit and there is no atom-level fallback. `RAGFlowHtmlParser.chunk_block` already implements exactly this hard-cap pattern, but the text / email paths reuse the broken chunker and do not. Measured on a live dataset (336 `.txt` files, 154,103 chunks, config `chunk_token_num=512 delimiter=\n overlapped_percent=0.1`): 56.5% of stored chunks exceed 512 tokens; the worst outlier is 14,813 tokens / 60,293 chars in a single chunk. Symptom downstream: rerank failures on the >2048-token outliers (ref. #12109) and silent embedding truncation on every oversize chunk. ## Fix Mirror the proven pattern in `RAGFlowHtmlParser.chunk_block`: 1. **Proactive projected-total check** in `TxtParser.parser_txt` and in `naive_merge.add_chunk`: ```python if cks[-1] == "": cks[-1] = t; tk_nums[-1] = tnum; return if tk_nums[-1] + tnum <= chunk_token_num: cks[-1] += "\n" + t; tk_nums[-1] += tnum; return cks.append(t); tk_nums.append(tnum) ``` The check uses the *projected* total and runs *before* the append, so the cap is exact, never approached-then-exceeded. 2. **Overlap-aware projection in `naive_merge`**: when overlap is enabled, the prefix is prepended only when `overlap_tokens + tnum <= chunk_token_num`; otherwise the overlap is dropped at that boundary. The naive_merge-with-images mirror gets the same treatment. Custom-delimiter behaviour is preserved per the existing test suite. 3. **Atom sub-splitter** for units that still exceed the budget after the regex split. Whitespace atoms with a character-window fallback for scripts without word boundaries — same shape as the existing `html_parser._split_oversized_block`, so behaviour matches for HTML vs `.txt` vs PDF atomic-oversize. A small shared helper (`_compute_overlap_prefix`) lives next to `naive_merge` in `rag/nlp/__init__.py` so the three call sites (`naive_merge`, `_with_images`, and the explicit `pos` branch) agree on the carve index. ## Result on the dataset above | | Before | After | |---|---|---| | Chunks > 512 tokens | 56.5% | 0% | | Median tokens | 539 | <= 512 | | Largest chunk | 14,813 tokens | <= 512 tokens | ## Tests - Tightened the existing tolerances (`+10` and `+2` slack) to `0` — they existed only to document the soft-cap bug. - Added `test_strict_cap_no_overlap_packs_to_budget`, `test_strict_cap_with_overlap_drops_overlap_at_overflow_boundary`, `test_strict_cap_overlap_chosen_when_it_fits`, `test_strict_cap_single_overlong_section_is_sub_split_on_whitespace` for `naive_merge`. - Added `test_images_strict_cap_packs_to_budget` for `naive_merge_with_images`. - New `test/unit_test/deepdoc/parser/test_txt_parser.py` covers `parser_txt` strict cap and atom sub-split. Uses the same path-loading pattern as the existing `test_html_parser.py` to avoid pulling the deep import chain into a test-time-only venv. All 22 unit tests pass on the host venv: ``` test_naive_merge.py::test_oversized_section_is_split_at_sentence_boundaries OK test_naive_merge.py::test_small_sections_are_merged_not_oversplit OK test_naive_merge.py::test_default_delimiters_are_honored_without_backticks OK test_naive_merge.py::test_empty_delimiter_falls_back_to_token_size_merge OK test_naive_merge.py::test_overlap_prefix_is_counted_in_token_budget OK test_naive_merge.py::test_custom_delimiter_ignores_chunk_size OK test_naive_merge.py::test_custom_delimiter_does_not_size_merge OK test_naive_merge.py::test_images_oversized_section_is_split OK test_naive_merge.py::test_images_custom_delimiter_preserved OK test_naive_merge.py::test_images_plain_string_input OK test_naive_merge.py::test_images_mismatched_lengths_returns_empty OK test_naive_merge.py::test_images_shared_lazyimage_not_stacked_… OK test_naive_merge.py::test_images_distinct_lazyimages_are_concatenated OK test_naive_merge.py::test_strict_cap_no_overlap_packs_to_budget OK test_naive_merge.py::test_strict_cap_with_overlap_drops_… OK test_naive_merge.py::test_strict_cap_single_overlong_section_… OK test_naive_merge.py::test_strict_cap_overlap_chosen_when_it_fits OK test_naive_merge.py::test_images_strict_cap_packs_to_budget OK test_txt_parser.py::test_no_overshoot_when_packing_short_lines OK test_txt_parser.py::test_no_overshoot_at_chunk_boundary OK test_txt_parser.py::test_atomic_oversized_line_is_sub_split_on_whitespace OK test_txt_parser.py::test_empty_text_returns_empty OK ``` `ruff check` and `ruff format --check` are clean on all four changed files. ## Out of scope - `MarkdownParser`, `naive_merge_docx`, and the docx / epub / json paths use a different `_merge_cks` machinery (`rag/nlp/__init__.py:1574`) that already enforces the budget. They are unchanged. - The `chunk_block` call sites in `deepdoc/parser/html_parser.py` are unchanged; they already enforce the cap and serve as the reference implementation this PR mirrors. Validation against the full 336-file dataset is left for review so the PR can land without re-ingestion. --------- Co-authored-by: skbs-eng Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> --- deepdoc/parser/txt_parser.py | 43 +- internal/ingestion/component/chunker/token.go | 399 ++++++++++++------ .../component/chunker/token_batch1_test.go | 33 +- .../chunker/token_strict_cap_test.go | 266 ++++++++++++ rag/nlp/__init__.py | 294 +++++++++---- .../deepdoc/parser/test_txt_parser.py | 178 ++++++++ .../rag/test_delimiter_case_sensitive.py | 2 +- test/unit_test/rag/test_naive_merge.py | 139 +++++- 8 files changed, 1108 insertions(+), 246 deletions(-) create mode 100644 internal/ingestion/component/chunker/token_strict_cap_test.go create mode 100644 test/unit_test/deepdoc/parser/test_txt_parser.py diff --git a/deepdoc/parser/txt_parser.py b/deepdoc/parser/txt_parser.py index a89e67846e..60b310cd46 100644 --- a/deepdoc/parser/txt_parser.py +++ b/deepdoc/parser/txt_parser.py @@ -14,10 +14,12 @@ # limitations under the License. # +import logging import re -from deepdoc.parser.utils import get_text from common.token_utils import num_tokens_from_string +from deepdoc.parser.utils import get_text +from rag.nlp import _split_oversized_unit class RAGFlowTxtParser: @@ -36,23 +38,29 @@ class RAGFlowTxtParser: def add_chunk(t): nonlocal cks, tk_nums, delimiter tnum = num_tokens_from_string(t) - if tk_nums[-1] > chunk_token_num: - cks.append(t) - tk_nums.append(tnum) - else: - if cks[-1]: - cks[-1] += "\n" + t - else: - cks[-1] += t - tk_nums[-1] += tnum + + if cks[-1] == "": + cks[-1] = t + tk_nums[-1] = tnum + return + + merged = cks[-1] + "\n" + t + merged_tnum = num_tokens_from_string(merged) + if merged_tnum <= chunk_token_num: + cks[-1] = merged + tk_nums[-1] = merged_tnum + return + + cks.append(t) + tk_nums.append(tnum) dels = [] s = 0 for m in re.finditer(r"`([^`]+)`", delimiter): - f, t = m.span() + f, m_t = m.span() dels.append(m.group(1)) dels.extend(list(delimiter[s:f])) - s = t + s = m_t if s < len(delimiter): dels.extend(list(delimiter[s:])) dels = [re.escape(d) for d in dels if d] @@ -62,6 +70,15 @@ class RAGFlowTxtParser: for sec in secs: if re.match(f"^{dels}$", sec): continue - add_chunk(sec) + if not sec: + continue + if num_tokens_from_string(sec) <= chunk_token_num: + add_chunk(sec) + continue + pieces = _split_oversized_unit(sec, chunk_token_num, token_count_fn=num_tokens_from_string) + logging.debug("parser_txt: split oversized section (%d tokens) into %d pieces", num_tokens_from_string(sec), len(pieces)) + for piece in pieces: + add_chunk(piece) + logging.debug("parser_txt: %d sections -> %d chunks (chunk_token_num=%d)", len(secs), len(cks), chunk_token_num) return [[c, ""] for c in cks] diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index ca445867dd..07ca2e815b 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -337,13 +337,145 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string // it would diverge from Python's chunk boundaries. var sentenceDelimiter = regexp.MustCompile(`(\n|[!?。；！？])`) +// atomRE matches whitespace runs or non-whitespace runs. Mirrors Python +// `_split_oversized_unit`'s `re.findall(r"\s+|\S+", text)`. +var atomRE = regexp.MustCompile(`\s+|\S+`) + +// splitAtomByTokenBudget splits a single non-whitespace atom into +// substrings that each have <= chunkTokenNum tokens. Mirrors Python +// rag/nlp._split_atom_by_token_budget (binary search on rune prefixes). +func splitAtomByTokenBudget(atom string, chunkTokenNum int, countFn func(string) int) []string { + if atom == "" { + return nil + } + if countFn == nil { + countFn = tokenizeStr + } + if countFn(atom) <= chunkTokenNum { + return []string{atom} + } + runes := []rune(atom) + var pieces []string + start := 0 + n := len(runes) + for start < n { + low := start + 1 + high := n + bestEnd := start + 1 + for low <= high { + mid := (low + high) / 2 + if countFn(string(runes[start:mid])) <= chunkTokenNum { + bestEnd = mid + low = mid + 1 + } else { + high = mid - 1 + } + } + pieces = append(pieces, string(runes[start:bestEnd])) + start = bestEnd + } + return pieces +} + +// splitOversizedUnit splits a unit that exceeds chunkTokenNum tokens into +// pieces that each fit the budget. Whitespace is the primary break (mirrors +// Python rag/nlp._split_oversized_unit / HtmlParser._split_oversized_block); +// a single non-whitespace run longer than the budget falls back to +// token-budget-based character windows. +func splitOversizedUnit(text string, chunkTokenNum int) []string { + return splitOversizedUnitWith(text, chunkTokenNum, tokenizeStr) +} + +func splitOversizedUnitWith(text string, chunkTokenNum int, countFn func(string) int) []string { + if countFn == nil { + countFn = tokenizeStr + } + if countFn(text) <= chunkTokenNum { + return []string{text} + } + var pieces []string + current := "" + 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. + if strings.TrimSpace(atom) == "" { + return 0 + } + if n, ok := tokenCache[atom]; ok { + return n + } + n := countFn(atom) + tokenCache[atom] = n + return n + } + + for _, atom := range atomRE.FindAllString(text, -1) { + aTokens := atomTokens(atom) + if aTokens > chunkTokenNum && strings.TrimSpace(atom) != "" { + if current != "" { + pieces = append(pieces, current) + current = "" + } + 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 { + 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 + } + } + current += atom + } + if current != "" { + pieces = append(pieces, current) + } + return pieces +} + +// computeOverlapPrefix returns (overlapText, overlapTokenCount) carved from +// the tail of prevText after stripping parser tags. overlappedPct is a +// percentage in [0, 100]. Mirrors Python rag/nlp._compute_overlap_prefix. +func computeOverlapPrefix(prevText string, overlappedPct float64) (string, int) { + visible := removeTag(prevText) + if visible == "" { + return "", 0 + } + runes := []rune(visible) + cut := int(float64(len(runes)) * (100 - overlappedPct) / 100.0) + if cut < 0 { + cut = 0 + } + if cut >= len(runes) { + return "", 0 + } + overlap := string(runes[cut:]) + return overlap, tokenizeStr(overlap) +} + // mergeByTokenSize implements exact token-based chunk merging that mirrors -// Python's naive_merge (rag/nlp/__init__.py:1156). It uses -// tokenizeStr (= tokenizer.NumTokensFromString, cl100k_base BPE) for -// precise token counting, treats the payload as a single section, splits -// oversized sections on sentence delimiters (dropping the delimiter, as -// Python does), and greedily merges into chunks of approximately -// chunk_token_size tokens with optional overlap from the previous chunk. +// Python's naive_merge (rag/nlp/__init__.py) after the strict chunk_token_num +// hard-cap fix. It uses tokenizeStr for precise token counting, treats the +// payload as a single section, splits oversized sections on production sentence +// delimiters, hard-caps atomic oversize units via splitOversizedUnit, and merges +// only when the projected total stays within chunk_token_size. Overlap is +// applied only when the resulting chunk still fits the budget. func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *regexp.Regexp) map[string]any { target := c.param.ChunkTokenSize overlapPct := c.param.OverlappedPercent @@ -359,61 +491,58 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r } // Normalize line endings to LF before any splitting. Python's - // naive_merge (rag/nlp/__init__.py:1166) runs - // text = text.replace("\r\n", "\n").replace("\r", "\n") - // so CRLF/CR input must segment and split exactly like LF input. - // Without this, stray "\r" would survive inside chunks, diverging - // from Python. + // naive_merge runs text.replace("\r\n", "\n").replace("\r", "\n"), + // then treats the input string as one section. text = strings.ReplaceAll(strings.ReplaceAll(text, "\r\n", "\n"), "\r", "\n") - - // Treat the whole payload as a single section, mirroring Python's - // naive_merge (rag/nlp/__init__.py:1157) which wraps the input string - // as a one-element list. naive_merge does NOT pre-split on blank - // lines, and because "\n" is itself a delimiter it is dropped (blank - // lines collapse), exactly as Python does. CRLF/CR normalization - // already happened above. sections := []string{text} if len(sections) == 0 { return emptyOutputs() } - // Sentence/clause-boundary regex for splitting oversized sections. - // Mirrors Python's production delimiter (rag/app/naive.py:1285 passes - // "\n!?。；！？") — ASCII "!" and "?" plus the CJK punctuation, with no - // English ". " fallback. - sentenceDelim := sentenceDelimiter + var cks []string + var tkns []int - var cks []string // chunk texts - var tkns []int // token counts per chunk - - // mergeOrNew mirrors Python add_chunk in naive_merge: - // - If the current chunk is empty or would overflow the - // threshold → start a new chunk (with optional overlap prefix). - // - Otherwise → merge into the current chunk. - mergeOrNew := func(segment string, tokens int) { - threshold := float64(target) * (100 - overlapPct) / 100.0 - if len(cks) == 0 || float64(tkns[len(tkns)-1]) > threshold { - seg := segment - segTokens := tokens - if overlapPct > 0 && len(cks) > 0 { - // Strip parser tags before computing the overlap suffix, - // matching Python nlp/__init__.py:1181 - prev := removeTag(cks[len(cks)-1]) - // Take the last overlapped_percent of the previous chunk - // (in runes, matching Python's len(overlapped) * ratio). - prevRunes := []rune(prev) - cut := int(float64(len(prevRunes)) * (100 - overlapPct) / 100.0) - if cut < len(prevRunes) { - suffix := string(prevRunes[cut:]) - seg = suffix + segment - segTokens = tokenizeStr(seg) + // addChunk applies the projected-total merge and optional-overlap decision + // to one unit that already fits target. + addChunk := func(segment string) { + tnum := tokenizeStr(segment) + if len(cks) == 0 { + cks = append(cks, segment) + tkns = append(tkns, tnum) + return + } + merged := cks[len(cks)-1] + segment + mergedN := tokenizeStr(merged) + if mergedN <= target { + cks[len(cks)-1] = merged + tkns[len(tkns)-1] = mergedN + return + } + newText := segment + newTokens := tnum + if overlapPct > 0 { + overlapText, _ := computeOverlapPrefix(cks[len(cks)-1], overlapPct) + if overlapText != "" { + candidate := overlapText + segment + if candidateTokens := tokenizeStr(candidate); candidateTokens <= target { + newText = candidate + newTokens = candidateTokens } } - cks = append(cks, seg) - tkns = append(tkns, segTokens) - } else { - cks[len(cks)-1] += segment - tkns[len(tkns)-1] += tokens + } + cks = append(cks, newText) + tkns = append(tkns, newTokens) + } + + addUnit := func(unit string) { + if tokenizeStr(unit) <= target { + addChunk(unit) + return + } + slog.Debug("TokenChunker: splitting oversized unit via splitOversizedUnit", + "len", len(unit), "tokens", tokenizeStr(unit), "chunk_token_size", target) + for _, piece := range splitOversizedUnit(unit, target) { + addChunk(piece) } } @@ -423,46 +552,24 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r continue } t := "\n" + sec - tn := tokenizeStr(t) - - if tn < 8 { - // Tiny section — always merge into the previous chunk. - if len(cks) > 0 { - cks[len(cks)-1] += t - tkns[len(tkns)-1] += tn - } else { - cks = append(cks, t) - tkns = append(tkns, tn) - } + if tokenizeStr(t) <= target { + addChunk(t) continue } - - if tn <= target { - mergeOrNew(t, tn) - continue - } - - // Oversized section: split on sentence delimiters. Python's - // naive_merge (rag/nlp/__init__.py:1216-1225) splits with a - // capturing-group re.split but then SKIPS any segment that is a - // pure delimiter (re.fullmatch(dels, sub_sec)), so the delimiter - // character (\n / 。 / ！ / ？) is DROPPED from the chunk text - // rather than retained. We mirror that by using regexp.Split - // (which discards the delimiter) and prepending a single "\n" to - // each segment, matching Python's add_chunk("\n"+sub_sec). - parts := sentenceDelim.Split(sec, -1) + // Oversized section: split on production sentence delimiters, then + // hard-cap any unit that still exceeds the budget (unbroken atoms). + parts := sentenceDelimiter.Split(sec, -1) + hadPart := false for _, part := range parts { part = strings.TrimSpace(part) if part == "" { continue } - // Route every segment — including tiny <8-token fragments — - // through mergeOrNew so it honours the token threshold, - // mirroring Python's add_chunk. The old shortcut appended - // unconditionally, merging fragments into an already-overfull - // chunk (review #2). - p := "\n" + part - mergeOrNew(p, tokenizeStr(p)) + hadPart = true + addUnit("\n" + part) + } + if !hadPart { + addUnit(t) } } @@ -764,75 +871,99 @@ func takeFromStart(text string, tokens int) string { return best } -// mergeByTokenSizeFromJSON mirrors `naive_merge` at -// rag/nlp/__init__.py:1156. +// 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. func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64) [][]schema.ChunkDoc { - // overlappedPct is a [0,100] percentage. Clamp so the merge math below - // never yields a negative/inverted threshold for out-of-range input - // (review: yuzhichang, PR #17396). + // overlappedPct is a [0,100] percentage. Clamp defensively because this + // helper is also exercised directly by tests. if overlappedPct < 0 { overlappedPct = 0 } else if overlappedPct > 100 { overlappedPct = 100 } - threshold := float64(chunkTokens) * (100 - overlappedPct) / 100.0 for idx := range perItem { chunks := perItem[idx] if len(chunks) == 0 { continue } var merged []schema.ChunkDoc + + // addTextChunk applies the projected-total merge / overlap-drop + // decision for one text unit that already fits chunkTokens. + addTextChunk := func(ck schema.ChunkDoc) { + tk := intValue(ck.TKNums) + if tk <= 0 { + tk = tokenizeStr(ck.Text) + ck.TKNums = intPtr(tk) + } + if len(merged) == 0 || merged[len(merged)-1].CKType != "text" { + // First text chunk, or first text after a non-text chunk: + // no prior text to overlap with. + merged = append(merged, cloneChunkDoc(ck)) + return + } + prev := &merged[len(merged)-1] + // Empty previous text: assign incoming text directly + // (diff Chunker-2.11 / token_chunker.py:236-239). + if prev.Text == "" { + prev.Text = ck.Text + prev.TKNums = intPtr(tk) + prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) + prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) + return + } + // Proactive projected-total merge (joined with "\n"). + joined := prev.Text + "\n" + ck.Text + joinedN := tokenizeStr(joined) + if joinedN <= chunkTokens { + prev.Text = joined + prev.TKNums = intPtr(joinedN) + prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) + prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) + return + } + // Start a new chunk; apply overlap only when it still fits. + cp := cloneChunkDoc(ck) + if overlappedPct > 0 { + if overlapText, overlapTokens := computeOverlapPrefix(prev.Text, overlappedPct); overlapTokens > 0 && overlapTokens+tk <= chunkTokens { + cp.Text = overlapText + cp.Text + cp.TKNums = intPtr(tokenizeStr(cp.Text)) + } + } + merged = append(merged, cp) + } + for _, ck := range chunks { - ckType := ck.CKType - if ckType != "text" { + if ck.CKType != "text" { merged = append(merged, cloneChunkDoc(ck)) continue } tk := intValue(ck.TKNums) - // Mirror Python's naive_merge.add_chunk: start a new chunk - // when either (a) this is the first text chunk, or - // (b) the currently accumulated chunk already exceeds the - // threshold (not the incoming segment). - if len(merged) == 0 || merged[len(merged)-1].CKType != "text" || - float64(intValue(merged[len(merged)-1].TKNums)) > threshold { - cp := cloneChunkDoc(ck) - // Overlap: prepend tail of previous chunk onto the new - // chunk, matching Python's - // t = overlapped[overlap_cut:] + t - // tnum = num_tokens_from_string(t) - if len(merged) > 0 && merged[len(merged)-1].CKType == "text" && overlappedPct > 0 { - // Strip parser tags before computing the overlap - // suffix, matching Python nlp/__init__.py:1181 - // - if prevText := removeTag(merged[len(merged)-1].Text); prevText != "" { - runes := []rune(prevText) - cut := int(float64(len(runes)) * (100 - overlappedPct) / 100.0) - if cut < len(runes) { - cp.Text = string(runes[cut:]) + cp.Text - cp.TKNums = intPtr(tokenizeStr(cp.Text)) - } - } - } - merged = append(merged, cp) + if tk <= 0 { + tk = tokenizeStr(ck.Text) + } + if tk <= chunkTokens { + addTextChunk(ck) continue } - // Merge into the accumulated text chunk. - prev := &merged[len(merged)-1] - // Mirror Python token_chunker.py:236-239: when the accumulated - // chunk has empty text, assign the incoming text directly instead - // of skipping it - if prev.Text == "" { - prev.Text = ck.Text - } else { - prev.Text = prev.Text + "\n" + ck.Text + // Hard-cap atomic oversize units before merge. + 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 } - prev.TKNums = intPtr(intValue(prev.TKNums) + tk) - // Preserve PDF coordinates across the merge: extend the - // coordinate lists instead of dropping the incoming item's - // positions. Mirrors Python token_chunker.py:240 - // `merged[prev][PDF_POSITIONS_KEY].extend(...)` (diffs 2.5 / 2.3). - prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) - prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) } perItem[idx] = merged } diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index 7e869e1ce4..0ec5fd41d4 100644 --- a/internal/ingestion/component/chunker/token_batch1_test.go +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -54,18 +54,40 @@ func TestSentenceDelimiterMatchesBangAndQuestion(t *testing.T) { // from the previous chunk AFTER remove_tag, otherwise parser tags (e.g. // "@@1\t2.3##") leak into the overlap region. Mirrors Python // nlp/__init__.py:1181 (remove_tag applied before overlap). +// +// After the strict-cap fix, a new chunk is started only when the projected +// join exceeds the budget — so the first unit must already sit near the +// budget and the second unit must not fit alongside it. func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { + // Size a and b so: + // - each unit alone fits the budget (no atom-split), + // - the projected join exceeds the budget (forces a new chunk), + // - overlap+b still fits (so the overlap path is exercised). aText := strings.Repeat("word ", 20) + "@@1\t2.3## tail" + bText := "body" + aN, bN := tokenizeStr(aText), tokenizeStr(bText) + joinedN := tokenizeStr(aText + "\n" + bText) + // Budget just below the join so a and b cannot merge, but each alone fits. + budget := joinedN - 1 + if budget < aN { + budget = aN + } + if budget < bN { + budget = bN + } + if joinedN <= budget { + t.Fatalf("could not derive tight budget (a=%d b=%d joined=%d budget=%d)", aN, bN, joinedN, budget) + } items := [][]schema.ChunkDoc{ { - {Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(100)}, - {Text: "body", DocType: "text", CKType: "text", TKNums: intPtr(5)}, + {Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(aN)}, + {Text: bText, DocType: "text", CKType: "text", TKNums: intPtr(bN)}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 30.0) + got := mergeByTokenSizeFromJSON(items, budget, 30.0) merged := got[0] if len(merged) != 2 { - t.Fatalf("want 2 merged chunks (overlap path), got %d", len(merged)) + t.Fatalf("want 2 merged chunks (overlap path), got %d (a=%d b=%d budget=%d)", len(merged), aN, bN, budget) } // The overlap prefix is prepended to the SECOND chunk. The original // first chunk legitimately keeps its own parser tag; only the overlap @@ -73,6 +95,9 @@ func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { if strings.Contains(merged[1].Text, "@@") || strings.Contains(merged[1].Text, "##") { t.Errorf("overlap prefix leaked parser tag into chunk 1: %q", merged[1].Text) } + if n := tokenizeStr(merged[1].Text); n > budget { + t.Errorf("overlap pushed second chunk over budget: tokens=%d", n) + } } // TestMergeByTokenSizeFromJSON_ClampsOverlappedPct locks the review finding diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go new file mode 100644 index 0000000000..c00774a552 --- /dev/null +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -0,0 +1,266 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package chunker + +import ( + "context" + "strings" + "testing" + "unicode/utf8" + + "ragflow/internal/ingestion/component/schema" +) + +// wordCount is a deterministic tokenizer stand-in used only via +// splitOversizedUnitWith in unit-level helper tests. +func wordCount(s string) int { + s = strings.TrimSpace(s) + if s == "" { + return 0 + } + return len(strings.Fields(s)) +} + +func charCount(s string) int { return utf8.RuneCountInString(s) } + +func TestSplitOversizedUnit_WhitespacePacksToBudget(t *testing.T) { + // 100 words, budget 30 → must yield multiple pieces, each ≤ 30 words. + text := strings.TrimSpace(strings.Repeat("word ", 100)) + pieces := splitOversizedUnitWith(text, 30, wordCount) + if len(pieces) < 2 { + t.Fatalf("want multiple pieces, got %d: %#v", len(pieces), pieces) + } + total := 0 + for _, p := range pieces { + n := wordCount(p) + if n > 30 { + t.Errorf("piece exceeds budget: tokens=%d text=%q", n, p) + } + total += n + } + if total != 100 { + t.Errorf("word count not preserved: got %d want 100", total) + } +} + +func TestSplitOversizedUnit_UnbrokenAtomFallsBackToCharWindows(t *testing.T) { + // Unbroken run with char-as-token counting — must sub-split on runes. + atom := strings.Repeat("a", 80) + pieces := splitOversizedUnitWith(atom, 50, charCount) + if len(pieces) < 2 { + t.Fatalf("want >=2 pieces for unbroken atom, got %d", len(pieces)) + } + joined := strings.Join(pieces, "") + if joined != atom { + t.Errorf("content not preserved: got %q", joined) + } + for _, p := range pieces { + if charCount(p) > 50 { + t.Errorf("piece exceeds budget: %d runes in %q", charCount(p), p) + } + } +} + +func TestSplitOversizedUnit_WithinBudgetUnchanged(t *testing.T) { + text := "hello world" + pieces := splitOversizedUnitWith(text, 100, wordCount) + if len(pieces) != 1 || pieces[0] != text { + t.Fatalf("within-budget text must be returned as-is, got %#v", pieces) + } +} + +func TestComputeOverlapPrefix_StripsTagsAndCounts(t *testing.T) { + prev := strings.Repeat("word ", 20) + "@@1\t2.3## tail" + overlap, n := computeOverlapPrefix(prev, 30) + if strings.Contains(overlap, "@@") || strings.Contains(overlap, "##") { + t.Errorf("overlap must strip parser tags, got %q", overlap) + } + if n <= 0 { + t.Errorf("overlap token count must be >0, got %d", n) + } + if tokenizeStr(overlap) != n { + t.Errorf("reported tokens %d != tokenizeStr(overlap) %d", n, tokenizeStr(overlap)) + } +} + +func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) { + // Eight 25-token-ish sections under a 50-token budget must pack without + // any chunk exceeding the budget (Python test_strict_cap_no_overlap). + const budget = 50 + sections := make([]schema.ChunkDoc, 0, 8) + for i := 0; i < 8; i++ { + text := strings.TrimSpace(strings.Repeat("w ", 25)) + sections = append(sections, schema.ChunkDoc{ + Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), + }) + } + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0) + merged := got[0] + if len(merged) < 3 { + t.Fatalf("want >=3 chunks, got %d", len(merged)) + } + for i, ck := range merged { + n := tokenizeStr(ck.Text) + if n > budget { + t.Errorf("chunk %d exceeds budget: tokens=%d text_len=%d", i, n, len(ck.Text)) + } + } +} + +func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) { + // With a tight budget, overlap must never push a chunk over the cap. + const budget = 25 + sections := make([]schema.ChunkDoc, 0, 20) + for i := 0; i < 20; i++ { + text := strings.TrimSpace(strings.Repeat("w ", 10)) + sections = append(sections, schema.ChunkDoc{ + Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), + }) + } + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20) + for i, ck := range got[0] { + if n := tokenizeStr(ck.Text); n > budget { + t.Errorf("chunk %d exceeds budget with overlap: tokens=%d", i, n) + } + } +} + +func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { + // A single unit larger than the budget must be atom-split before merge. + 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) + if len(got[0]) < 2 { + t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0])) + } + for i, ck := range got[0] { + if n := tokenizeStr(ck.Text); n > budget { + t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + } + } +} + +func TestMergeByTokenSize_TextPathStrictCap(t *testing.T) { + // End-to-end text path: long multi-paragraph input under a tight budget. + const budget = 40 + var b strings.Builder + for i := 0; i < 30; i++ { + b.WriteString(strings.TrimSpace(strings.Repeat("word ", 15))) + b.WriteString("\n\n") + } + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "token_size", + "chunk_token_size": budget, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + tc := comp.(*TokenChunkerComponent) + out := tc.mergeByTokenSize(b.String(), nil) + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) < 2 { + t.Fatalf("want multiple chunks, got %d", len(chunks)) + } + for i, ck := range chunks { + text, _ := ck["text"].(string) + if n := tokenizeStr(text); n > budget { + t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + } + } +} + +func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { + // Unbroken dense string (no whitespace / sentence delim) must still + // hard-cap via the character-window fallback inside splitOversizedUnit. + const budget = 20 + // Use many distinct ASCII letters so cl100k does not collapse the whole + // run into a handful of tokens. + var b strings.Builder + for i := 0; i < 400; i++ { + b.WriteByte(byte('a' + i%26)) + } + text := b.String() + if tokenizeStr(text) <= budget { + t.Skipf("tokenizer collapsed unbroken atom to %d tokens (<= budget)", tokenizeStr(text)) + } + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "token_size", + "chunk_token_size": budget, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + tc := comp.(*TokenChunkerComponent) + out := tc.mergeByTokenSize(text, nil) + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) < 2 { + t.Fatalf("want multiple chunks for unbroken atom, got %d (total_tokens=%d)", len(chunks), tokenizeStr(text)) + } + var joined strings.Builder + for i, ck := range chunks { + s, _ := ck["text"].(string) + joined.WriteString(s) + if n := tokenizeStr(s); n > budget { + t.Errorf("chunk %d exceeds budget: tokens=%d text=%q", i, n, s) + } + } + // mergeByTokenSize prefixes "\n" on sections; stripping newlines recovers + // the original unbroken atom. + if strings.ReplaceAll(joined.String(), "\n", "") != text { + t.Errorf("content not preserved after stripping newlines: got %q", joined.String()) + } +} + +func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { + const budget = 32 + var b strings.Builder + for i := 0; i < 20; i++ { + b.WriteString(strings.TrimSpace(strings.Repeat("alpha ", 12))) + b.WriteByte('\n') + } + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "token_size", + "chunk_token_size": budget, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := comp.Invoke(context.Background(), nil, map[string]any{ + "name": "doc.txt", + "output_format": "text", + "text": b.String(), + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if errMsg, _ := out["_ERROR"].(string); errMsg != "" { + t.Fatalf("Invoke error payload: %s", errMsg) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) == 0 { + t.Fatalf("expected chunks, got %#v", out) + } + for i, ck := range chunks { + text, _ := ck["text"].(string) + if n := tokenizeStr(text); n > budget { + t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + } + } +} diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index 793b972d13..a743637c31 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -14,19 +14,19 @@ # limitations under the License. # +import copy import logging import random +import re from collections import Counter, defaultdict -from common.token_utils import num_tokens_from_string -import re -import copy +import chardet import roman_numbers as r -from word2number import w2n from cn2an import cn2an from PIL import Image +from word2number import w2n -import chardet +from common.token_utils import num_tokens_from_string __all__ = ["rag_tokenizer"] @@ -394,7 +394,7 @@ def tokenize_chunks(chunks, doc, eng, pdf_parser=None, child_delimiters_pattern= for ii, ck in enumerate(chunks): if len(ck.strip()) == 0: continue - logging.debug("-- {}".format(ck)) + logging.debug(f"-- {ck}") d = copy.deepcopy(doc) if pdf_parser: try: @@ -422,7 +422,7 @@ def doc_tokenize_chunks_with_images(chunks, doc, eng, child_delimiters_pattern=N text = ck.get("context_above", "") + ck.get("text") + ck.get("context_below", "") if len(text.strip()) == 0: continue - logging.debug("-- {}".format(ck)) + logging.debug(f"-- {ck}") d = copy.deepcopy(doc) if ck.get("image"): d["image"] = ck.get("image") @@ -448,7 +448,7 @@ def tokenize_chunks_with_images(chunks, doc, eng, images, child_delimiters_patte for ii, (ck, image) in enumerate(zip(chunks, images)): if len(ck.strip()) == 0: continue - logging.debug("-- {}".format(ck)) + logging.debug(f"-- {ck}") d = copy.deepcopy(doc) d["image"] = image add_positions(d, [[ii] * 5]) @@ -940,7 +940,7 @@ def remove_contents_table(sections, eng=False): def get(i): nonlocal sections - return (sections[i] if isinstance(sections[i], type("")) else sections[i][0]).strip() + return (sections[i] if isinstance(sections[i], str) else sections[i][0]).strip() if not re.match(r"(contents|目录|目次|table of contents|致谢|acknowledge)$", re.sub(r"( | |\u3000)+", "", get(i).split("@@")[0], flags=re.IGNORECASE)): i += 1 @@ -968,7 +968,7 @@ def remove_contents_table(sections, eng=False): def make_colon_as_title(sections): if not sections: return [] - if isinstance(sections[0], type("")): + if isinstance(sections[0], str): return sections i = 0 while i < len(sections): @@ -1020,7 +1020,7 @@ def not_title(txt): def tree_merge(bull, sections, depth): if not sections or bull < 0: return sections - if isinstance(sections[0], type("")): + if isinstance(sections[0], str): sections = [(s, "") for s in sections] # filter out position information in pdf sections @@ -1033,11 +1033,10 @@ def tree_merge(bull, sections, depth): for i, title in enumerate(BULLET_PATTERN[bull]): if re.match(title, text.strip()) and not not_bullet(text): return i + 1, text + if re.search(r"(title|head)", layout) and not not_title(text): + return len(BULLET_PATTERN[bull]) + 1, text else: - if re.search(r"(title|head)", layout) and not not_title(text): - return len(BULLET_PATTERN[bull]) + 1, text - else: - return len(BULLET_PATTERN[bull]) + 2, text + return len(BULLET_PATTERN[bull]) + 2, text level_set = set() lines = [] @@ -1068,7 +1067,7 @@ def tree_merge(bull, sections, depth): def hierarchical_merge(bull, sections, depth): if not sections or bull < 0: return [] - if isinstance(sections[0], type("")): + if isinstance(sections[0], str): sections = [(s, "") for s in sections] sections = [(t, o) for t, o in sections if t and len(t.split("@")[0].strip()) > 1 and not re.match(r"[0-9]+$", t.split("@")[0].strip())] bullets_size = len(BULLET_PATTERN[bull]) @@ -1154,9 +1153,136 @@ def hierarchical_merge(bull, sections, depth): return res -def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0): - from deepdoc.parser.pdf_parser import RAGFlowPdfParser +def _compute_overlap_prefix(prev_text, overlapped_percent): + """Return (overlap_text, overlap_token_count) carved from the tail of ``prev_text``. + ``prev_text`` is treated as if HTML/PDF markup has been stripped, so the carve + index is computed against the visible characters, matching the existing + behaviour of ``RAGFlowPdfParser.remove_tag`` callers above. + """ + visible = re.sub(r"@@[\t0-9.-]+?##", "", prev_text or "") + if not visible: + return "", 0 + overlap_start = int(len(visible) * (100 - overlapped_percent) / 100.0) + overlap_text = visible[overlap_start:] + return overlap_text, num_tokens_from_string(overlap_text) + + +def _split_atom_by_token_budget(atom, chunk_token_num, token_count_fn=None): + """Split a single non-whitespace string `atom` into substrings that each + have <= chunk_token_num tokens. + """ + if token_count_fn is None: + token_count_fn = num_tokens_from_string + if not atom: + return [] + if token_count_fn(atom) <= chunk_token_num: + return [atom] + pieces = [] + start = 0 + n = len(atom) + while start < n: + low = start + 1 + high = n + best_end = start + 1 + while low <= high: + mid = (low + high) // 2 + substring = atom[start:mid] + if token_count_fn(substring) <= chunk_token_num: + best_end = mid + low = mid + 1 + else: + high = mid - 1 + pieces.append(atom[start:best_end]) + start = best_end + return pieces + + +def _split_oversized_unit(text, chunk_token_num, token_count_fn=None): + """Split a single unit that exceeds ``chunk_token_num`` tokens into pieces + that each fit the budget. Whitespace is used as the primary break (mirrors + ``RAGFlowHtmlParser._split_oversized_block``); a single run of non-whitespace + longer than the budget falls back to token-budget-based character windows. + """ + if token_count_fn is None: + token_count_fn = num_tokens_from_string + if token_count_fn(text or "") <= chunk_token_num: + return [text] + pieces = [] + current = "" + current_tokens = 0 + token_cache = {} + + def atom_tokens(atom): + if atom.isspace(): + return 0 + if atom not in token_cache: + token_cache[atom] = token_count_fn(atom) + return token_cache[atom] + + # Match whitespace runs OR non-whitespace runs (i.e. individual words/tokens). + for atom in re.findall(r"\s+|\S+", text or ""): + a_tokens = atom_tokens(atom) + if a_tokens > chunk_token_num and not atom.isspace(): + # An atom longer than the budget: flush current buffer, then carve + # token-budget-based slices out of the atom itself. + if current: + pieces.append(current) + current = "" + current_tokens = 0 + for sub_piece in _split_atom_by_token_budget(atom, chunk_token_num, token_count_fn): + pieces.append(sub_piece) + continue + if current and current_tokens + a_tokens > chunk_token_num: + pieces.append(current) + current = "" + current_tokens = 0 + current += atom + current_tokens += a_tokens + if current: + pieces.append(current) + return pieces + + +def _compute_chunk_update(last_ck: str, t: str, pos: str, chunk_token_num: int, overlapped_percent: float): + tnum = num_tokens_from_string(t) + if not pos or tnum < 8: + pos = "" + + # First chunk ever — no previous content to overlap with. + if last_ck == "": + new_t = t + pos if t.find(pos) < 0 else t + final_t = new_t if num_tokens_from_string(new_t) <= chunk_token_num else t + return "first", final_t, num_tokens_from_string(final_t) + + # Proactive merge: append only if the *projected* total still fits. + merged = last_ck + t + merged_pos = merged + pos if last_ck.find(pos) < 0 else merged + if num_tokens_from_string(merged_pos) <= chunk_token_num: + return "merge", merged_pos, num_tokens_from_string(merged_pos) + elif num_tokens_from_string(merged) <= chunk_token_num: + return "merge", merged, num_tokens_from_string(merged) + + # Need a new chunk. Apply overlap prefix from the previous chunk — + # but only when the projected size (overlap + t) fits — otherwise drop + # the overlap for this boundary so the chunk stays within budget. + new_t = t + new_tnum = tnum + if overlapped_percent > 0: + overlap_text, overlap_tokens = _compute_overlap_prefix(last_ck, overlapped_percent) + if overlap_tokens + new_tnum <= chunk_token_num: + new_t = overlap_text + t + new_tnum = num_tokens_from_string(new_t) + if t.find(pos) < 0: + new_t_with_pos = new_t + pos + new_tnum_with_pos = num_tokens_from_string(new_t_with_pos) + if new_tnum_with_pos <= chunk_token_num: + new_t = new_t_with_pos + new_tnum = new_tnum_with_pos + return "append", new_t, new_tnum + + +def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0): if not sections: return [] if isinstance(sections, str): @@ -1169,28 +1295,14 @@ def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。； tk_nums = [0] def add_chunk(t, pos): - nonlocal cks, tk_nums, delimiter - tnum = num_tokens_from_string(t) - if not pos: - pos = "" - if tnum < 8: - pos = "" - # Ensure that the length of the merged chunk does not exceed chunk_token_num - if cks[-1] == "" or tk_nums[-1] > chunk_token_num * (100 - overlapped_percent) / 100.0: - if cks: - overlapped = RAGFlowPdfParser.remove_tag(cks[-1]) - t = overlapped[int(len(overlapped) * (100 - overlapped_percent) / 100.0) :] + t - # Recount with the overlap prefix included, else chunks overshoot chunk_token_num. - tnum = num_tokens_from_string(t) - if t.find(pos) < 0: - t += pos - cks.append(t) - tk_nums.append(tnum) + nonlocal cks, tk_nums + action, text, tk_num = _compute_chunk_update(cks[-1], t, pos, chunk_token_num, overlapped_percent) + if action in ("first", "merge"): + cks[-1] = text + tk_nums[-1] = tk_num else: - if cks[-1].find(pos) < 0: - t += pos - cks[-1] += t - tk_nums[-1] += tnum + cks.append(text) + tk_nums.append(tk_num) custom_delimiters = [m.group(1) for m in re.finditer(r"`([^`]+)`", delimiter)] has_custom = bool(custom_delimiters) @@ -1214,23 +1326,41 @@ def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。； return cks # Split oversized sections at sentence delimiters; add_chunk re-merges to size. + # Units that exceed the budget after the regex split (a single long line with + # no delimiter, e.g. PDF / .txt runs of unbroken text) are sub-split on + # whitespace atoms with a character-window fallback, mirroring the html path. dels = get_delimiters(delimiter) for sec, pos in sections: - if not dels or num_tokens_from_string(sec) < chunk_token_num: - add_chunk("\n" + sec, pos) + sec_text = "\n" + sec + if num_tokens_from_string(sec_text) <= chunk_token_num: + add_chunk(sec_text, pos) continue - for sub_sec in re.split(r"(%s)" % dels, sec, flags=re.DOTALL): - if not sub_sec or re.fullmatch(dels, sub_sec): - continue - add_chunk("\n" + sub_sec, pos) + if dels: + for sub_sec in re.split(r"(%s)" % dels, sec, flags=re.DOTALL): + if not sub_sec or re.fullmatch(dels, sub_sec): + continue + text = "\n" + sub_sec + if num_tokens_from_string(text) <= chunk_token_num: + add_chunk(text, pos) + else: + logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit", len(text), num_tokens_from_string(text)) + for piece in _split_oversized_unit(text, chunk_token_num): + add_chunk(piece, pos) + else: + logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit (no delimiters)", len(sec_text), num_tokens_from_string(sec_text)) + for piece in _split_oversized_unit(sec_text, chunk_token_num): + add_chunk(piece, pos) logging.debug("naive_merge: %d sections -> %d chunks (delimiter=%r)", len(sections), len(cks), delimiter) + # Drop the leading empty placeholder that exists only so ``add_chunk`` could + # detect "first chunk ever" without an extra flag. + if cks and cks[0] == "": + cks = cks[1:] + tk_nums = tk_nums[1:] return cks def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0): - from deepdoc.parser.pdf_parser import RAGFlowPdfParser - if not texts or len(texts) != len(images): return [], [] cks = [""] @@ -1238,33 +1368,23 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 tk_nums = [0] def add_chunk(t, image, pos=""): - nonlocal cks, result_images, tk_nums, delimiter - tnum = num_tokens_from_string(t) - if not pos: - pos = "" - if tnum < 8: - pos = "" - # Ensure that the length of the merged chunk does not exceed chunk_token_num - if cks[-1] == "" or tk_nums[-1] > chunk_token_num * (100 - overlapped_percent) / 100.0: - if cks: - overlapped = RAGFlowPdfParser.remove_tag(cks[-1]) - t = overlapped[int(len(overlapped) * (100 - overlapped_percent) / 100.0) :] + t - # Recount with the overlap prefix included, else chunks overshoot chunk_token_num. - tnum = num_tokens_from_string(t) - if t.find(pos) < 0: - t += pos - cks.append(t) - result_images.append(image) - tk_nums.append(tnum) - else: - if cks[-1].find(pos) < 0: - t += pos - cks[-1] += t + nonlocal cks, result_images, tk_nums + action, text, tk_num = _compute_chunk_update(cks[-1], t, pos, chunk_token_num, overlapped_percent) + if action == "first": + cks[-1] = text + tk_nums[-1] = tk_num + result_images[-1] = image + elif action == "merge": + cks[-1] = text + tk_nums[-1] = tk_num if result_images[-1] is None: result_images[-1] = image else: result_images[-1] = concat_img(result_images[-1], image) - tk_nums[-1] += tnum + else: + cks.append(text) + result_images.append(image) + tk_nums.append(tk_num) custom_delimiters = [m.group(1) for m in re.finditer(r"`([^`]+)`", delimiter)] has_custom = bool(custom_delimiters) @@ -1294,6 +1414,8 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 # Split oversized sections at sentence delimiters; the section's image rides # along on every piece (concat_img dedupes when pieces re-merge into a chunk). + # Units still exceeding the budget after the regex split are sub-split on + # whitespace atoms so they cannot blow past the token cap. dels = get_delimiters(delimiter) for text, image in zip(texts, images): # if text is tuple, unpack it @@ -1303,15 +1425,32 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 else: text_str = text or "" text_pos = "" - if not dels or num_tokens_from_string(text_str) < chunk_token_num: - add_chunk("\n" + text_str, image, text_pos) + + text_seg = "\n" + text_str + if num_tokens_from_string(text_seg) <= chunk_token_num: + add_chunk(text_seg, image, text_pos) continue - for sub_sec in re.split(r"(%s)" % dels, text_str, flags=re.DOTALL): - if not sub_sec or re.fullmatch(dels, sub_sec): - continue - add_chunk("\n" + sub_sec, image, text_pos) + if dels: + for sub_sec in re.split(r"(%s)" % dels, text_str, flags=re.DOTALL): + if not sub_sec or re.fullmatch(dels, sub_sec): + continue + sub_text = "\n" + sub_sec + if num_tokens_from_string(sub_text) <= chunk_token_num: + add_chunk(sub_text, image, text_pos) + else: + logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit", len(sub_text), num_tokens_from_string(sub_text)) + for piece in _split_oversized_unit(sub_text, chunk_token_num): + add_chunk(piece, image, text_pos) + else: + logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit (no delimiters)", len(text_seg), num_tokens_from_string(text_seg)) + for piece in _split_oversized_unit(text_seg, chunk_token_num): + add_chunk(piece, image, text_pos) logging.debug("naive_merge_with_images: %d texts -> %d chunks (delimiter=%r)", len(texts), len(cks), delimiter) + if cks and cks[0] == "": + cks = cks[1:] + result_images = result_images[1:] + tk_nums = tk_nums[1:] return cks, result_images @@ -1334,7 +1473,7 @@ def docx_question_level(p, bull=-1): def concat_img(img1, img2): - from rag.utils.lazy_image import ensure_pil_image, LazyImage + from rag.utils.lazy_image import LazyImage, ensure_pil_image # Same image must not stack with itself (the LazyImage branch would otherwise # concatenate its blob list); mirrors the PIL branch's same-reference guard. @@ -1603,7 +1742,6 @@ def naive_merge_docx( table_context_size=0, image_context_size=0, ): - if not sections: return [], [] diff --git a/test/unit_test/deepdoc/parser/test_txt_parser.py b/test/unit_test/deepdoc/parser/test_txt_parser.py new file mode 100644 index 0000000000..289f77ac8e --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_txt_parser.py @@ -0,0 +1,178 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Unit tests for ``RAGFlowTxtParser.parser_txt`` strict-cap behaviour. + +The pre-fix ``add_chunk`` fired its size check *after* the append, so each +chunk could overshoot ``chunk_token_num`` by up to the size of one line. These +tests assert the proactive projected-total invariant: no produced chunk may +contain more than ``chunk_token_num`` tokens. +""" + +import importlib.util +import os +import sys +from unittest import mock + +_MOCK_MODULES = [ + "xgboost", + "pdfplumber", + "huggingface_hub", + "PIL", + "PIL.Image", + "pypdf", + "sklearn", + "deepdoc.vision", + "deepdoc", + "deepdoc.parser", + "deepdoc.parser.utils", +] +_orig_modules = {m: sys.modules.get(m) for m in _MOCK_MODULES} +_orig_get_text = getattr(sys.modules.get("deepdoc.parser.utils"), "get_text", None) + +try: + for _m in _MOCK_MODULES: + if _m not in sys.modules: + sys.modules[_m] = mock.MagicMock() + + # ``get_text`` is invoked by ``RAGFlowTxtParser.__call__`` only, not by + # ``parser_txt``. Provide a permissive stub so the module loads. + sys.modules["deepdoc.parser.utils"].get_text = lambda *a, **kw: "" + + def _find_project_root(marker="pyproject.toml"): + d = os.path.dirname(os.path.abspath(__file__)) + while d != os.path.dirname(d): + if os.path.exists(os.path.join(d, marker)): + return d + d = os.path.dirname(d) + return None + + _PROJECT_ROOT = _find_project_root() + + _spec = importlib.util.spec_from_file_location( + "deepdoc.parser._txt_parser_under_test", + os.path.join(_PROJECT_ROOT, "deepdoc", "parser", "txt_parser.py"), + ) + _mod = importlib.util.module_from_spec(_spec) + sys.modules["deepdoc.parser._txt_parser_under_test"] = _mod + _spec.loader.exec_module(_mod) + + RAGFlowTxtParser = _mod.RAGFlowTxtParser +finally: + for _m, _orig in _orig_modules.items(): + if _orig is None: + sys.modules.pop(_m, None) + else: + sys.modules[_m] = _orig + if _orig_modules.get("deepdoc.parser.utils") is not None and _orig_get_text is not None: + _orig_modules["deepdoc.parser.utils"].get_text = _orig_get_text + if _orig_modules.get("deepdoc.parser") is not None and _orig_modules.get("deepdoc.parser.utils") is not None: + _orig_modules["deepdoc.parser"].utils = _orig_modules["deepdoc.parser.utils"] + + +# A deterministic, tokenizer-free stand-in for ``num_tokens_from_string`` so +# the assertions below reason in plain words and are independent of tiktoken. + + +def _patch_word_count(monkeypatch_module): + def fake_num_tokens(s): + return len((s or "").split()) + + monkeypatch_module.setattr(_mod, "num_tokens_from_string", fake_num_tokens) + + +def test_no_overshoot_when_packing_short_lines(monkeypatch): + """Lines of 25 tokens, budget 100 — every chunk must be <= 100 tokens.""" + _patch_word_count(monkeypatch) + txt = " ".join(["alpha"] * 25) + "\n" + " ".join(["beta"] * 25) + "\n" + " ".join(["gamma"] * 25) + chunks = RAGFlowTxtParser.parser_txt(txt, chunk_token_num=100, delimiter="\n") + sizes = [len(c[0].split()) for c in chunks if c[0].strip()] + assert all(s <= 100 for s in sizes), sizes + # 75 tokens of content, expected a single 75-token chunk. + assert sum(sizes) == 75 + + +def test_no_overshoot_at_chunk_boundary(monkeypatch): + """Lines of 30 tokens, budget 100. Pre-fix the boundary chunk was 130 tokens.""" + _patch_word_count(monkeypatch) + lines = [" ".join([f"w{i}"] * 30) for i in range(10)] # 10 lines, 300 tokens + chunks = RAGFlowTxtParser.parser_txt("\n".join(lines), chunk_token_num=100, delimiter="\n") + sizes = [len(c[0].split()) for c in chunks if c[0].strip()] + assert all(s <= 100 for s in sizes), sizes + + +def test_atomic_oversized_line_is_sub_split_on_whitespace(monkeypatch): + """A single line that exceeds the budget is split on whitespace atoms.""" + _patch_word_count(monkeypatch) + huge_line = " ".join(["alpha"] * 80) # 80 tokens, no internal delimiter + chunks = RAGFlowTxtParser.parser_txt(huge_line, chunk_token_num=50, delimiter="\n") + sizes = [len(c[0].split()) for c in chunks if c[0].strip()] + assert all(s <= 50 for s in sizes), sizes + assert sum(sizes) == 80 + assert len(chunks) >= 2 + + +def test_empty_text_returns_empty(monkeypatch): + _patch_word_count(monkeypatch) + # Empty input produces a single empty chunk placeholder (existing + # behaviour the callers rely on). The hard-cap guarantee is that any + # chunk carrying content stays within the budget. + result = RAGFlowTxtParser.parser_txt("", chunk_token_num=128, delimiter="\n") + non_empty = [c for c in result if c[0].strip()] + assert non_empty == [] + result2 = RAGFlowTxtParser.parser_txt(" \n\n ", chunk_token_num=128, delimiter="\n") + non_empty2 = [c for c in result2 if c[0].strip()] + assert non_empty2 == [] + + +def test_unbroken_token_exceeding_budget_fallback(monkeypatch): + """A single unbroken non-whitespace string exceeding the budget is split + via the character-window/token-slicing fallback. + """ + + def char_count_tokens(s): + return len(s or "") + + monkeypatch.setattr(_mod, "num_tokens_from_string", char_count_tokens) + + huge_word = "a" * 80 # 80 characters/tokens, no whitespace + chunks = RAGFlowTxtParser.parser_txt(huge_word, chunk_token_num=30, delimiter="\n") + + non_empty = [c[0] for c in chunks if c[0].strip()] + assert len(non_empty) >= 3 + assert all(char_count_tokens(c) <= 30 for c in non_empty) + assert "".join(non_empty) == huge_word + + +def test_newline_join_token_count_strict_cap(monkeypatch): + """Verify that joining chunks with newline does not overshoot chunk_token_num + even when individual token counts sum to <= budget but the newline pushes it over. + """ + + def char_count_tokens(s): + return len(s or "") + + monkeypatch.setattr(_mod, "num_tokens_from_string", char_count_tokens) + + # Two lines of 10 chars each. Budget = 20. + # line1 + "\n" + line2 = 10 + 1 + 10 = 21 chars/tokens, exceeding budget of 20. + line1 = "a" * 10 + line2 = "b" * 10 + txt = f"{line1}\n{line2}" + chunks = RAGFlowTxtParser.parser_txt(txt, chunk_token_num=20, delimiter="\n") + non_empty = [c[0] for c in chunks if c[0].strip()] + assert all(char_count_tokens(c) <= 20 for c in non_empty) + assert len(non_empty) == 2 diff --git a/test/unit_test/rag/test_delimiter_case_sensitive.py b/test/unit_test/rag/test_delimiter_case_sensitive.py index 66c48553d7..6715293b94 100644 --- a/test/unit_test/rag/test_delimiter_case_sensitive.py +++ b/test/unit_test/rag/test_delimiter_case_sensitive.py @@ -119,7 +119,7 @@ def force_every_section_above_budget(monkeypatch): chunk-size heuristics.""" def fake(_s): - return 10**9 + return 9 if len(_s) >= 4 else 8 monkeypatch.setattr(nlp, "num_tokens_from_string", fake) diff --git a/test/unit_test/rag/test_naive_merge.py b/test/unit_test/rag/test_naive_merge.py index 6ec4abb57f..c61e78024d 100644 --- a/test/unit_test/rag/test_naive_merge.py +++ b/test/unit_test/rag/test_naive_merge.py @@ -16,17 +16,22 @@ """Regression tests for ``naive_merge`` / ``naive_merge_with_images``. -Guards against the regression introduced by commit db0f6840d (#11434) where the -default (non-custom-delimiter) path stopped splitting oversized sections at -sentence boundaries, and the overlap prefix was not counted toward a chunk's -token budget. +Guards against: + +* the regression introduced by commit db0f6840d (#11434) where the default + (non-custom-delimiter) path stopped splitting oversized sections at sentence + boundaries, and the overlap prefix was not counted toward a chunk's token + budget; +* the soft-cap bug where chunks systematically overshot ``chunk_token_num`` by + up to one unit (sentence / line) because the size check fired *after* the + append instead of using a projected-total check. """ import re import pytest -import rag.nlp as nlp +from rag import nlp from rag.nlp import naive_merge, naive_merge_with_images DEFAULT_DELIMITER = "\n!?。；！？" @@ -72,8 +77,10 @@ def test_oversized_section_is_split_at_sentence_boundaries(): # Pre-regression behaviour: the section is broken into several chunks # instead of a single oversized one. assert len(chunks) > 1 - # No chunk should greatly exceed the budget (allow one trailing sentence of slack). - assert all(_tok(c) <= 50 + 10 for c in chunks) + # Hard cap: no chunk may exceed the budget. ``<=`` is exact; the slack + # previously allowed (one trailing sentence) is no longer permitted because + # the projected-total check fires before the append. + assert all(_tok(c) <= 50 for c in chunks) # Content is preserved. assert "".join(chunks).count("word") == 200 @@ -107,17 +114,19 @@ def test_empty_delimiter_falls_back_to_token_size_merge(): @pytest.mark.p2 def test_overlap_prefix_is_counted_in_token_budget(): - # With overlap, each chunk = overlap-prefix + new content. The fix recomputes - # the chunk's token count after prepending the prefix, so chunks stay bounded. - # Pre-fix, the prefix tokens were not counted, so the per-chunk budget check - # fired late and chunks systematically overshot chunk_token_num. + # With overlap, each chunk = overlap-prefix + new content. The proactive + # projected-total check rejects a section that, even after prepending the + # overlap prefix, would exceed chunk_token_num; the overlap is dropped at + # that boundary instead of letting the chunk overshoot. Pre-fix, the prefix + # tokens were not counted, so the per-chunk budget check fired late and + # chunks systematically overshot chunk_token_num (observed up to 63). sentences = [" ".join(["w"] * 10) for _ in range(30)] chunks = _nonempty(naive_merge(sentences, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, overlapped_percent=20)) assert len(chunks) > 1 - # Each 10-token sentence divides chunk_token_num evenly, so a correct - # accounting yields chunks of exactly the budget. The buggy version - # overshot (observed up to 63). A small tolerance guards tokenizer rounding. - assert all(_tok(c) <= 50 + 2 for c in chunks) + # Each chunk stays within the budget. Sentences are 10 tokens, the budget + # is 50, so even a 10-token overlap prefix (20% of 50) fits a 40-token + # remainder and the projected-total guarantee holds exactly. + assert all(_tok(c) <= 50 for c in chunks) # --------------------------------------------------------------------------- # @@ -159,7 +168,7 @@ def test_images_oversized_section_is_split(): assert len(nonempty) > 1 # Returned lists stay aligned. assert len(chunks) == len(imgs) - assert all(_tok(c) <= 50 + 10 for c in nonempty) + assert all(_tok(c) <= 50 for c in nonempty) @pytest.mark.p2 @@ -215,3 +224,101 @@ def test_images_distinct_lazyimages_are_concatenated(): merged = nonempty_imgs[0] assert isinstance(merged, LazyImage) assert merged._blobs == [b"BLOB_A", b"BLOB_B"] + + +# --------------------------------------------------------------------------- # +# Hard cap on chunk size (overshoot bug fix) +# --------------------------------------------------------------------------- # + + +@pytest.mark.p2 +def test_strict_cap_no_overlap_packs_to_budget(): + sections = [" ".join(["w"] * 25) for _ in range(8)] + chunks = _nonempty(naive_merge(sections, chunk_token_num=50, delimiter=DEFAULT_DELIMITER)) + assert len(chunks) >= 3 + assert all(_tok(c) <= 50 for c in chunks) + + +@pytest.mark.p2 +def test_strict_cap_with_overlap_drops_overlap_at_overflow_boundary(): + sentences = [" ".join(["w"] * 10) for _ in range(20)] + chunks = _nonempty(naive_merge(sentences, chunk_token_num=25, delimiter=DEFAULT_DELIMITER, overlapped_percent=20)) + assert all(_tok(c) <= 25 for c in chunks) + + +@pytest.mark.p2 +def test_strict_cap_single_overlong_section_is_sub_split_on_whitespace(monkeypatch): + # Override tokenizer in nlp to treat characters as tokens for testing character fallback + def char_count_tokens(s): + return len(s or "") + + monkeypatch.setattr(nlp, "num_tokens_from_string", char_count_tokens) + + big_section = "a" * 80 # unbroken, token-dense string + chunks = _nonempty(naive_merge([big_section], chunk_token_num=50, delimiter=DEFAULT_DELIMITER)) + assert len(chunks) >= 2 + assert all(char_count_tokens(c) <= 50 for c in chunks) + assert "".join(chunks) == big_section + + +@pytest.mark.p2 +def test_strict_cap_overlap_chosen_when_it_fits(): + sentences = [" ".join(["w"] * 5) for _ in range(20)] + chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20)) + assert all(_tok(c) <= 20 for c in chunks) + overlap_seen = False + for a, b in zip(chunks, chunks[1:]): + a_tokens = a.split() + b_tokens = b.split() + if a_tokens and b_tokens and any(t in b_tokens for t in a_tokens): + overlap_seen = True + break + assert overlap_seen + + +@pytest.mark.p2 +def test_images_strict_cap_packs_to_budget(): + sections = [" ".join(["w"] * 25) for _ in range(6)] + images = [None] * len(sections) + chunks, imgs = naive_merge_with_images(sections, images, chunk_token_num=50, delimiter=DEFAULT_DELIMITER) + nonempty = _nonempty(chunks) + assert all(_tok(c) <= 50 for c in nonempty) + assert len(chunks) == len(imgs) + + +@pytest.mark.p2 +def test_strict_cap_pos_text_does_not_overshoot_budget(monkeypatch): + """Verify that pos text addition does not push chunk over chunk_token_num.""" + + def char_count_tokens(s): + return len(s or "") + + monkeypatch.setattr(nlp, "num_tokens_from_string", char_count_tokens) + + # section is 15 chars, pos is 10 chars. chunk_token_num is 20. + # section + pos = 25 > 20, so pos should be omitted or chunk kept <= 20. + pos_tag = "@@12345678" + sections = [("\na" * 15, pos_tag)] + chunks = _nonempty(naive_merge(sections, chunk_token_num=20, delimiter=DEFAULT_DELIMITER)) + assert all(char_count_tokens(c) <= 20 for c in chunks) + + +@pytest.mark.p2 +def test_empty_delimiter_oversized_section_strictly_capped(): + # When delimiter="" and a section exceeds chunk_token_num, it must be sub-split + # so no chunk exceeds chunk_token_num. + long_section = "word " * 100 # ~100 tokens + chunks = _nonempty(naive_merge([long_section], chunk_token_num=30, delimiter="")) + assert len(chunks) > 1 + assert all(_tok(c) <= 30 for c in chunks) + + +@pytest.mark.p2 +def test_images_empty_delimiter_oversized_section_strictly_capped(): + long_section = "word " * 100 + images = [None] + chunks, imgs = naive_merge_with_images([long_section], images, chunk_token_num=30, delimiter="") + nonempty = _nonempty(chunks) + assert len(nonempty) > 1 + assert all(_tok(c) <= 30 for c in nonempty) + assert len(chunks) == len(imgs)