docling-core 1.6.2__py3-none-any.whl → 1.7.0__py3-none-any.whl

docling_core/search/json_schema_to_search_mapper.py

@@ -8,7 +8,7 @@ import re
 from copy import deepcopy
 from typing import Any, Optional, Pattern, Tuple, TypedDict
 
-from jsonref import JsonRef
+from jsonref import replace_refs
 
 
 class SearchIndexDefinition(TypedDict):
@@ -95,7 +95,11 @@ class JsonSchemaToSearchMapper:
         which define the fields, their data types, and other specifications to index
         JSON documents into a Lucene index.
         """
-        mapping = JsonRef.replace_refs(schema)
+        mapping = deepcopy(schema)
+
+        mapping = self._suppress(mapping, self._suppress_key)
+
+        mapping = replace_refs(mapping)
 
         mapping = self._merge_unions(mapping)
 
@@ -105,8 +109,6 @@ class JsonSchemaToSearchMapper:
 
         mapping = self._remove_keys(mapping, self._rm_keys)
 
-        mapping = self._suppress(mapping, self._suppress_key)
-
         mapping = self._translate_keys_re(mapping)
 
         mapping = self._clean(mapping)

docling_core/transforms/chunker/base.py

@@ -22,8 +22,9 @@ class Chunk(BaseModel):
 class ChunkWithMetadata(Chunk):
     """Data model for Chunk including metadata."""
 
-    page: Optional[int]
-    bbox: Optional[BoundingBox]
+    page: Optional[int] = None
+    bbox: Optional[BoundingBox] = None
+    heading: Optional[str] = None
 
 
 class BaseChunker(BaseModel, ABC):

docling_core/transforms/chunker/hierarchical_chunker.py

@@ -26,6 +26,7 @@ class HierarchicalChunker(BaseChunker):
     """Chunker implementation leveraging the document layout."""
 
     include_metadata: bool = True
+    heading_as_metadata: bool = False
     min_chunk_len: PositiveInt = 64
 
     class _NodeType(str, Enum):
@@ -184,7 +185,7 @@ class HierarchicalChunker(BaseChunker):
 
     def _build_chunk_impl(
         self, doc: DLDocument, doc_map: _DocContext, idx: int, rec: bool = False
-    ) -> list[_TextEntry]:
+    ) -> tuple[list[_TextEntry], Optional[str]]:
         if doc.main_text:
             item = doc.main_text[idx]
             item_type = _HC._norm(item.obj_type)
@@ -193,7 +194,7 @@ class HierarchicalChunker(BaseChunker):
                 item_type not in self._allowed_types
                 or item_name in self._disallowed_names_by_type.get(item_type, [])
             ):
-                return []
+                return [], None
 
             c2p = doc_map.dmap
 
@@ -219,7 +220,7 @@ class HierarchicalChunker(BaseChunker):
                         else []
                     )
                 else:
-                    return []
+                    return [], None
             elif isinstance(item, BaseText):
                 text_entries = [
                     self._TextEntry(
@@ -248,21 +249,29 @@ class HierarchicalChunker(BaseChunker):
                     _HC._NodeName.LIST_ITEM,
                     _HC._NodeName.SUBTITLE_LEVEL_1,
                 ]:
-                    return []
+                    return [], None
 
             if (parent := c2p[idx].parent) is not None:
                 # prepend with ancestors
+
+                parent_res = self._build_chunk_impl(
+                    doc=doc, doc_map=doc_map, idx=parent, rec=True
+                )
                 return (
-                    self._build_chunk_impl(
-                        doc=doc, doc_map=doc_map, idx=parent, rec=True
-                    )
-                    + text_entries
+                    parent_res[0] + text_entries,  # expanded text
+                    parent_res[1],  # heading
                 )
             else:
-                # if root, augment with title (if available and different)
-                return text_entries
+                if (
+                    self.heading_as_metadata
+                    and isinstance(item, BaseText)
+                    and _HC._norm(item.obj_type) == _HC._NodeType.SUBTITLE_LEVEL_1
+                ):
+                    return [], text_entries[0].text
+                else:
+                    return text_entries, None
         else:
-            return []
+            return [], None
 
     def _build_chunk(
         self,
@@ -272,7 +281,9 @@ class HierarchicalChunker(BaseChunker):
         delim: str,
         rec: bool = False,
     ) -> Optional[Chunk]:
-        texts = self._build_chunk_impl(doc=doc, doc_map=doc_map, idx=idx, rec=rec)
+        res = self._build_chunk_impl(doc=doc, doc_map=doc_map, idx=idx, rec=rec)
+        texts = res[0]
+        heading = res[1]
         concat = delim.join([t.text for t in texts if t.text])
         assert doc.main_text is not None
         if len(concat) >= self.min_chunk_len:
@@ -295,6 +306,7 @@ class HierarchicalChunker(BaseChunker):
                     path=path,
                     page=item.prov[0].page if item.prov else None,
                     bbox=item.prov[0].bbox if item.prov else None,
+                    heading=heading,
                 )
             else:
                 return Chunk(
@@ -315,6 +327,11 @@ class HierarchicalChunker(BaseChunker):
         Yields:
             Iterator[Chunk]: iterator over extracted chunks
         """
