use Field constraints instead of conlist, refactor chunking types

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

Author: vagenas

docling_core/transforms/chunker/__init__.py

@@ -5,8 +5,5 @@
 
 """Define the chunker types."""
 
-from docling_core.transforms.chunker.base import BaseChunk, BaseChunker  # noqa
-from docling_core.transforms.chunker.hierarchical_chunker import (  # noqa
-    Chunk,
-    HierarchicalChunker,
-)
+from docling_core.transforms.chunker.base import BaseChunk, BaseChunker, BaseMeta
+from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker

docling_core/transforms/chunker/base.py

@@ -5,22 +5,45 @@
 
 """Define base classes for chunking."""
 from abc import ABC, abstractmethod
-from typing import Any, Iterator
+from typing import Any, ClassVar, Iterator
 
 from pydantic import BaseModel
 
 from docling_core.types.doc import DoclingDocument as DLDocument
 
 
+class BaseMeta(BaseModel):
+    """Metadata base class."""
+
+    excluded_embed: ClassVar[list[str]] = []
+    excluded_llm: ClassVar[list[str]] = []
+
+    def export_json_dict(self) -> dict[str, Any]:
+        """Helper method for exporting non-None keys to JSON mode.
+
+        Returns:
+            dict[str, Any]: The exported dictionary.
+        """
+        return self.model_dump(mode="json", by_alias=True, exclude_none=True)
+
+
 class BaseChunk(BaseModel):
-    """Data model for base chunk."""
+    """Chunk base class."""
 
     text: str
-    meta: Any = None
+    meta: BaseMeta
+
+    def export_json_dict(self) -> dict[str, Any]:
+        """Helper method for exporting non-None keys to JSON mode.
+
+        Returns:
+            dict[str, Any]: The exported dictionary.
+        """
+        return self.model_dump(mode="json", by_alias=True, exclude_none=True)
 
 
 class BaseChunker(BaseModel, ABC):
-    """Base class for Chunker."""
+    """Chunker base class."""
 
     @abstractmethod
     def chunk(self, dl_doc: DLDocument, **kwargs) -> Iterator[BaseChunk]:

docling_core/transforms/chunker/hierarchical_chunker.py

@@ -11,10 +11,9 @@
 from typing import Any, ClassVar, Iterator, Optional
 
 from pandas import DataFrame
-from pydantic import BaseModel, Field, conlist
+from pydantic import Field
 
