Coverage for src/qdrant_loader/core/chunking/docling/docling_chunker.py: 89%
47 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""The docling chunker — the one module that touches ``HybridChunker`` directly.
3:class:`DoclingChunker` is a thin facade over a single docling ``HybridChunker``.
4It owns construction *once* via ``cached_property`` (class-enforced lazy-once, not a
5lazy-null singleton) and composes the pure :class:`~.tokenizer.TokenizerFactory` and
6:class:`~.structure.StructureProjector` rather than inlining their logic.
8It structurally satisfies :class:`~.chunker.DocumentChunker` — no base class, no
9registration. It runs HybridChunker's two-pass split/merge, maps ``contextualize``
10output to ``embed_text``, and surfaces the post-context budget check via
11``_warn_if_over_budget``.
12"""
14from __future__ import annotations
16from functools import cached_property
17from typing import TYPE_CHECKING
19from qdrant_loader.utils.logging import LoggingConfig
21from .config import ChunkingConfig, TableSerialization
22from .exceptions import EmptyDocumentError
23from .outcome import Chunk
24from .structure import StructureProjector
25from .tokenizer import TokenizerFactory
27if TYPE_CHECKING:
28 from docling_core.transforms.chunker import BaseChunk, HybridChunker
30 from ...conversion import ConvertedDocument
32logger = LoggingConfig.get_logger(__name__)
35class DoclingChunker:
36 """ConvertedDocument -> list[Chunk] via docling's HybridChunker."""
38 def __init__(self, config: ChunkingConfig) -> None:
39 self._config = config
40 self._tokenizer_factory = TokenizerFactory(config.tokenizer)
41 self._projector = StructureProjector()
43 @cached_property
44 def _chunker(self) -> HybridChunker:
45 """Construct the docling ``HybridChunker`` once, then reuse it.
47 The tokenizer (built from our config) carries the token budget; ``merge_peers``
48 and ``delim`` come from the config, and ``table_serialization`` selects the
49 serializer provider. Built lazily so the budget/tokenizer cost is paid on
50 first chunk, not at import.
51 """
52 from docling_core.transforms.chunker import HybridChunker
54 return HybridChunker(
55 tokenizer=self._tokenizer_factory.build(),
56 merge_peers=self._config.merge_peers,
57 delim=self._config.delim,
58 serializer_provider=self._serializer_provider(),
59 )
61 def warm_up(self) -> None:
62 """Build the docling ``HybridChunker`` + tokenizer eagerly.
64 Used at pipeline startup so the first document's chunking timeout never
65 covers the tokenizer build (HuggingFaceTokenizer.from_pretrained / tiktoken
66 load). Touching the ``_chunker`` cached_property triggers and caches the
67 one-time construction — mirroring :meth:`GlinerEntityExtractor.warm_up`.
68 """
69 _ = self._chunker
71 def _serializer_provider(self):
72 """Map the table-serialization knob onto a docling serializer provider.
74 TRIPLETS is docling's own chunking default; MARKDOWN swaps in the
75 GitHub-table serializer while keeping the rest of the chunking
76 serialization (picture descriptions, code fences) unchanged.
77 """
78 from docling_core.transforms.chunker.hierarchical_chunker import (
79 ChunkingDocSerializer,
80 ChunkingSerializerProvider,
81 )
83 if self._config.table_serialization is not TableSerialization.MARKDOWN:
84 return ChunkingSerializerProvider()
86 from docling_core.transforms.serializer.markdown import MarkdownTableSerializer
88 class _MarkdownTableProvider(ChunkingSerializerProvider):
89 def get_serializer(self, doc):
90 return ChunkingDocSerializer(
91 doc=doc, table_serializer=MarkdownTableSerializer()
92 )
94 return _MarkdownTableProvider()
96 def chunk(self, document: ConvertedDocument) -> list[Chunk]:
97 """Chunk a converted document into engine-neutral chunks.
99 Runs ``HybridChunker.chunk`` over the structured ``DoclingDocument``, projects
100 each chunk's ``meta`` via :class:`StructureProjector`, and counts tokens on the
101 embedded text. A contentless document fails loud with
102 :class:`EmptyDocumentError` rather than emitting a fake empty chunk.
104 ``embed_text`` mirrors the body while contextual embedding is disabled (the
105 deferred toggle); once enabled it becomes ``contextualize`` output and the
106 ``enforce_token_budget`` check (which HybridChunker cannot guarantee for
107 prepended context) becomes reachable.
108 """
109 chunker = self._chunker
110 chunks = []
111 for doc_chunk in chunker.chunk(document.document):
112 chunk = self._to_chunk(doc_chunk, chunker)
113 self._warn_if_over_budget(chunk)
114 chunks.append(chunk)
115 if not chunks:
116 raise EmptyDocumentError(
117 f"converted {document.source_format!r} document produced no chunks"
118 )
119 return chunks
121 def _warn_if_over_budget(self, chunk: Chunk) -> None:
122 """Surface — but do not raise on — a chunk that exceeds the token budget.
124 HybridChunker only guarantees the budget for the bare body; ``contextualize``
125 (when ``include_context_in_embed`` is on) and atomic oversized doc items can
126 push a chunk over. The overflow is allowed; when
127 ``enforce_token_budget`` is set we log it, because the embedder may silently
128 truncate or reject oversized input.
129 """
130 max_tokens = self._config.tokenizer.max_tokens
131 if (
132 self._config.enforce_token_budget
133 and chunk.token_count is not None
134 and chunk.token_count > max_tokens
135 ):
136 logger.warning(
137 "docling chunk exceeds token budget; embedder may truncate it",
138 token_count=chunk.token_count,
139 max_tokens=max_tokens,
140 heading_path=list(chunk.structure.heading_path),
141 )
143 def _to_chunk(self, doc_chunk: BaseChunk, chunker: HybridChunker) -> Chunk:
144 """Project one docling chunk into our engine-neutral :class:`Chunk`."""
145 embed_text = (
146 chunker.contextualize(doc_chunk)
147 if self._config.include_context_in_embed
148 else doc_chunk.text
149 )
150 return Chunk(
151 text=doc_chunk.text,
152 embed_text=embed_text,
153 structure=self._projector.project(doc_chunk),
154 token_count=chunker.tokenizer.count_tokens(embed_text),
155 )