+        if (not self.include_metadata) and self.heading_as_metadata:
+            raise RuntimeError(
+                "To enable `heading_as_metadata`, also `include_metadata` must be True."
+            )
+
         if dl_doc.main_text:
             # extract doc structure incl. metadata for
             # each item (e.g. parent, children)

docling_core/transforms/id_generator/__init__.py

@@ -0,0 +1,12 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Define the ID generator types."""
+
+from docling_core.transforms.id_generator.base import BaseIDGenerator  # noqa
+from docling_core.transforms.id_generator.doc_hash_id_generator import (  # noqa
+    DocHashIDGenerator,
+)
+from docling_core.transforms.id_generator.uuid_generator import UUIDGenerator  # noqa

docling_core/transforms/id_generator/base.py

@@ -0,0 +1,30 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Base document ID generator module."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from docling_core.types import Document as DLDocument
+
+
+class BaseIDGenerator(ABC):
+    """Document ID generator base class."""
+
+    @abstractmethod
+    def generate_id(self, doc: DLDocument, *args: Any, **kwargs: Any) -> str:
+        """Generate an ID for the given document.
+
+        Args:
+            doc (DLDocument): document to generate ID for
+
+        Raises:
+            NotImplementedError: in this abstract implementation
+
+        Returns:
+            str: the generated ID
+        """
+        raise NotImplementedError()

docling_core/transforms/id_generator/doc_hash_id_generator.py

@@ -0,0 +1,27 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Doc-hash-based ID generator module."""
+
+
+from typing import Any
+
+from docling_core.transforms.id_generator import BaseIDGenerator
+from docling_core.types import Document as DLDocument
+
+
+class DocHashIDGenerator(BaseIDGenerator):
+    """Doc-hash-based ID generator class."""
+
+    def generate_id(self, doc: DLDocument, *args: Any, **kwargs: Any) -> str:
+        """Generate an ID for the given document.
+
+        Args:
+            doc (DLDocument): document to generate ID for
+
+        Returns:
+            str: the generated ID
+        """
+        return doc.file_info.document_hash

docling_core/transforms/id_generator/uuid_generator.py

@@ -0,0 +1,34 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""UUID-based ID generator module."""
+
+from random import Random
+from typing import Annotated, Any, Optional
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from docling_core.transforms.id_generator import BaseIDGenerator
+from docling_core.types import Document as DLDocument
+
+
+class UUIDGenerator(BaseModel, BaseIDGenerator):
+    """UUID-based ID generator class."""
+
+    seed: Optional[int] = None
+    uuid_version: Annotated[int, Field(strict=True, ge=1, le=5)] = 4
+
+    def generate_id(self, doc: DLDocument, *args: Any, **kwargs: Any) -> str:
+        """Generate an ID for the given document.
+
+        Args:
+            doc (DLDocument): document to generate ID for
+
+        Returns:
+            str: the generated ID
+        """
+        rd = Random(x=self.seed)
+        return str(UUID(int=rd.getrandbits(128), version=self.uuid_version))

docling_core/transforms/metadata_extractor/__init__.py

@@ -0,0 +1,13 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Define the metadata extractor types."""
+
+from docling_core.transforms.metadata_extractor.base import (  # noqa
+    BaseMetadataExtractor,
+)
+from docling_core.transforms.metadata_extractor.simple_metadata_extractor import (  # noqa
+    SimpleMetadataExtractor,
+)

docling_core/transforms/metadata_extractor/base.py