-from docling_core.transforms.chunker import BaseChunker
-from docling_core.transforms.chunker.base import BaseChunk
+from docling_core.transforms.chunker import BaseChunk, BaseChunker, BaseMeta
 from docling_core.types.doc import DoclingDocument as DLDocument
 from docling_core.types.doc.document import (
     DocItem,
@@ -33,50 +32,42 @@
 _logger = logging.getLogger(__name__)
 
 
-class ChunkMeta(BaseModel):
-    """Data model for specific chunk metadata."""
+class DocMeta(BaseMeta):
+    """Data model for Hierarchical Chunker metadata."""
 
-    # TODO align paths typewith _JSON_POINTER_REGEX
-    doc_items: conlist(DocItem, min_length=1) = Field(  # type: ignore
+    doc_items: list[DocItem] = Field(
         alias=_KEY_DOC_ITEMS,
+        min_length=1,
     )
-    headings: Optional[conlist(str, min_length=1)] = Field(  # type: ignore
+    headings: Optional[list[str]] = Field(
         default=None,
         alias=_KEY_HEADINGS,
+        min_length=1,
     )
-    captions: Optional[conlist(str, min_length=1)] = Field(  # type: ignore
+    captions: Optional[list[str]] = Field(
         default=None,
         alias=_KEY_CAPTIONS,
+        min_length=1,
     )
 
     excluded_embed: ClassVar[list[str]] = [_KEY_DOC_ITEMS]
     excluded_llm: ClassVar[list[str]] = [_KEY_DOC_ITEMS]
 
-    def export_json_dict(self) -> dict[str, Any]:
-        """Helper method for exporting non-None keys to JSON mode.
-
-        Returns:
-            dict[str, Any]: The exported dictionary.
-        """
-        return self.model_dump(mode="json", by_alias=True, exclude_none=True)
 
+class DocChunk(BaseChunk):
+    """Data model for Hierarchical Chunker chunks."""
 
-class Chunk(BaseChunk):
-    """Data model for specific chunk."""
-
-    meta: ChunkMeta
-
-    def export_json_dict(self) -> dict[str, Any]:
-        """Helper method for exporting non-None keys to JSON mode.
-
-        Returns:
-            dict[str, Any]: The exported dictionary.
-        """
-        return self.model_dump(mode="json", by_alias=True, exclude_none=True)
+    meta: DocMeta
 
 
 class HierarchicalChunker(BaseChunker):
-    """Chunker implementation leveraging the document layout."""
+    r"""Chunker implementation leveraging the document layout.
+
+    Args:
+        merge_list_items (bool): Whether to merge successive list items.
+            Defaults to True.
+        delim (str): Delimiter to use for merging text. Defaults to "\n".
+    """
 
     merge_list_items: bool = True
     delim: str = "\n"
@@ -129,9 +120,9 @@ def chunk(self, dl_doc: DLDocument, **kwargs: Any) -> Iterator[BaseChunk]:
                         list_items.append(item)
                         continue
                     elif list_items:  # need to yield
-                        yield Chunk(
+                        yield DocChunk(
                             text=self.delim.join([i.text for i in list_items]),
-                            meta=ChunkMeta(
+                            meta=DocMeta(
                                 doc_items=list_items,
                                 headings=[
                                     heading_by_level[k]
@@ -148,7 +139,7 @@ def chunk(self, dl_doc: DLDocument, **kwargs: Any) -> Iterator[BaseChunk]:
                     isinstance(item, TextItem)
                     and item.label == DocItemLabel.SECTION_HEADER
                 ):
-                    # TODO second branch not needed after cleanup above:
+                    # TODO second branch not needed once cleanup above complete:
                     level = item.level if isinstance(item, SectionHeaderItem) else 1
                     heading_by_level[level] = item.text
 
@@ -173,9 +164,9 @@ def chunk(self, dl_doc: DLDocument, **kwargs: Any) -> Iterator[BaseChunk]:
                     ] or None
                 else:
                     continue
-                c = Chunk(
+                c = DocChunk(
                     text=text,
-                    meta=ChunkMeta(
+                    meta=DocMeta(
                         doc_items=[item],
                         headings=[heading_by_level[k] for k in sorted(heading_by_level)]
                         or None,
@@ -185,9 +176,9 @@ def chunk(self, dl_doc: DLDocument, **kwargs: Any) -> Iterator[BaseChunk]:
                 yield c
 
         if self.merge_list_items and list_items:  # need to yield
-            yield Chunk(
+            yield DocChunk(
                 text=self.delim.join([i.text for i in list_items]),
-                meta=ChunkMeta(
+                meta=DocMeta(
                     doc_items=list_items,
                     headings=[heading_by_level[k] for k in sorted(heading_by_level)]
                     or None,

docling_core/transforms/id_generator/__init__.py

@@ -6,7 +6,7 @@
 import json
 
 from docling_core.transforms.chunker import HierarchicalChunker
-from docling_core.transforms.chunker.hierarchical_chunker import Chunk
+from docling_core.transforms.chunker.hierarchical_chunker import DocChunk
 from docling_core.types.doc import DoclingDocument as DLDocument
 
 
@@ -18,7 +18,9 @@ def test_chunk_merge_list_items():
         merge_list_items=True,
     )
     chunks = chunker.chunk(dl_doc=dl_doc)
-    act_data = dict(root=[Chunk.model_validate(n).export_json_dict() for n in chunks])
+    act_data = dict(
+        root=[DocChunk.model_validate(n).export_json_dict() for n in chunks]
+    )
     with open("test/data/chunker/0_out_chunks.json") as f:
         exp_data = json.load(fp=f)
     assert exp_data == act_data
@@ -32,7 +34,9 @@ def test_chunk_no_merge_list_items():
         merge_list_items=False,
     )
     chunks = chunker.chunk(dl_doc=dl_doc)
-    act_data = dict(root=[Chunk.model_validate(n).export_json_dict() for n in chunks])
+    act_data = dict(
+        root=[DocChunk.model_validate(n).export_json_dict() for n in chunks]
+    )
     with open("test/data/chunker/1_out_chunks.json") as f:
         exp_data = json.load(fp=f)
     assert exp_data == act_data