mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-03 14:27:32 +08:00
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:
@@ -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]
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
266
internal/ingestion/component/chunker/token_strict_cap_test.go
Normal file
266
internal/ingestion/component/chunker/token_strict_cap_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 [], []
|
||||
|
||||
|
||||
178
test/unit_test/deepdoc/parser/test_txt_parser.py
Normal file
178
test/unit_test/deepdoc/parser/test_txt_parser.py
Normal 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
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user