fix(chunker): enforce strict chunk_token_num cap on .txt / PDF / email paths (#17203)

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 <skbs-eng@users.noreply.github.com>
Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
This commit is contained in:
S
2026-08-01 20:18:51 +05:30
committed by GitHub
parent f621b4c7b4
commit deb3d0c201
8 changed files with 1108 additions and 246 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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)