@@ -0,0 +1,59 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Base metadata extractor module."""
+
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pydantic import BaseModel
+
+from docling_core.types import Document as DLDocument
+
+
+class BaseMetadataExtractor(BaseModel, ABC):
+    """Metadata extractor base class."""
+
+    @abstractmethod
+    def get_metadata(
+        self, doc: DLDocument, *args: Any, **kwargs: Any
+    ) -> dict[str, Any]:
+        """Extract metadata for the given document.
+
+        Args:
+            doc (DLDocument): document to extract metadata for
+
+        Raises:
+            NotImplementedError: in this abstract implementation
+
+        Returns:
+            dict[str, Any]: the extracted metadata
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_excluded_embed_metadata_keys(self) -> list[str]:
+        """Get metadata keys to exclude from embedding.
+
+        Raises:
+            NotImplementedError: in this abstract implementation
+
+        Returns:
+            list[str]: the metadata to exclude
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_excluded_llm_metadata_keys(self) -> list[str]:
+        """Get metadata keys to exclude from LLM generation.
+
+        Raises:
+            NotImplementedError: in this abstract implementation
+
+        Returns:
+            list[str]: the metadata to exclude
+        """
+        raise NotImplementedError()

docling_core/transforms/metadata_extractor/simple_metadata_extractor.py

@@ -0,0 +1,61 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Simple metadata extractor module."""
+
+
+from enum import Enum
+from typing import Any
+
+from docling_core.transforms.metadata_extractor import BaseMetadataExtractor
+from docling_core.types import Document as DLDocument
+
+
+class SimpleMetadataExtractor(BaseMetadataExtractor):
+    """Simple metadata extractor class."""
+
+    class _Keys(str, Enum):
+        DL_DOC_HASH = "dl_doc_hash"
+        ORIGIN = "origin"
+
+    include_origin: bool = False
+
+    def get_metadata(
+        self, doc: DLDocument, origin: str, *args: Any, **kwargs: Any
+    ) -> dict[str, Any]:
+        """Extract metadata for the given document.
+
+        Args:
+            doc (DLDocument): document to extract metadata for
+            origin (str): the document origin
+
+        Returns:
+            dict[str, Any]: the extracted metadata
+        """
+        meta: dict[str, Any] = {
+            self._Keys.DL_DOC_HASH: doc.file_info.document_hash,
+        }
+        if self.include_origin:
+            meta[self._Keys.ORIGIN] = origin
+        return meta
+
+    def get_excluded_embed_metadata_keys(self) -> list[str]:
+        """Get metadata keys to exclude from embedding.
+
+        Returns:
+            list[str]: the metadata to exclude
+        """
+        excl_keys: list[str] = [self._Keys.DL_DOC_HASH]
+        if self.include_origin:
+            excl_keys.append(self._Keys.ORIGIN)
+        return excl_keys
+
+    def get_excluded_llm_metadata_keys(self) -> list[str]:
+        """Get metadata keys to exclude from LLM generation.
+
+        Returns:
+            list[str]: the metadata to exclude
+        """
+        return self.get_excluded_embed_metadata_keys()

docling_core/types/doc/base.py

@@ -440,7 +440,6 @@ class BaseText(BaseCell):
     ):
         """Export text element to document tokens format."""
         body = f"<{self.obj_type}>"
