docling-core 2.23.3__py3-none-any.whl → 2.24.0__py3-none-any.whl

docling_core/experimental/serializer/base.py

@@ -197,7 +197,7 @@ class BaseDocSerializer(ABC):
     @abstractmethod
     def get_parts(
         self,
-        node: Optional[NodeItem] = None,
+        item: Optional[NodeItem] = None,
         **kwargs,
     ) -> list[SerializationResult]:
         """Get the components to be combined for serializing this node."""
@@ -222,6 +222,6 @@ class BaseDocSerializer(ABC):
         ...
 
     @abstractmethod
-    def get_excluded_refs(self) -> list[str]:
+    def get_excluded_refs(self, **kwargs) -> list[str]:
         """Get references to excluded items."""
         ...

docling_core/experimental/serializer/common.py

@@ -5,12 +5,14 @@
 
 """Define base classes for serialization."""
 import sys
+from abc import abstractmethod
+from copy import deepcopy
 from functools import cached_property
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
-from pydantic import AnyUrl, BaseModel, computed_field
-from typing_extensions import override
+from pydantic import AnyUrl, BaseModel, NonNegativeInt, computed_field
+from typing_extensions import Self, override
 
 from docling_core.experimental.serializer.base import (
     BaseDocSerializer,
@@ -24,9 +26,7 @@ from docling_core.experimental.serializer.base import (
     BaseTextSerializer,
     SerializationResult,
 )
-from docling_core.types.doc.base import ImageRefMode
 from docling_core.types.doc.document import (
-    DEFAULT_CONTENT_LAYERS,
     DOCUMENT_TOKENS_EXPORT_LABELS,
     ContentLayer,
     DocItem,
@@ -38,10 +38,7 @@ from docling_core.types.doc.document import (
     KeyValueItem,
     NodeItem,
     OrderedList,
-    PictureClassificationData,
-    PictureDescriptionData,
     PictureItem,
-    PictureMoleculeData,
     TableItem,
     TextItem,
     UnorderedList,
@@ -49,6 +46,30 @@ from docling_core.types.doc.document import (
 from docling_core.types.doc.labels import DocItemLabel
 
 _DEFAULT_LABELS = DOCUMENT_TOKENS_EXPORT_LABELS
+_DEFAULT_LAYERS = {cl for cl in ContentLayer}
+
+
+class CommonParams(BaseModel):
+    """Common serialization parameters."""
+
+    # allowlists with non-recursive semantics, i.e. if a list group node is outside the
+    # range and some of its children items are within, they will be serialized
+    labels: set[DocItemLabel] = _DEFAULT_LABELS
+    layers: set[ContentLayer] = _DEFAULT_LAYERS
+    pages: Optional[set[int]] = None  # None means all pages are allowed
+
+    # slice-like semantics: start is included, stop is excluded
+    start_idx: NonNegativeInt = 0
+    stop_idx: NonNegativeInt = sys.maxsize
+
+    include_formatting: bool = True
+    include_hyperlinks: bool = True
+    caption_delim: str = " "
+
+    def merge_with_patch(self, patch: dict[str, Any]) -> Self:
+        """Create an instance by merging the provided patch dict on top of self."""
+        res = self.model_validate({**self.model_dump(), **patch})
+        return res
 
 
 class DocSerializer(BaseModel, BaseDocSerializer):
@@ -58,22 +79,10 @@ class DocSerializer(BaseModel, BaseDocSerializer):
         """Pydantic config."""
 
         arbitrary_types_allowed = True
+        extra = "forbid"
 
     doc: DoclingDocument
 
-    include_formatting: bool = True
-    include_hyperlinks: bool = True
-    escape_underscores: bool = True
-
-    # this filtering criteria are non-recursive;
-    # e.g. if a list group node is outside the range and some of its children items are
-    # within, they will be serialized
-    start: int = 0
-    stop: int = sys.maxsize
-    labels: set[DocItemLabel] = _DEFAULT_LABELS
-    layers: set[ContentLayer] = DEFAULT_CONTENT_LAYERS
-    pages: Optional[set[int]] = None
-
     text_serializer: BaseTextSerializer
     table_serializer: BaseTableSerializer
     picture_serializer: BasePictureSerializer
@@ -84,150 +93,242 @@ class DocSerializer(BaseModel, BaseDocSerializer):
     list_serializer: BaseListSerializer
     inline_serializer: BaseInlineSerializer
 
-    # these will be passed to the picture serializer (None defers/delegates fallback
-    # setting to callee):
-    image_placeholder: Optional[str] = None
-    image_mode: Optional[ImageRefMode] = None
+    params: CommonParams = CommonParams()
+
+    _excluded_refs_cache: dict[str, list[str]] = {}
 
     @computed_field  # type: ignore[misc]
     @cached_property
-    def _excluded_refs(self) -> list[str]:
-        refs: list[str] = [
-            item.self_ref
-            for ix, (item, _) in enumerate(
-                self.doc.iterate_items(
-                    with_groups=True,
-                    traverse_pictures=True,
-                )
+    def _captions_of_some_item(self) -> set[str]:
+        layers = {cl for cl in ContentLayer}  # TODO review
+        refs = {
+            cap.cref
+            for (item, _) in self.doc.iterate_items(
+                with_groups=True,
+                traverse_pictures=True,
+                included_content_layers=layers,
             )
-            if (
-                (ix < self.start or ix >= self.stop)
-                or (
-                    isinstance(item, DocItem)
-                    and (
-                        item.label not in self.labels
-                        or item.content_layer not in self.layers
-                        or (
-                            self.pages is not None
-                            and (
-                                (not item.prov)
-                                or item.prov[0].page_no not in self.pages
+            for cap in (item.captions if isinstance(item, FloatingItem) else [])
+        }
+        return refs
+
+    @override
+    def get_excluded_refs(self, **kwargs) -> list[str]:
+        """References to excluded items."""
+        params = self.params.merge_with_patch(patch=kwargs)
+        params_json = params.model_dump_json()
+        refs = self._excluded_refs_cache.get(params_json)
+        if refs is None:
+            refs = [
+                item.self_ref
+                for ix, (item, _) in enumerate(
+                    self.doc.iterate_items(
+                        with_groups=True,
+                        traverse_pictures=True,
+                        included_content_layers=params.layers,
+                    )
+                )
+                if (
+                    (ix < params.start_idx or ix >= params.stop_idx)
+                    or (
+                        isinstance(item, DocItem)
+                        and (
+                            item.label not in params.labels
+                            or item.content_layer not in params.layers
+                            or (
+                                params.pages is not None
+                                and (
+                                    (not item.prov)
+                                    or item.prov[0].page_no not in params.pages
+                                )
                             )
                         )
                     )
                 )
-            )
-        ]
+            ]
+            self._excluded_refs_cache[params_json] = refs
         return refs
 
-    @override
-    def get_excluded_refs(self) -> list[str]:
-        """References to excluded items."""
-        return self._excluded_refs
+    @abstractmethod
+    def serialize_page(self, parts: list[SerializationResult]) -> SerializationResult:
+        """Serialize a page out of its parts."""
+        ...
+
+    @abstractmethod
+    def serialize_doc(self, pages: list[SerializationResult]) -> SerializationResult:
+        """Serialize a document out of its pages."""
+        ...
+
+    def _serialize_body(self) -> SerializationResult:
+        """Serialize the document body."""
+        # find page ranges if available; otherwise regard whole doc as a single page
+        last_page: Optional[int] = None
+        starts: list[int] = []
+        for ix, (item, _) in enumerate(
+            self.doc.iterate_items(
+                with_groups=True,
+                traverse_pictures=True,
+                included_content_layers=self.params.layers,
+            )
+        ):
+            if isinstance(item, DocItem):
+                if item.prov:
+                    if last_page is None or item.prov[0].page_no > last_page:
+                        starts.append(ix)
+                        last_page = item.prov[0].page_no
+        page_ranges = [
+            (
+                (starts[i] if i > 0 else 0),
+                (starts[i + 1] if i < len(starts) - 1 else sys.maxsize),
+            )
+            for i, _ in enumerate(starts)
+        ] or [
+            (0, sys.maxsize)
+        ]  # use whole range if no pages detected
+
+        page_results: list[SerializationResult] = []
+        for page_range in page_ranges:
+            params_to_pass = deepcopy(self.params)
+            params_to_pass.start_idx = page_range[0]
+            params_to_pass.stop_idx = page_range[1]
+            subparts = self.get_parts(**params_to_pass.model_dump())
+            page_res = self.serialize_page(subparts)
+            page_results.append(page_res)
+        res = self.serialize_doc(page_results)
+        return res
 
-    # making some assumptions about the kwargs it can pass
     @override
-    def get_parts(
+    def serialize(
         self,
-        node: Optional[NodeItem] = None,
         *,
-        traverse_pictures: bool = False,
+        item: Optional[NodeItem] = None,
         list_level: int = 0,
         is_inline_scope: bool = False,
         visited: Optional[set[str]] = None,  # refs of visited items
         **kwargs,
-    ) -> list[SerializationResult]:
-        """Get the components to be combined for serializing this node."""
+    ) -> SerializationResult:
+        """Serialize a given node."""
         my_visited: set[str] = visited if visited is not None else set()
-        parts: list[SerializationResult] = []
+        empty_res = SerializationResult(text="")
+        if item is None or item == self.doc.body:
+            if self.doc.body.self_ref not in my_visited:
+                my_visited.add(self.doc.body.self_ref)
+                return self._serialize_body()
+            else:
+                return empty_res
 
-        label_blocklist = {
-            DocItemLabel.CAPTION,
-            DocItemLabel.FOOTNOTE,
-            # TODO handle differently as it clashes with self.labels
-        }
-        for ix, (item, _) in enumerate(
-            self.doc.iterate_items(
-                root=node,
-                with_groups=True,
-                traverse_pictures=traverse_pictures,
-                # ...
+        ########
+        # groups
+        ########
+        if isinstance(item, (UnorderedList, OrderedList)):
+            part = self.list_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                list_level=list_level,
+                is_inline_scope=is_inline_scope,
+                visited=my_visited,
+                **kwargs,
             )
-        ):
-            if item.self_ref in my_visited:
-                continue
+        elif isinstance(item, InlineGroup):
+            part = self.inline_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                list_level=list_level,
+                visited=my_visited,
+                **kwargs,
+            )
+        ###########
+        # doc items
+        ###########
+        elif isinstance(item, TextItem):
+            if item.self_ref in self._captions_of_some_item:
+                # those captions will be handled by the floating item holding them
+                return empty_res
             else:
-                my_visited.add(item.self_ref)
-
-            ########
-            # groups
-            ########
-            if isinstance(item, (UnorderedList, OrderedList)):
-                part = self.list_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                    list_level=list_level,
-                    is_inline_scope=is_inline_scope,
-                    visited=my_visited,
-                )
-            elif isinstance(item, InlineGroup):
-                part = self.inline_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                    list_level=list_level,
-                    visited=my_visited,
-                )
-            ###########
-            # doc items
-            ###########
-            elif isinstance(item, DocItem) and item.label in label_blocklist:
-                continue
-            elif isinstance(item, TextItem):
                 part = (
                     self.text_serializer.serialize(
                         item=item,
                         doc_serializer=self,
                         doc=self.doc,
                         is_inline_scope=is_inline_scope,
+                        **kwargs,
                     )
-                    if item.self_ref not in self.get_excluded_refs()
-                    else SerializationResult(text="")
-                )
-            elif isinstance(item, TableItem):
-                part = self.table_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                )
-            elif isinstance(item, PictureItem):
-                part = self.picture_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                    visited=my_visited,
-                    image_mode=self.image_mode,
-                    image_placeholder=self.image_placeholder,
-                )
-            elif isinstance(item, KeyValueItem):
-                part = self.key_value_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                )
-            elif isinstance(item, FormItem):
-                part = self.form_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
+                    if item.self_ref not in self.get_excluded_refs(**kwargs)
+                    else empty_res
                 )
+        elif isinstance(item, TableItem):
+            part = self.table_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                **kwargs,
+            )
+        elif isinstance(item, PictureItem):
+            part = self.picture_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                visited=my_visited,
+                **kwargs,
+            )
+        elif isinstance(item, KeyValueItem):
+            part = self.key_value_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                **kwargs,
+            )
+        elif isinstance(item, FormItem):
+            part = self.form_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                **kwargs,
+            )
+        else:
+            part = self.fallback_serializer.serialize(
+                item=item,
+                doc_serializer=self,
+                doc=self.doc,
+                **kwargs,
+            )
+        return part
+
+    # making some assumptions about the kwargs it can pass
+    @override
+    def get_parts(
+        self,
+        item: Optional[NodeItem] = None,
+        *,
+        traverse_pictures: bool = False,
+        list_level: int = 0,
+        is_inline_scope: bool = False,
+        visited: Optional[set[str]] = None,  # refs of visited items
+        **kwargs,
+    ) -> list[SerializationResult]:
+        """Get the components to be combined for serializing this node."""
+        parts: list[SerializationResult] = []
+        my_visited: set[str] = visited if visited is not None else set()
+        params = self.params.merge_with_patch(patch=kwargs)
+        for item, _ in self.doc.iterate_items(
+            root=item,
+            with_groups=True,
+            traverse_pictures=traverse_pictures,
+            included_content_layers=params.layers,
+        ):
+            if item.self_ref in my_visited:
+                continue
             else:
