fix(paper): bound title-pivot merge by chunk_token_num (#12109) (#16959)

rag/app/paper.py

@@ -20,6 +20,7 @@ import re
 
 from deepdoc.parser.figure_parser import vision_figure_parser_pdf_wrapper
 from common.constants import ParserType, MAXIMUM_PAGE_NUMBER
+from common.token_utils import num_tokens_from_string
 from rag.nlp import rag_tokenizer, tokenize, tokenize_table, add_positions, bullets_category, title_frequency, tokenize_chunks, attach_media_context, DEFAULT_DELIMITER
 from deepdoc.parser import PdfParser
 import numpy as np
@@ -132,6 +133,26 @@ class Pdf(PdfParser):
         }
 
 
+def _merge_sections_by_pivot(sorted_sections, sec_ids, chunk_token_num):
+    """Concatenate consecutive sections sharing a title pivot, bounded by
+    ``chunk_token_num`` so a long section is not emitted as one oversized chunk.
+    A non-positive budget disables the cap.
+    """
+    chunks = []
+    last_sid = -2
+    for (txt, _), sec_id in zip(sorted_sections, sec_ids, strict=True):
+        # Count body tokens only: the @@page\tx0\t...## position tags are metadata
+        # for pdf_parser.crop(), not content, so strip them before measuring the
+        # budget (the tags stay in chunks[-1] for cropping).
+        merged_body = PdfParser.remove_tag(chunks[-1] + "\n" + txt) if chunks else ""
+        if sec_id == last_sid and chunks and (chunk_token_num <= 0 or num_tokens_from_string(merged_body) <= chunk_token_num):
+            chunks[-1] += "\n" + txt
+            continue
+        chunks.append(txt)
+        last_sid = sec_id
+    return chunks
+
+
 def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang="Chinese", callback=None, **kwargs):
     """
     Only pdf is supported.
@@ -225,15 +246,10 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
         sec_ids.append(sid)
         logging.debug("{} {} {} {}".format(lvl, sorted_sections[i][0], most_level, sid))
 
-    chunks = []
-    last_sid = -2
-    for (txt, _), sec_id in zip(sorted_sections, sec_ids):
-        if sec_id == last_sid:
-            if chunks:
-                chunks[-1] += "\n" + txt
-                continue
-        chunks.append(txt)
-        last_sid = sec_id
+    # Concatenate sections sharing a title pivot, bounded by chunk_token_num so a
+    # long section is not emitted as one oversized chunk (closes #12109).
+    chunk_token_num = int(parser_config.get("chunk_token_num", 512) or 0)
+    chunks = _merge_sections_by_pivot(sorted_sections, sec_ids, chunk_token_num)
     res.extend(tokenize_chunks(chunks, doc, eng, pdf_parser, language=lang))
     table_ctx = max(0, int(parser_config.get("table_context_size", 0) or 0))
     image_ctx = max(0, int(parser_config.get("image_context_size", 0) or 0))

test/unit_test/rag/test_paper_chunk_size.py

@@ -0,0 +1,121 @@
+#
+#  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.
+#
+
+"""Tests for the paper chunk policy's size bound (rag/app/paper.py).
+
+``_merge_sections_by_pivot`` concatenates sections sharing a title pivot but
+bounds the accumulation at ``chunk_token_num`` so a long section is not emitted
+as one oversized chunk.
+"""
+
+import pytest
+
+import rag.app.paper as paper
+from rag.app.paper import _merge_sections_by_pivot
+
+
+@pytest.fixture(autouse=True)
+def word_count_tokens(monkeypatch):
+    monkeypatch.setattr(paper, "num_tokens_from_string", lambda s: len((s or "").split()))
+
+
+def _sections(n, words_each=4):
+    return [(" ".join(["word"] * words_each), "") for _ in range(n)]
+
+
+@pytest.mark.p2
+def test_long_pivot_is_split_into_bounded_chunks():
+    # One pivot, 10 sections x 4 tokens = 40 tokens.
+    chunks = _merge_sections_by_pivot(_sections(10), [0] * 10, chunk_token_num=12)
+    assert len(chunks) > 1
+    # A section is never split mid-way, so allow one section of slack over the budget.
+    assert all(len(c.split()) <= 12 + 4 for c in chunks)
+    assert sum(c.count("word") for c in chunks) == 40
+
+
+@pytest.mark.p2
+def test_nonpositive_budget_keeps_one_chunk_per_pivot():
+    chunks = _merge_sections_by_pivot(_sections(10), [0] * 10, chunk_token_num=0)
+    assert len(chunks) == 1
+    assert len(chunks[0].split()) == 40
+
+
+@pytest.mark.p2
+def test_distinct_pivots_are_never_merged():
+    sections = [("alpha", ""), ("beta", ""), ("gamma", "")]
+    assert _merge_sections_by_pivot(sections, [0, 1, 2], chunk_token_num=1000) == ["alpha", "beta", "gamma"]
+
+
+@pytest.mark.p2
+def test_position_tags_do_not_count_toward_budget():
+    # Each section is 4 body tokens plus a position tag. The @@...## tag is
+    # cropping metadata, not content, so it must not consume the budget: 3 x 4 = 12
+    # body tokens fit chunk_token_num=12 and stay one chunk. Were the tag counted
+    # (it splits into several whitespace-separated tokens) the pivot would break up.
+    tag = "@@1\t72.0\t523.0\t88.0\t101.0##"
+    sections = [("word word word word" + tag, "") for _ in range(3)]
+    chunks = _merge_sections_by_pivot(sections, [0, 0, 0], chunk_token_num=12)
+    assert len(chunks) == 1
+    # The tags stay in the chunk so pdf_parser.crop() can still map page + bbox.
+    assert chunks[0].count("@@") == 3
+
+
+@pytest.mark.p2
+def test_chunk_bounds_oversized_pivot(monkeypatch):
+    # End-to-end guard through paper.chunk(): a single title pivot spanning many
+    # sections must NOT collapse into one oversized chunk. Drives the real merge
+    # path; red against the pre-fix unbounded loop (one 24-token chunk), green
+    # once chunk_token_num caps it. Everything around the merge is stubbed so the
+    # test needs no PDF / model.
+    sections = _sections(6)  # 6 sections x 4 tokens, one pivot => 24 tokens
+
+    class FakePdf:
+        def __call__(self, *a, **k):
+            return {"title": "t", "authors": " ", "abstract": "", "sections": sections, "tables": []}
+
+        def remove_tag(self, s):
+            return s
+
+    captured = {}
+
+    def fake_tokenize_chunks(chunks, *a, **k):
+        captured["chunks"] = chunks
+        return []
+
+    class FakeTok:
+        def tokenize(self, s):
+            return ""
+
+        def fine_grained_tokenize(self, s):
+            return ""
+
+    monkeypatch.setattr(paper, "normalize_layout_recognizer", lambda x: ("DeepDOC", None))
+    monkeypatch.setattr(paper, "Pdf", FakePdf)
+    monkeypatch.setattr(paper, "vision_figure_parser_pdf_wrapper", lambda tbls, **k: tbls)
+    monkeypatch.setattr(paper, "bullets_category", lambda *_: 1)
+    # One pivot: every section shares a level that never advances the pivot id.
+    monkeypatch.setattr(paper, "title_frequency", lambda bull, secs: (1, [2] * len(secs)))
+    monkeypatch.setattr(paper, "tokenize_table", lambda *a, **k: [])
+    monkeypatch.setattr(paper, "tokenize_chunks", fake_tokenize_chunks)
+    monkeypatch.setattr(paper, "rag_tokenizer", FakeTok())
+
+    paper.chunk("doc.pdf", binary=b"x", callback=lambda *a, **k: None, parser_config={"chunk_token_num": 12, "layout_recognize": "DeepDOC"})
+
+    chunks = captured["chunks"]
+    assert len(chunks) > 1, f"oversized pivot collapsed into {len(chunks)} chunk(s)"
+    # A section is never split mid-way, so allow one section of slack over budget.
+    assert all(len(c.split()) <= 12 + 4 for c in chunks)
+    assert sum(c.count("word") for c in chunks) == 24  # nothing lost