-        # body = f"<{self.name}>"
 
         assert DocumentToken.is_known_token(
             body

docling_core/types/experimental/__init__.py

@@ -0,0 +1,30 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+
+"""Package for models defined by the Document type."""
+
+from .base import BoundingBox, CoordOrigin, Size
+from .document import (
+    BasePictureData,
+    BaseTableData,
+    DescriptionItem,
+    DocItem,
+    DoclingDocument,
+    DocumentOrigin,
+    FloatingItem,
+    GroupItem,
+    ImageRef,
+    KeyValueItem,
+    NodeItem,
+    PageItem,
+    PictureItem,
+    ProvenanceItem,
+    RefItem,
+    SectionHeaderItem,
+    TableCell,
+    TableItem,
+    TextItem,
+)
+from .labels import DocItemLabel, GroupLabel, TableCellLabel

docling_core/types/experimental/base.py

@@ -0,0 +1,167 @@
+"""Models for the base data types."""
+
+import copy
+from enum import Enum
+from typing import Tuple
+
+from pydantic import BaseModel
+
+
+class CoordOrigin(str, Enum):
+    """CoordOrigin."""
+
+    TOPLEFT = "TOPLEFT"
+    BOTTOMLEFT = "BOTTOMLEFT"
+
+
+class Size(BaseModel):
+    """Size."""
+
+    width: float = 0.0
+    height: float = 0.0
+
+    def as_tuple(self):
+        """as_tuple."""
+        return (self.width, self.height)
+
+
+class BoundingBox(BaseModel):
+    """BoundingBox."""
+
+    l: float  # left
+    t: float  # top
+    r: float  # right
+    b: float  # bottom
+
+    coord_origin: CoordOrigin = CoordOrigin.TOPLEFT
+
+    @property
+    def width(self):
+        """width."""
+        return self.r - self.l
+
+    @property
+    def height(self):
+        """height."""
+        return abs(self.t - self.b)
+
+    def scaled(self, scale: float) -> "BoundingBox":
+        """scaled.
+
+        :param scale: float:
+
+        """
+        out_bbox = copy.deepcopy(self)
+        out_bbox.l *= scale
+        out_bbox.r *= scale
+        out_bbox.t *= scale
+        out_bbox.b *= scale
+
+        return out_bbox
+
+    def normalized(self, page_size: Size) -> "BoundingBox":
+        """normalized.
+
+        :param page_size: Size:
+
+        """
+        out_bbox = copy.deepcopy(self)
+        out_bbox.l /= page_size.width
+        out_bbox.r /= page_size.width
+        out_bbox.t /= page_size.height
+        out_bbox.b /= page_size.height
+
+        return out_bbox
+
+    def as_tuple(self):
+        """as_tuple."""
+        if self.coord_origin == CoordOrigin.TOPLEFT:
+            return (self.l, self.t, self.r, self.b)
+        elif self.coord_origin == CoordOrigin.BOTTOMLEFT:
+            return (self.l, self.b, self.r, self.t)
+
+    @classmethod
+    def from_tuple(cls, coord: Tuple[float, ...], origin: CoordOrigin):
+        """from_tuple.
+
+        :param coord: Tuple[float:
+        :param ...]:
+        :param origin: CoordOrigin:
+
+        """
+        if origin == CoordOrigin.TOPLEFT:
+            l, t, r, b = coord[0], coord[1], coord[2], coord[3]
+            if r < l:
+                l, r = r, l
+            if b < t:
+                b, t = t, b
+
+            return BoundingBox(l=l, t=t, r=r, b=b, coord_origin=origin)
+        elif origin == CoordOrigin.BOTTOMLEFT:
+            l, b, r, t = coord[0], coord[1], coord[2], coord[3]
+            if r < l:
+                l, r = r, l
+            if b > t:
+                b, t = t, b
+
+            return BoundingBox(l=l, t=t, r=r, b=b, coord_origin=origin)
+
+    def area(self) -> float:
+        """area."""
+        return (self.r - self.l) * (self.b - self.t)
+
+    def intersection_area_with(self, other: "BoundingBox") -> float:
+        """intersection_area_with.
+
+        :param other: "BoundingBox":
+
+        """
+        # Calculate intersection coordinates
+        left = max(self.l, other.l)
+        top = max(self.t, other.t)
+        right = min(self.r, other.r)
+        bottom = min(self.b, other.b)
+
+        # Calculate intersection dimensions
+        width = right - left
+        height = bottom - top
+
+        # If the bounding boxes do not overlap, width or height will be negative
+        if width <= 0 or height <= 0:
+            return 0.0
+
+        return width * height
+
+    def to_bottom_left_origin(self, page_height) -> "BoundingBox":
+        """to_bottom_left_origin.
+
+        :param page_height:
+
+        """
+        if self.coord_origin == CoordOrigin.BOTTOMLEFT:
+            return self
+        elif self.coord_origin == CoordOrigin.TOPLEFT:
+            return BoundingBox(
+                l=self.l,
+                r=self.r,
+                t=page_height - self.t,
+                b=page_height - self.b,
+                coord_origin=CoordOrigin.BOTTOMLEFT,
+            )
+
+    def to_top_left_origin(self, page_height):
+        """to_top_left_origin.
+
+        :param page_height:
+
+        """
+        if self.coord_origin == CoordOrigin.TOPLEFT:
+            return self
+        elif self.coord_origin == CoordOrigin.BOTTOMLEFT:
+            return BoundingBox(
+                l=self.l,
+                r=self.r,
+                t=page_height - self.t,  # self.b
+                b=page_height - self.b,  # self.t
+                coord_origin=CoordOrigin.TOPLEFT,
+            )