feat: add page chunking (#337)

Signed-off-by: Panos Vagenas <[email protected]>

Author: vagenas

docling_core/transforms/chunker/__init__.py

@@ -11,3 +11,4 @@
     DocMeta,
     HierarchicalChunker,
 )
+from docling_core.transforms.chunker.page_chunker import PageChunker

docling_core/transforms/chunker/page_chunker.py

@@ -0,0 +1,59 @@
+"""Page-based chunker implementation: each chunk corresponds to a single page."""
+
+from __future__ import annotations
+
+from typing import Any, Iterator
+
+from pydantic import ConfigDict
+from typing_extensions import override
+
+from docling_core.transforms.chunker import BaseChunker, DocChunk, DocMeta
+from docling_core.transforms.chunker.hierarchical_chunker import (
+    ChunkingSerializerProvider,
+)
+from docling_core.types import DoclingDocument as DLDocument
+
+
+class PageChunker(BaseChunker):
+    r"""Chunker implementation that yields one chunk per page."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    serializer_provider: ChunkingSerializerProvider = ChunkingSerializerProvider()
+
+    @override
+    def chunk(
+        self,
+        dl_doc: DLDocument,
+        **kwargs: Any,
+    ) -> Iterator[DocChunk]:
+        """Chunk the provided document by page."""
+        my_doc_ser = self.serializer_provider.get_serializer(doc=dl_doc)
+        if dl_doc.pages:
+            # chunk by page
+            for page_no in sorted(dl_doc.pages.keys()):
+                ser_res = my_doc_ser.serialize(pages={page_no})
+                if not ser_res.text:
+                    continue
+                yield DocChunk(
+                    text=ser_res.text,
+                    meta=DocMeta(
+                        doc_items=ser_res.get_unique_doc_items(),
+                        headings=None,
+                        captions=None,
+                        origin=dl_doc.origin,
+                    ),
+                )
+        else:
+            # if no pages, treat whole document as single chunk
+            ser_res = my_doc_ser.serialize()
+            if ser_res.text:
+                yield DocChunk(
+                    text=ser_res.text,
+                    meta=DocMeta(
+                        doc_items=ser_res.get_unique_doc_items(),
+                        headings=None,
+                        captions=None,
+                        origin=dl_doc.origin,
+                    ),
+                )

docling_core/transforms/serializer/base.py

@@ -39,6 +39,16 @@ class SerializationResult(BaseModel):
     spans: list[Span] = []
     # group: Optional[GroupItem] = None  # set when result reflects specific group item
 
+    def get_unique_doc_items(self) -> list[DocItem]:
+        """Get the doc items corresponding to this result."""
+        seen_doc_item_refs: set[str] = set()
+        doc_items: list[DocItem] = []
+        for span in self.spans:
+            if span.item.self_ref not in seen_doc_item_refs:
+                seen_doc_item_refs.add(span.item.self_ref)
+                doc_items.append(span.item)
+        return doc_items
+
 
 class BaseTextSerializer(ABC):
     """Base class for text item serializers."""

docling_core/transforms/serializer/common.py

@@ -285,7 +285,7 @@ def serialize_doc(
 
     def _serialize_body(self, **kwargs) -> SerializationResult:
         """Serialize the document body."""
-        subparts = self.get_parts()
+        subparts = self.get_parts(**kwargs)
         res = self.serialize_doc(parts=subparts, **kwargs)
         return res