From afb092a34d71d6ea43ba0825b04a8163c4eb5d13 Mon Sep 17 00:00:00 2001 From: Pablo Pupo Date: Mon, 20 Jul 2026 14:54:13 -0400 Subject: [PATCH 1/2] fix(chunker): respect contextualized table token limits Signed-off-by: Pablo Pupo --- .../transforms/chunker/hybrid_chunker.py | 6 +- .../transforms/chunker/line_chunker.py | 26 +++- test/data/chunker/0d_out_chunks.json | 28 +++- test/test_hybrid_chunker.py | 129 +++++++++++++++++- 4 files changed, 175 insertions(+), 14 deletions(-) diff --git a/docling_core/transforms/chunker/hybrid_chunker.py b/docling_core/transforms/chunker/hybrid_chunker.py index dfa19799..d1a8cfc4 100644 --- a/docling_core/transforms/chunker/hybrid_chunker.py +++ b/docling_core/transforms/chunker/hybrid_chunker.py @@ -258,9 +258,6 @@ def segment(self, doc_chunk: DocChunk, available_length: int, doc_serializer: Ba Args: doc_chunk: The chunk to segment. - available_length: Maximum token budget for the semantic-splitting - path (ignored for the table path, which uses ``self.tokenizer`` - and ``self.max_tokens`` internally). doc_serializer: Serializer for the current document; must be a ``ChunkingDocSerializer`` for the table path to activate. @@ -288,13 +285,14 @@ def segment(self, doc_chunk: DocChunk, available_length: int, doc_serializer: Ba line_chunker = LineBasedTokenChunker( tokenizer=self.tokenizer, + max_tokens_override=available_length, prefix=full_prefix, omit_prefix_on_overflow=self.omit_header_on_overflow, serializer_provider=self.serializer_provider, ) segments = line_chunker.chunk_text(lines=body_lines) if preamble: - segments = segments[:1] + [s[len(preamble) :] for s in segments[1:]] + segments = segments[:1] + [s[len(preamble) :] if s.startswith(full_prefix) else s for s in segments[1:]] else: sem_chunker = semchunk.chunkerify(self.tokenizer.get_tokenizer(), chunk_size=available_length) sem_segments = sem_chunker(doc_chunk.text) diff --git a/docling_core/transforms/chunker/line_chunker.py b/docling_core/transforms/chunker/line_chunker.py index 26a97005..4eacf173 100644 --- a/docling_core/transforms/chunker/line_chunker.py +++ b/docling_core/transforms/chunker/line_chunker.py @@ -47,6 +47,8 @@ class LineBasedTokenChunker(BaseChunker): ), ] + max_tokens_override: Annotated[int | None, Field(default=None, gt=0)] + prefix: Annotated[ str, Field( @@ -88,6 +90,7 @@ def prefix_chunks(self) -> list[str]: # Split the prefix into chunks using a temporary chunker with no prefix temp_chunker = LineBasedTokenChunker( tokenizer=self.tokenizer, + max_tokens_override=self.max_tokens_override, prefix="", omit_prefix_on_overflow=False, serializer_provider=self.serializer_provider, @@ -113,8 +116,7 @@ def prefix_len(self) -> int: @property def max_tokens(self) -> int: - """Maximum number of tokens allowed in a chunk, as reported by the tokenizer.""" - return self.tokenizer.get_max_tokens() + return self.max_tokens_override or self.tokenizer.get_max_tokens() def model_post_init(self, __context) -> None: # Trigger computation of prefix_chunks to validate prefix length @@ -190,14 +192,14 @@ def chunk_text(self, lines: list[str]) -> list[str]: available = self.max_tokens - current_len # If the remaining part fits entirely into current chunk → append and stop - if line_tokens <= available: + if self.tokenizer.count_tokens(current + remaining) <= self.max_tokens: current += remaining - current_len += line_tokens + current_len = self.tokenizer.count_tokens(current) break # Remaining does NOT fit into current chunk. # If it CAN fit into a fresh chunk → flush current and start new one. - if line_tokens + self.prefix_len <= self.max_tokens: + if self.tokenizer.count_tokens(self.prefix + remaining) <= self.max_tokens: chunks.append(current) # Only add prefix to new chunks if it fits (prefix_len > 0) if self.prefix_len > 0: @@ -236,8 +238,18 @@ def chunk_text(self, lines: list[str]) -> list[str]: # Split off the first segment that fits into current. take, remaining = self.split_by_token_limit(remaining, available) + while take and self.tokenizer.count_tokens(current + take) > self.max_tokens: + take_limit = self.tokenizer.count_tokens(take) - 1 + take, returned = self.split_by_token_limit(take, take_limit) + remaining = returned + remaining + # Zero-progress detection: if take is empty, force character-level split if not take: + if current: + chunks.append(current) + current = "" + current_len = 0 + continue # Fallback: take at least one character to ensure progress if remaining: take = remaining[0] @@ -247,8 +259,8 @@ def chunk_text(self, lines: list[str]) -> list[str]: break # Add the taken part - current += "\n" + take - current_len += self.tokenizer.count_tokens(take) + current += take + current_len = self.tokenizer.count_tokens(current) # flush the current chunk (full) chunks.append(current) diff --git a/test/data/chunker/0d_out_chunks.json b/test/data/chunker/0d_out_chunks.json index d214b374..7fca55ed 100644 --- a/test/data/chunker/0d_out_chunks.json +++ b/test/data/chunker/0d_out_chunks.json @@ -770,7 +770,19 @@ } }, { - "text": "| class label | Count | % of Total | % of Total | % of Total | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| class label | Count | Train | Test | Val | All | Fin | Man | Sci | Law | Pat | Ten |\n| Caption | 22524 | 2.04 | 1.77 | 2.32 | 84-89 | 40-61 | 86-92 | 94-99 | 95-99 | 69-78 | n/a |\n", + "text": "| class label | Count | % of Total | % of Total | % of Total | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| class label | Count | Train | Test | Val | All | Fin | Man | Sci | Law | Pat | Ten |\n", + "meta": { + "doc_items": [ + "#/tables/3" + ], + "headings": [ + "Docling Technical Report", + "Baselines for Object Detection" + ] + } + }, + { + "text": "| class label | Count | % of Total | % of Total | % of Total | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) | triple inter-annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| Caption | 22524 | 2.04 | 1.77 | 2.32 | 84-89 | 40-61 | 86-92 | 94-99 | 95-99 | 69-78 | n/a |\n", "meta": { "doc_items": [ "#/tables/3" @@ -930,7 +942,19 @@ } }, { - "text": "| | | % of Total | % of Total | % of Total | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| class label | Count | Train | Test | Val | All | Fin | Man | Sci | Law | Pat | Ten |\n| Caption | 22524 | 2.04 | 1.77 | 2.32 | 84-89 | 40-61 | 86-92 | 94-99 | 95-99 | 69-78 | n/a |\n", + "text": "| | | % of Total | % of Total | % of Total | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| class label | Count | Train | Test | Val | All | Fin | Man | Sci | Law | Pat | Ten |\n", + "meta": { + "doc_items": [ + "#/tables/4" + ], + "headings": [ + "Docling Technical Report", + "Baselines for Object Detection" + ] + } + }, + { + "text": "| | | % of Total | % of Total | % of Total | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) | triple inter- annotator mAP @ 0.5-0.95 (%) |\n| - | - | - | - | - | - | - | - | - | - | - | - |\n| Caption | 22524 | 2.04 | 1.77 | 2.32 | 84-89 | 40-61 | 86-92 | 94-99 | 95-99 | 69-78 | n/a |\n", "meta": { "doc_items": [ "#/tables/4" diff --git a/test/test_hybrid_chunker.py b/test/test_hybrid_chunker.py index a7fbd2c5..8d4235e8 100644 --- a/test/test_hybrid_chunker.py +++ b/test/test_hybrid_chunker.py @@ -18,7 +18,7 @@ from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer from docling_core.transforms.serializer.html import HTMLTableSerializer from docling_core.transforms.serializer.markdown import MarkdownParams, MarkdownTableSerializer -from docling_core.types.doc import DocItemLabel, DoclingDocument +from docling_core.types.doc import DocItemLabel, DoclingDocument, TableItem from docling_core.types.doc.document import TableCell, TableData from .test_utils import assert_or_generate_json_ground_truth, build_single_cell_rich_table_doc @@ -58,6 +58,33 @@ def get_serializer(self, doc: DoclingDocument): ) +def _build_wide_header_table( + caption_text: str | None = None, +) -> tuple[DoclingDocument, TableItem]: + doc = DoclingDocument(name="wide_table") + doc.add_heading(text="Section heading", level=1) + caption = doc.add_text(label=DocItemLabel.CAPTION, text=caption_text) if caption_text is not None else None + cells = [ + TableCell( + text=f"Column {col} heading" if row == 0 else f"Value {col}", + row_span=1, + col_span=1, + start_row_offset_idx=row, + end_row_offset_idx=row + 1, + start_col_offset_idx=col, + end_col_offset_idx=col + 1, + column_header=row == 0, + ) + for row in range(2) + for col in range(50) + ] + table = doc.add_table( + data=TableData(num_rows=2, num_cols=50, table_cells=cells), + caption=caption, + ) + return doc, table + + # --------------------------------------------------------------------------- # Shared fixtures # --------------------------------------------------------------------------- @@ -742,3 +769,103 @@ def test_html_parser_error_handling(dl_doc_0): for chunk in chunks: assert isinstance(chunk.text, str), "Chunk text should be a string" assert chunk.meta is not None, "Chunk should have metadata" + + +def test_contextualized_markdown_table_chunks_respect_token_limit(): + max_tokens = 64 + doc, _ = _build_wide_header_table() + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.encoding_for_model("text-embedding-3-small"), + max_tokens=max_tokens, + ) + chunker = HybridChunker( + tokenizer=tokenizer, + merge_peers=False, + repeat_table_header=True, + serializer_provider=CompactMarkdownSerializerProvider(), + ) + + chunks = list(chunker.chunk(dl_doc=doc)) + + assert len(chunks) > 1 + assert all(tokenizer.count_tokens(chunker.contextualize(chunk)) <= max_tokens for chunk in chunks) + + +@pytest.mark.parametrize( + ("max_tokens", "cell_prefix", "heading_text"), + [ + (96, "x", None), + (128, "a-b/", "Section heading"), + ], +) +def test_split_markdown_table_chunks_respect_exact_token_boundaries( + max_tokens: int, + cell_prefix: str, + heading_text: str | None, +): + num_cols = 12 + num_rows = 3 + doc = DoclingDocument(name="overflow") + if heading_text is not None: + doc.add_heading(text=heading_text, level=1) + cells = [ + TableCell( + text=f"{cell_prefix}{row}_{col}", + row_span=1, + col_span=1, + start_row_offset_idx=row, + end_row_offset_idx=row + 1, + start_col_offset_idx=col, + end_col_offset_idx=col + 1, + column_header=row == 0, + ) + for row in range(num_rows) + for col in range(num_cols) + ] + doc.add_table( + data=TableData( + num_rows=num_rows, + num_cols=num_cols, + table_cells=cells, + ) + ) + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.encoding_for_model("text-embedding-3-small"), + max_tokens=max_tokens, + ) + chunker = HybridChunker( + tokenizer=tokenizer, + merge_peers=False, + repeat_table_header=True, + serializer_provider=CompactMarkdownSerializerProvider(), + ) + + chunks = list(chunker.chunk(dl_doc=doc)) + counts = [tokenizer.count_tokens(chunker.contextualize(chunk)) for chunk in chunks] + + assert max(counts) <= max_tokens + + +def test_split_markdown_table_prefix_preserves_content(): + doc, table = _build_wide_header_table(caption_text="Wide table caption") + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.encoding_for_model("text-embedding-3-small"), + max_tokens=64, + ) + chunker = HybridChunker( + tokenizer=tokenizer, + merge_peers=False, + repeat_table_header=True, + serializer_provider=CompactMarkdownSerializerProvider(), + ) + serializer = chunker.serializer_provider.get_serializer(doc) + expected_text = serializer.serialize(item=table).text + + chunks = [ + chunk + for chunk in chunker.chunk(dl_doc=doc) + if any(item.self_ref == table.self_ref for item in chunk.meta.doc_items) + ] + + assert len(chunks) > 1 + assert "".join(chunk.text for chunk in chunks).replace("\n", "") == expected_text.replace("\n", "") From c95ce2fb18b7b9fe86fa49decf381040d23dc942 Mon Sep 17 00:00:00 2001 From: Pablo Pupo Date: Mon, 20 Jul 2026 16:07:47 -0400 Subject: [PATCH 2/2] test(chunker): cover concatenated token boundary Signed-off-by: Pablo Pupo --- test/test_line_chunker.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/test_line_chunker.py b/test/test_line_chunker.py index fedb0b20..3b212866 100644 --- a/test/test_line_chunker.py +++ b/test/test_line_chunker.py @@ -189,6 +189,24 @@ def test_chunk_text_with_prefix_and_long_lines(default_tokenizer): assert token_count <= MAX_TOKENS +def test_chunk_text_rechecks_concatenated_token_count(default_tokenizer): + chunker = LineBasedTokenChunker( + tokenizer=default_tokenizer, + max_tokens_override=4, + prefix="a", + ) + text = "Column Column" + + assert chunker.tokenizer.count_tokens(chunker.prefix) + chunker.tokenizer.count_tokens(text) <= chunker.max_tokens + assert chunker.tokenizer.count_tokens(chunker.prefix + text) > chunker.max_tokens + + chunks = chunker.chunk_text([text]) + + assert chunks == ["aColumn", "a Column"] + assert all(chunker.tokenizer.count_tokens(chunk) <= chunker.max_tokens for chunk in chunks) + assert "".join(chunk.removeprefix(chunker.prefix) for chunk in chunks) == text + + def test_chunk_document(default_tokenizer): """Test the chunk() method with a DoclingDocument.""" # Create a simple DoclingDocument