diff --git a/docling_core/transforms/chunker/hybrid_chunker.py b/docling_core/transforms/chunker/hybrid_chunker.py index c881d3dbb..f8f9f8f1b 100644 --- a/docling_core/transforms/chunker/hybrid_chunker.py +++ b/docling_core/transforms/chunker/hybrid_chunker.py @@ -243,20 +243,23 @@ def _split_using_plain_text( # captions: available_length = self.max_tokens - lengths.other_len - if available_length <= 0: - warnings.warn( - "Headers and captions for this chunk are longer than the total " - "available size for the chunk, so they will be ignored: " - f"{doc_chunk.text=}, {doc_chunk.meta=}" - ) - new_chunk = DocChunk(**doc_chunk.export_json_dict()) - new_chunk.meta.captions = None - new_chunk.meta.headings = None - return self._split_using_plain_text(doc_chunk=new_chunk, doc_serializer=doc_serializer) - - segments = self.segment(doc_chunk, available_length, doc_serializer) - chunks = [DocChunk(text=s, meta=doc_chunk.meta) for s in segments] - return chunks + while available_length > 0: + segments = self.segment(doc_chunk, available_length, doc_serializer) + chunks = [DocChunk(text=s, meta=doc_chunk.meta) for s in segments] + overflow = max(self._count_chunk_tokens(chunk) - self.max_tokens for chunk in chunks) + if overflow <= 0: + return chunks + available_length -= overflow + + warnings.warn( + "Headers and captions for this chunk are longer than the total " + "available size for the chunk, so they will be ignored: " + f"{doc_chunk.text=}, {doc_chunk.meta=}" + ) + new_chunk = DocChunk(**doc_chunk.export_json_dict()) + new_chunk.meta.captions = None + new_chunk.meta.headings = None + return self._split_using_plain_text(doc_chunk=new_chunk, doc_serializer=doc_serializer) def segment(self, doc_chunk: DocChunk, available_length: int, doc_serializer: BaseDocSerializer) -> list[str]: """Split a single doc chunk into a list of text segments. @@ -276,9 +279,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. @@ -306,13 +306,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: if not _SEMCHUNK_AVAILABLE: raise ImportError(_SEMCHUNK_INSTALL_HINT) from _SEMCHUNK_IMPORT_ERROR diff --git a/docling_core/transforms/chunker/line_chunker.py b/docling_core/transforms/chunker/line_chunker.py index 26a970057..0ef98e5d9 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 @@ -167,9 +169,8 @@ def chunk_text(self, lines: list[str]) -> list[str]: # Check if first line would overflow with prefix when omit_prefix_on_overflow=True # If yes, add prefix as a standalone chunk first to ensure it's visible if self.omit_prefix_on_overflow and self.prefix_len > 0 and lines: - first_line_tokens = self.tokenizer.count_tokens(lines[0]) # If first line would overflow with prefix, add prefix as standalone chunk - if first_line_tokens + self.prefix_len > self.max_tokens: + if self.tokenizer.count_tokens(self.prefix + lines[0]) > self.max_tokens: chunks.append(self.prefix) current = "" current_len = 0 @@ -190,22 +191,20 @@ 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: - chunks.append(current) + fresh_prefix = self.prefix if self.prefix_len > 0 else "" + if self.tokenizer.count_tokens(fresh_prefix + remaining) <= self.max_tokens: + if current: + chunks.append(current) # Only add prefix to new chunks if it fits (prefix_len > 0) - if self.prefix_len > 0: - current = self.prefix - current_len = self.prefix_len - else: - current = "" - current_len = 0 + current = fresh_prefix + current_len = self.prefix_len # loop continues to retry fitting `remaining` continue @@ -236,19 +235,32 @@ 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) - # Zero-progress detection: if take is empty, force character-level split + 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 + + if not take: + take, remaining = self._split_by_token_limit( + remaining, + token_limit=self.max_tokens, + prefer_word_boundary=True, + prefix=current, + ) + if not take: - # Fallback: take at least one character to ensure progress - if remaining: - take = remaining[0] - remaining = remaining[1:] - else: - # Should not happen, but break to prevent infinite loop - break + if current: + chunks.append(current) + current = "" + current_len = 0 + continue + raise ValueError( + f"Character {remaining[0]!r} cannot fit within the token budget of {self.max_tokens}" + ) # 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) @@ -301,11 +313,25 @@ def split_by_token_limit( Returns: (head, tail) where `head` contains at most `token_limit` tokens, `tail` is the remaining suffix. If `token_limit <= 0`, returns ("", text). """ + return self._split_by_token_limit( + text, + token_limit=token_limit, + prefer_word_boundary=prefer_word_boundary, + prefix="", + ) + + def _split_by_token_limit( + self, + text: str, + token_limit: int, + prefer_word_boundary: bool, + prefix: str, + ) -> tuple[str, str]: if token_limit <= 0 or not text: return "", text # if the whole text already fits, return as is. - if self.tokenizer.count_tokens(text) <= token_limit: + if self.tokenizer.count_tokens(prefix + text) <= token_limit: return text, "" # Binary search over character indices [0, len(text)] @@ -315,7 +341,7 @@ def split_by_token_limit( while lo <= hi: mid = (lo + hi) // 2 head = text[:mid] - tok_count = self.tokenizer.count_tokens(head) + tok_count = self.tokenizer.count_tokens(prefix + head) if tok_count <= token_limit: best_idx = mid # feasible; try to extend @@ -324,16 +350,23 @@ def split_by_token_limit( hi = mid - 1 if best_idx is None or best_idx <= 0: - # Even the first character exceeds the limit (e.g., tokenizer behavior). - # Return nothing in head, everything in tail. - return "", text + best_idx = next( + ( + index + for index in range(len(text) - 1, 0, -1) + if self.tokenizer.count_tokens(prefix + text[:index]) <= token_limit + ), + None, + ) + if best_idx is None: + return "", text # Optionally adjust to a previous whitespace boundary without violating the limit if prefer_word_boundary: # Search backwards from best_idx to find whitespace; keep within token limit. # Only snap back if it produces a non-empty head (last_space_index > 0) last_space_index = text[:best_idx].rfind(" ") - if last_space_index > 0: + if last_space_index > 0 and self.tokenizer.count_tokens(prefix + text[:last_space_index]) <= token_limit: best_idx = last_space_index head, tail = text[:best_idx], text[best_idx:] diff --git a/test/data/chunker/0d_out_chunks.json b/test/data/chunker/0d_out_chunks.json index d214b3747..7fca55ed0 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 7c9b17f12..d09c5fee3 100644 --- a/test/test_hybrid_chunker.py +++ b/test/test_hybrid_chunker.py @@ -19,7 +19,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.items.table.table_data import TableCell, TableData from .test_utils import assert_or_generate_json_ground_truth, build_single_cell_rich_table_doc @@ -59,6 +59,31 @@ 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 # --------------------------------------------------------------------------- @@ -785,3 +810,123 @@ def test_chunk_raises_on_missing_semchunk(monkeypatch): with pytest.raises(ImportError, match="semchunk"): list(chunker.chunk(dl_doc=dl_doc)) + + +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) + + +def test_contextualized_markdown_table_chunks_recheck_final_count(): + max_tokens = 70 + doc, _ = _build_wide_header_table(caption_text="\nA") + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + 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 + + +@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) == expected_text diff --git a/test/test_line_chunker.py b/test/test_line_chunker.py index fedb0b204..e9ad9c162 100644 --- a/test/test_line_chunker.py +++ b/test/test_line_chunker.py @@ -1,8 +1,10 @@ import pytest +import tiktoken from transformers import AutoTokenizer from docling_core.transforms.chunker.line_chunker import LineBasedTokenChunker from docling_core.transforms.chunker.tokenizer.huggingface import HuggingFaceTokenizer +from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer from docling_core.types.doc import DoclingDocument as DLDocument from docling_core.types.doc.labels import DocItemLabel @@ -189,6 +191,130 @@ 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_text_flushes_full_chunk_before_splitting_long_line(default_tokenizer): + chunker = LineBasedTokenChunker( + tokenizer=default_tokenizer, + max_tokens_override=4, + ) + lines = ["aColumn", "0 0 0 0 0"] + + assert chunker.tokenizer.count_tokens(lines[0]) == chunker.max_tokens + assert chunker.tokenizer.count_tokens(lines[1]) > chunker.max_tokens + + chunks = chunker.chunk_text(lines) + + assert "".join(chunks) == "".join(lines) + assert all(chunker.tokenizer.count_tokens(chunk) <= chunker.max_tokens for chunk in chunks) + + +def test_chunk_text_rejects_character_larger_than_budget(): + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + max_tokens=8192, + ) + chunker = LineBasedTokenChunker( + tokenizer=tokenizer, + max_tokens_override=1, + ) + + with pytest.raises(ValueError, match="cannot fit within the token budget"): + chunker.chunk_text(["πŸ™‚"]) + + +def test_chunk_text_makes_progress_after_standalone_prefix(): + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + max_tokens=1, + ) + chunker = LineBasedTokenChunker( + tokenizer=tokenizer, + prefix="a", + ) + + chunks = chunker.chunk_text(["tegory"]) + + assert chunks[0] == "a" + assert "".join(chunks[1:]) == "tegory" + assert "" not in chunks + assert all(chunker.tokenizer.count_tokens(chunk) <= 1 for chunk in chunks) + + +def test_chunk_text_splits_nonmonotone_token_prefixes(): + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + max_tokens=1, + ) + chunker = LineBasedTokenChunker(tokenizer=tokenizer) + + chunks = chunker.chunk_text(["εˆ ι™€a"]) + + assert chunks == ["εˆ ι™€", "a"] + + +def test_split_by_token_limit_rechecks_shorter_word_boundary(): + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + max_tokens=6, + ) + chunker = LineBasedTokenChunker(tokenizer=tokenizer) + text = "\néé\n0\n \n0a0" + + head, tail = chunker.split_by_token_limit(text, token_limit=6) + + assert head + tail == text + assert tokenizer.count_tokens(head) <= 6 + + +def test_chunk_text_retains_prefix_when_exact_split_fits(): + tokenizer = OpenAITokenizer( + tokenizer=tiktoken.get_encoding("cl100k_base"), + max_tokens=3, + ) + chunker = LineBasedTokenChunker( + tokenizer=tokenizer, + prefix="ΒΆ ", + ) + + chunks = chunker.chunk_text(["πŸ™‚πŸ™‚"]) + + assert chunks == ["ΒΆ πŸ™‚", "ΒΆ πŸ™‚"] + + +def test_omit_prefix_uses_concatenated_token_count(default_tokenizer): + chunker = LineBasedTokenChunker( + tokenizer=default_tokenizer, + max_tokens_override=4, + prefix="a", + omit_prefix_on_overflow=True, + ) + line = "Column Column" + + assert chunker.tokenizer.count_tokens(chunker.prefix) + chunker.tokenizer.count_tokens(line) <= 4 + assert chunker.tokenizer.count_tokens(chunker.prefix + line) > 4 + + chunks = chunker.chunk_text([line]) + + assert chunks == ["a", line] + + def test_chunk_document(default_tokenizer): """Test the chunk() method with a DoclingDocument.""" # Create a simple DoclingDocument @@ -272,30 +398,16 @@ def test_chunk_document_with_long_content(default_tokenizer): assert token_count <= MAX_TOKENS -def test_infinite_loop_regression_long_unbreakable_token(default_tokenizer): - """Regression test for infinite loop bug when processing text with a long - unbreakable token sequence preceded by a space. - - This test reproduces the issue where LineBasedTokenChunker.chunk_text() - would enter an infinite loop when the prefer_word_boundary logic in - split_by_token_limit() snapped best_idx back to 0, producing an empty - head and returning the tail unchanged. - - The fix ensures: - 1. split_by_token_limit only snaps to word boundary if it produces non-empty head - 2. chunk_text detects zero-progress and forces character-level splitting as fallback - """ +def test_chunk_text_preserves_long_unbreakable_token(default_tokenizer): chunker = LineBasedTokenChunker( tokenizer=default_tokenizer, ) - # Create text with leading space followed by long unbreakable token long_word = "a" * 200 text = "Header " + long_word + " Footer\n" token_count = chunker.tokenizer.count_tokens(text) assert token_count <= MAX_TOKENS - # This should complete without hanging result = chunker.chunk_text(lines=[text]) assert len(result) == 1 @@ -332,28 +444,14 @@ def test_split_by_token_limit_leading_space_regression(default_tokenizer): assert head_tokens <= 10 -def test_character_level_fallback_on_zero_available(default_tokenizer): - """Test that chunk_text uses character-level fallback when available space is 0. - - This test demonstrates a real scenario where the fallback is needed: - 1. Current chunk is exactly at max_tokens (available = 0) - 2. Remaining text is too long to fit in a fresh chunk (exceeds max_tokens) - 3. split_by_token_limit is called with token_limit=0 - 4. It returns ("", text) in split_by_token_limit of line_chunker.py - 5. The fallback takes 1 character to ensure progress - """ - # Use a very small max_tokens to make it easier to create the scenario +def test_chunk_text_flushes_when_no_space_remains(default_tokenizer): chunker = LineBasedTokenChunker( tokenizer=default_tokenizer, ) - # First line: 8 tokens (leaves room for 2 more) first_line = "word " * (MAX_TOKENS - 2) - - # Second line: more than 10 tokens (so it can't fit in a fresh chunk) second_line = "x" * 100 - # Verify second line is indeed > max_tokens second_line_tokens = chunker.tokenizer.count_tokens(second_line) assert second_line_tokens > MAX_TOKENS, ( f"Second line must exceed max_tokens for test to work: {second_line_tokens} <= {MAX_TOKENS}" @@ -362,19 +460,14 @@ def test_character_level_fallback_on_zero_available(default_tokenizer): lines = [first_line, second_line] result = chunker.chunk_text(lines=lines) - # Verify we got multiple chunks (the long second line should be split) assert len(result) > 1, f"Should have multiple chunks, got {len(result)}" - # Verify each chunk respects token limit (allow small overflow due to newline addition) for i, chunk in enumerate(result): token_count = chunker.tokenizer.count_tokens(chunk) assert token_count <= MAX_TOKENS, f"Chunk {i} exceeds token limit: {token_count} > {MAX_TOKENS}" - # Verify all content is preserved combined = "".join(result) - # First line should be in first chunk assert first_line in result[0], "First line should be in first chunk" - # Second line should be split across remaining chunks combined_without_newlines = combined.replace("\n", "") assert second_line in combined_without_newlines, "Second line should be preserved (possibly split)"