-                part = self.fallback_serializer.serialize(
-                    item=item,
-                    doc_serializer=self,
-                    doc=self.doc,
-                )
+                my_visited.add(item.self_ref)
+            part = self.serialize(
+                item=item,
+                list_level=list_level,
+                is_inline_scope=is_inline_scope,
+                visited=my_visited,
+                **kwargs,
+            )
             if part.text:
                 parts.append(part)
         return parts
@@ -242,8 +343,9 @@ class DocSerializer(BaseModel, BaseDocSerializer):
         **kwargs,
     ) -> str:
         """Apply some text post-processing steps."""
+        params = self.params.merge_with_patch(patch=kwargs)
         res = text
-        if self.include_formatting and formatting:
+        if params.include_formatting and formatting:
             if formatting.bold:
                 res = self.serialize_bold(text=res)
             if formatting.italic:
@@ -252,7 +354,7 @@ class DocSerializer(BaseModel, BaseDocSerializer):
                 res = self.serialize_underline(text=res)
             if formatting.strikethrough:
                 res = self.serialize_strikethrough(text=res)
-        if self.include_hyperlinks and hyperlink:
+        if params.include_hyperlinks and hyperlink:
             res = self.serialize_hyperlink(text=res, hyperlink=hyperlink)
         return res
 
