feat: add serializer provider to chunkers (#239)

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

Author: vagenas

docling_core/experimental/serializer/base.py

@@ -237,3 +237,12 @@ def serialize_captions(
     def get_excluded_refs(self, **kwargs) -> list[str]:
         """Get references to excluded items."""
         ...
+
+
+class BaseSerializerProvider(ABC):
+    """Base class for document serializer providers."""
+
+    @abstractmethod
+    def get_serializer(self, doc: DoclingDocument) -> BaseDocSerializer:
+        """Get a the associated serializer."""
+        ...

docling_core/transforms/chunker/hierarchical_chunker.py

@@ -11,11 +11,12 @@
 import re
 from typing import Any, ClassVar, Final, Iterator, Literal, Optional
 
-from pydantic import Field, StringConstraints, field_validator
+from pydantic import ConfigDict, Field, StringConstraints, field_validator
 from typing_extensions import Annotated, override
 
 from docling_core.experimental.serializer.base import (
     BaseDocSerializer,
+    BaseSerializerProvider,
     BaseTableSerializer,
     SerializationResult,
 )
@@ -183,6 +184,15 @@ class ChunkingDocSerializer(MarkdownDocSerializer):
     )
 
 
+class ChunkingSerializerProvider(BaseSerializerProvider):
+    """Serializer provider used for chunking purposes."""
+
+    @override
+    def get_serializer(self, doc: DoclingDocument) -> BaseDocSerializer:
+        """Get the associated serializer."""
+        return ChunkingDocSerializer(doc=doc)
+
+
 class HierarchicalChunker(BaseChunker):
     r"""Chunker implementation leveraging the document layout.
 
@@ -192,12 +202,16 @@ class HierarchicalChunker(BaseChunker):
         delim (str): Delimiter to use for merging text. Defaults to "\n".
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    serializer_provider: BaseSerializerProvider = ChunkingSerializerProvider()
+
+    # deprecated:
     merge_list_items: Annotated[bool, Field(deprecated=True)] = True
 
     def chunk(
         self,
         dl_doc: DLDocument,
-        doc_serializer: Optional[BaseDocSerializer] = None,
         **kwargs: Any,
     ) -> Iterator[BaseChunk]:
         r"""Chunk the provided document.
@@ -208,7 +222,7 @@ def chunk(
         Yields:
             Iterator[Chunk]: iterator over extracted chunks
         """
-        my_doc_ser = doc_serializer or ChunkingDocSerializer(doc=dl_doc)
+        my_doc_ser = self.serializer_provider.get_serializer(doc=dl_doc)
         heading_by_level: dict[LevelNumber, str] = {}
         visited: set[str] = set()
         ser_res = create_ser_result()

docling_core/transforms/chunker/hybrid_chunker.py

@@ -4,14 +4,23 @@
 #
 
 """Hybrid chunker implementation leveraging both doc structure & token awareness."""
-
 import warnings
+from functools import cached_property
 from typing import Any, Iterable, Iterator, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, PositiveInt, TypeAdapter, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    PositiveInt,
+    TypeAdapter,
+    computed_field,
+    model_validator,
+)
 from typing_extensions import Self
 
-from docling_core.transforms.chunker.hierarchical_chunker import ChunkingDocSerializer
+from docling_core.transforms.chunker.hierarchical_chunker import (
+    ChunkingSerializerProvider,
+)
 
 try:
     import semchunk
@@ -22,7 +31,10 @@
         "`pip install 'docling-core[chunking]'`"
     )
 
-from docling_core.experimental.serializer.base import BaseDocSerializer
+from docling_core.experimental.serializer.base import (
+    BaseDocSerializer,
+    BaseSerializerProvider,
+)
 from docling_core.transforms.chunker import (
     BaseChunk,
     BaseChunker,
@@ -52,7 +64,7 @@ class HybridChunker(BaseChunker):
     max_tokens: int = None  # type: ignore[assignment]
     merge_peers: bool = True
 
-    _inner_chunker: HierarchicalChunker = HierarchicalChunker()
+    serializer_provider: BaseSerializerProvider = ChunkingSerializerProvider()
 
     @model_validator(mode="after")
     def _patch_tokenizer_and_max_tokens(self) -> Self:
@@ -67,6 +79,11 @@ def _patch_tokenizer_and_max_tokens(self) -> Self:
             )
         return self
 
+    @computed_field  # type: ignore[misc]
+    @cached_property
+    def _inner_chunker(self) -> HierarchicalChunker:
+        return HierarchicalChunker(serializer_provider=self.serializer_provider)
+
     def _count_text_tokens(self, text: Optional[Union[str, list[str]]]):
         if text is None:
             return 0
@@ -246,7 +263,6 @@ def _merge_chunks_with_matching_metadata(self, chunks: list[DocChunk]):
     def chunk(
         self,
         dl_doc: DoclingDocument,
-        doc_serializer: Optional[BaseDocSerializer] = None,
         **kwargs: Any,
     ) -> Iterator[BaseChunk]:
         r"""Chunk the provided document.
@@ -257,7 +273,7 @@ def chunk(
         Yields:
             Iterator[Chunk]: iterator over extracted chunks
         """
-        my_doc_ser = doc_serializer or ChunkingDocSerializer(doc=dl_doc)
+        my_doc_ser = self.serializer_provider.get_serializer(doc=dl_doc)
         res: Iterable[DocChunk]
         res = self._inner_chunker.chunk(
             dl_doc=dl_doc,

examples/chunking_and_serialization.ipynb

@@ -66,22 +66,12 @@
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "chunker.max_tokens=512\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from transformers import AutoTokenizer\n",
     "\n",
     "EMBED_MODEL_ID = \"sentence-transformers/all-MiniLM-L6-v2\"\n",
-    "tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_ID)\n",
-    "chunker = HybridChunker(tokenizer=tokenizer)\n",
-    "print(f\"{chunker.max_tokens=}\")"
+    "tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_ID)"
    ]
   },
   {
@@ -202,6 +192,8 @@
     }
    ],
    "source": [
+    "chunker = HybridChunker(tokenizer=tokenizer)\n",
+    "\n",
     "chunk_iter = chunker.chunk(dl_doc=doc)\n",
     "\n",
     "chunks = list(chunk_iter)\n",
@@ -279,12 +271,22 @@
     }
    ],
    "source": [
-    "doc_serializer = ChunkingDocSerializer(\n",
-    "    doc=doc,\n",
-    "    table_serializer=MarkdownTableSerializer(),  # configuring a different table serializer\n",
+    "from docling_core.transforms.chunker.hierarchical_chunker import ChunkingSerializerProvider\n",
+    "\n",
+    "\n",
+    "class MDTableSerializerProvider(ChunkingSerializerProvider):\n",
+    "    def get_serializer(self, doc):\n",
+    "        return ChunkingDocSerializer(\n",
+    "            doc=doc,\n",
+    "            table_serializer=MarkdownTableSerializer(),  # configuring a different table serializer\n",
+    "        )\n",
+    "\n",
+    "chunker = HybridChunker(\n",
+    "    tokenizer=tokenizer,\n",
+    "    serializer_provider=MDTableSerializerProvider(),\n",
     ")\n",
     "\n",
-    "chunk_iter = chunker.chunk(dl_doc=doc, doc_serializer=doc_serializer)\n",
+    "chunk_iter = chunker.chunk(dl_doc=doc)\n",
     "\n",
     "chunks = list(chunk_iter)\n",
     "i, chunk = find_n_th_chunk_with_label(chunks, n=0, label=DocItemLabel.TABLE)\n",
@@ -355,14 +357,21 @@
    "source": [
     "from docling_core.experimental.serializer.markdown import MarkdownParams\n",
     "\n",
-    "doc_serializer = ChunkingDocSerializer(\n",
-    "    doc=doc,\n",
-    "    params=MarkdownParams(\n",
-    "        image_placeholder=\"<!-- image -->\",\n",
-    "    ),\n",
+    "class ImgPlaceholderSerializerProvider(ChunkingSerializerProvider):\n",
+    "    def get_serializer(self, doc):\n",
+    "        return ChunkingDocSerializer(\n",
+    "            doc=doc,\n",
+    "            params=MarkdownParams(\n",
+    "                image_placeholder=\"<!-- image -->\",\n",
+    "            ),\n",
+    "        )\n",
+    "\n",
+    "chunker = HybridChunker(\n",
+    "    tokenizer=tokenizer,\n",
+    "    serializer_provider=ImgPlaceholderSerializerProvider(),\n",
     ")\n",
     "\n",
-    "chunk_iter = chunker.chunk(dl_doc=doc, doc_serializer=doc_serializer)\n",
+    "chunk_iter = chunker.chunk(dl_doc=doc)\n",
     "\n",
     "chunks = list(chunk_iter)\n",
     "i, chunk = find_n_th_chunk_with_label(chunks, n=0, label=DocItemLabel.PICTURE)\n",
@@ -466,12 +475,19 @@
     }
    ],
    "source": [
-    "doc_serializer = ChunkingDocSerializer(\n",
-    "    doc=doc,\n",
-    "    picture_serializer=AnnotationPictureSerializer(),  # configuring a different picture serializer\n",
+    "class ImgAnnotationSerializerProvider(ChunkingSerializerProvider):\n",
+    "    def get_serializer(self, doc):\n",
+    "        return ChunkingDocSerializer(\n",
+    "            doc=doc,\n",
+    "            picture_serializer=AnnotationPictureSerializer(),  # configuring a different picture serializer\n",
+    "        )\n",
+    "\n",
+    "chunker = HybridChunker(\n",
+    "    tokenizer=tokenizer,\n",
+    "    serializer_provider=ImgAnnotationSerializerProvider(),\n",
     ")\n",
     "\n",
-    "chunk_iter = chunker.chunk(dl_doc=doc, doc_serializer=doc_serializer)\n",
+    "chunk_iter = chunker.chunk(dl_doc=doc)\n",
     "\n",
     "chunks = list(chunk_iter)\n",
     "i, chunk = find_n_th_chunk_with_label(chunks, n=0, label=DocItemLabel.PICTURE)\n",

test/test_hierarchical_chunker.py

@@ -9,9 +9,11 @@
 from docling_core.transforms.chunker import HierarchicalChunker
 from docling_core.transforms.chunker.hierarchical_chunker import (
     ChunkingDocSerializer,
+    ChunkingSerializerProvider,
     DocChunk,
 )
 from docling_core.types.doc import DoclingDocument as DLDocument
+from docling_core.types.doc.document import DoclingDocument
 
 from .test_data_gen_flag import GEN_TEST_DATA
 
@@ -48,18 +50,20 @@ def test_chunk_custom_serializer():
     with open("test/data/chunker/0_inp_dl_doc.json", encoding="utf-8") as f:
         data_json = f.read()
     dl_doc = DLDocument.model_validate_json(data_json)
+
+    class MySerializerProvider(ChunkingSerializerProvider):
+        def get_serializer(self, doc: DoclingDocument):
+            return ChunkingDocSerializer(
+                doc=doc,
+                table_serializer=MarkdownTableSerializer(),
+            )
+
     chunker = HierarchicalChunker(
         merge_list_items=True,
-    )
-    doc_serializer = ChunkingDocSerializer(
-        doc=dl_doc,
-        table_serializer=MarkdownTableSerializer(),
+        serializer_provider=MySerializerProvider(),
     )
 
-    chunks = chunker.chunk(
-        dl_doc=dl_doc,
-        doc_serializer=doc_serializer,
-    )
+    chunks = chunker.chunk(dl_doc=dl_doc)
     act_data = dict(
         root=[DocChunk.model_validate(n).export_json_dict() for n in chunks]
     )

test/test_hybrid_chunker.py

@@ -10,10 +10,12 @@
 from docling_core.experimental.serializer.markdown import MarkdownTableSerializer
 from docling_core.transforms.chunker.hierarchical_chunker import (
     ChunkingDocSerializer,
+    ChunkingSerializerProvider,
     DocChunk,
 )
 from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
 from docling_core.types.doc import DoclingDocument as DLDocument
+from docling_core.types.doc.document import DoclingDocument
 
 from .test_data_gen_flag import GEN_TEST_DATA
 
@@ -193,20 +195,21 @@ def test_chunk_custom_serializer():
         data_json = f.read()
     dl_doc = DLDocument.model_validate_json(data_json)
 
+    class MySerializerProvider(ChunkingSerializerProvider):
+        def get_serializer(self, doc: DoclingDocument):
+            return ChunkingDocSerializer(
+                doc=doc,
+                table_serializer=MarkdownTableSerializer(),
+            )
+
     chunker = HybridChunker(
         tokenizer=TOKENIZER,
         max_tokens=MAX_TOKENS,
         merge_peers=True,
-    )
-    doc_serializer = ChunkingDocSerializer(
-        doc=dl_doc,
-        table_serializer=MarkdownTableSerializer(),  # configuring a different table serializer
+        serializer_provider=MySerializerProvider(),
     )
 
-    chunk_iter = chunker.chunk(
-        dl_doc=dl_doc,
-        doc_serializer=doc_serializer,
-    )
+    chunk_iter = chunker.chunk(dl_doc=dl_doc)
     chunks = list(chunk_iter)
     act_data = dict(
         root=[DocChunk.model_validate(n).export_json_dict() for n in chunks]