@@ -287,67 +389,19 @@ class DocSerializer(BaseModel, BaseDocSerializer):
     def serialize_captions(
         self,
         item: FloatingItem,
-        separator: Optional[str] = None,
         **kwargs,
     ) -> SerializationResult:
         """Serialize the item's captions."""
-        text_parts: list[str] = [
-            it.text
-            for cap in item.captions
-            if isinstance(it := cap.resolve(self.doc), TextItem)
-            and it.self_ref not in self.get_excluded_refs()
-        ]
-        text_res = (separator or "\n").join(text_parts)
-        text_res = self.post_process(text=text_res)
-        return SerializationResult(text=text_res)
-
-
-class PictureSerializer(BasePictureSerializer):
-    """Class for picture serializers."""
-
-    # helper function
-    def _serialize_content(
-        self,
-        item: PictureItem,
-        doc_serializer: "BaseDocSerializer",
-        doc: DoclingDocument,
-        separator: Optional[str] = None,
-        visited: Optional[set[str]] = None,
-        **kwargs,
-    ) -> SerializationResult:
-        parts = doc_serializer.get_parts(
-            node=item,
-            traverse_pictures=True,
-            visited=visited,
-        )
-        text_res = (separator or " ").join([p.text for p in parts])
-        # NOTE: we do no postprocessing since already done as needed
-        return SerializationResult(text=text_res)
-
-    # helper function
-    def _serialize_annotations(
-        self,
-        item: PictureItem,
-        doc_serializer: "BaseDocSerializer",
-        doc: DoclingDocument,
-        separator: Optional[str] = None,
-        **kwargs,
-    ) -> SerializationResult:
-        text_parts: list[str] = []
-        for annotation in item.annotations:
-            if isinstance(annotation, PictureClassificationData):
-                predicted_class = (
-                    annotation.predicted_classes[0].class_name
-                    if annotation.predicted_classes
-                    else None
-                )
-                if predicted_class is not None:
-                    text_parts.append(f"Picture type: {predicted_class}")
-            elif isinstance(annotation, PictureMoleculeData):
-                text_parts.append(f"SMILES: {annotation.smi}")
-            elif isinstance(annotation, PictureDescriptionData):
-                text_parts.append(f"Description: {annotation.text}")
-
-        text_res = (separator or "\n").join(text_parts)
-        text_res = doc_serializer.post_process(text=text_res)
+        params = self.params.merge_with_patch(patch=kwargs)
+        if DocItemLabel.CAPTION in params.labels:
+            text_parts: list[str] = [
+                it.text
+                for cap in item.captions
+                if isinstance(it := cap.resolve(self.doc), TextItem)
+                and it.self_ref not in self.get_excluded_refs(**kwargs)
+            ]
+            text_res = params.caption_delim.join(text_parts)
+            text_res = self.post_process(text=text_res)
+        else:
+            text_res = ""
         return SerializationResult(text=text_res)