first draft toc-based pdf header hierarchy

Author: krrome

docling/models/header_hierarchy/__init__.py

@@ -0,0 +1,153 @@
+import re
+from collections.abc import Generator
+from contextlib import contextmanager
+from functools import cached_property
+from io import BytesIO
+from logging import Logger
+from pathlib import Path, PurePath
+from typing import Optional, Union
+
+from docling_core.types.doc import BoundingBox, DocItemLabel, ListItem, TextItem
+
+from docling.backend.docling_parse_v4_backend import DoclingParseV4DocumentBackend
+from docling.datamodel.base_models import DocumentStream, PageElement
+from docling.datamodel.document import ConversionResult
+from docling.models.readingorder_model import ReadingOrderPageElement
+
+from .types.hierarchical_header import HierarchicalHeader
+
+logger = Logger(__name__)
+
+
+class HeaderNotFoundException(Exception):
+    def __init__(self, heading: str):
+        super().__init__(f"Following heading was not found in the document: {heading}")
+
+
+class ImplausibleHeadingStructureException(Exception):
+    def __init__(self) -> None:
+        super().__init__(
+            "Hierarchy demands equal level heading, but no common parent was found!"
+        )
+
+
+class PdfBackendIncompatible(Exception):
+    def __init__(self, backend) -> None:
+        super().__init__(
+            f"The selected backend is '{type(backend)}' instead of 'DoclingParseV4DocumentBackend'."
+        )
+
+
+class InvalidSourceTypeException(Exception):
+    pass
+
+
+class HierarchyBuilderMetadata:
+    def __init__(
+        self,
+        conv_res: ConversionResult,
+        sorted_elements: list[ReadingOrderPageElement],
+        raise_on_error: bool = False,
+    ):
+        """Correct
+
+        Args:
+            toc_meta (list[dict]): _description_
+            headers (list[PageElement]): _description_
+            raise_on_error (bool, optional): _description_. Defaults to False.
+        """
+        from docling_parse.pdf_parser import DoclingPdfParser
+
+        if not isinstance(conv_res.input._backend, DoclingParseV4DocumentBackend):
+            raise PdfBackendIncompatible(conv_res.input._backend)
+        backend: DoclingParseV4DocumentBackend = conv_res.input._backend
+        pdf_parser: DoclingPdfParser = backend.parser
+        self.toc_meta: list[dict] = pdf_parser.parser.get_table_of_contents(
+            pdf_parser.list_loaded_keys()[0]
+        )
+        self.conv_res: ConversionResult = conv_res
+        self.all_elements: list[PageElement] = conv_res.assembled.elements
+        self.all_cids: list[str] = [
+            f"#/{element.page_no}/{element.cluster.id}" for element in self.all_elements
+        ]
+        self.sorted_ro_elements: list[ReadingOrderPageElement] = sorted_elements
+        self.raise_on_error: bool = raise_on_error
+        self.cid_to_page_element: dict[str, PageElement] = dict(
+            zip(self.all_cids, self.all_elements)
+        )
+
+    def _iterate_toc(self, toc_list: Optional[list[dict]] = None):
+        if toc_list is None:
+            toc_list = self.toc_meta
+        for toc_el in toc_list:
+            yield toc_el["level"], toc_el["title"]
+            if "children" in toc_el:
+                yield from self._iterate_toc(toc_el["children"])
+
+    def infer(self) -> HierarchicalHeader:
+        root = HierarchicalHeader()
+        current = root
+
+        # my problem is that I will need the font information in PdfTextCell, but at the same time I need the ordered text elements (with self refs ideally)
+
+        for level, title in self._iterate_toc():
+            new_parent = None
+            this_element = None
+            orig_text = ""
+            ref = None
+            last_i: int = 0
+            # identify the text item in the document
+            for _i, ro_element in enumerate(self.sorted_ro_elements[last_i:]):
+                element = self.cid_to_page_element[ro_element.ref.cref]
+                # skip all page elements that are before the last ("current") header
+                # if element.page_no < last_page or element.cluster.id <= last_cid:
+                #     continue
+                # Future to do: fixme - better to look for an overlap with the 'to' pointer if possible...
+                from docling.datamodel.base_models import TextElement
+
+                if not isinstance(element, TextElement):
+                    continue
+                orig_text = "".join([c.orig for c in element.cluster.cells])
+
+                if re.sub(r"[^A-Za-z0-9]", "", title) == re.sub(
+                    r"[^A-Za-z0-9]", "", orig_text
+                ):
+                    this_element = element
+                    last_i = last_i + _i
+                    ref = ro_element.ref.cref
+                    break
+            if this_element is None:
+                if self.raise_on_error:
+                    raise HeaderNotFoundException(title)
+                else:
+                    logger.warning(HeaderNotFoundException(title))
+                    continue
+
+            if this_element.label != DocItemLabel.SECTION_HEADER:
+                this_element.label = DocItemLabel.SECTION_HEADER
+
+            if current.level_toc is None or level > current.level_toc:
+                # print(f"gt: {this_fs_level, this_style_attr} VS: {current.level_fontsize, current.style_attrs}")
+                new_parent = current
+            elif level == current.level_toc:
+                # print(f"eq: {this_fs_level, this_style_attr} VS: {current.level_fontsize, current.style_attrs}")
+                if current.parent is not None:
+                    new_parent = current.parent
+                else:
+                    raise ImplausibleHeadingStructureException()
+            else:
+                # go back up in hierarchy and try to find a new parent
+                new_parent = current
+                while new_parent.parent is not None and (level <= new_parent.level_toc):
+                    new_parent = new_parent.parent
+                # print(f"fit parent for : {this_fs_level, this_style_attr} parent: {new_parent.level_fontsize, new_parent.style_attrs}")
+            new_obj = HierarchicalHeader(
+                text=orig_text,
+                parent=new_parent,
+                level_toc=level,
+                doc_ref=ref,
+            )
+            new_parent.children.append(new_obj)
+            current = new_obj
+
+        return root

docling/models/header_hierarchy/metadata_hierarchy.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class NumberingLevel(Enum):
+    level_latin = "level_latin"
+    level_alpha = "level_alpha"
+    level_numerical = "level_numerical"
+
+
+class StyleAttributes(Enum):
+    font_size = "font_size"
+    bold = "bold"
+    italic = "italic"

docling/models/header_hierarchy/style_based_hierarchy.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+from .enums import NumberingLevel, StyleAttributes
+
+
+class UnkownNumberingLevel(Exception):
+    def __init__(self, level_name: NumberingLevel):
+        super().__init__(
+            f"Level kind must be one of {NumberingLevel.__members__.values()}, not '{level_name}'."
+        )
+
+
+@dataclass
+class HierarchicalHeader:
+    index: Optional[int] = None
+    level_toc: Optional[int] = None
+    level_fontsize: Optional[int] = None
+    style_attrs: list[StyleAttributes] = field(default_factory=list)
+    level_latin: list[int] = field(default_factory=list)
+    level_alpha: list[int] = field(default_factory=list)
+    level_numerical: list[int] = field(default_factory=list)
+    parent: Optional["HierarchicalHeader"] = None
+    children: list["HierarchicalHeader"] = field(default_factory=list)
+    doc_ref: Optional[str] = None
+    text: Optional[str] = None
+
+    def __post_init__(self):
+        self._doc_ref_to_parent_doc_ref: dict[Union[str, None], Union[str, None]] = None
+
+    def any_level(self) -> bool:
+        return bool(self.level_alpha or self.level_alpha or self.level_numerical)
+
+    def last_level_of_kind(
+        self, kind: NumberingLevel
+    ) -> tuple[list[int], Union["HierarchicalHeader", None]]:
+        if kind not in NumberingLevel.__members__.values():
+            raise UnkownNumberingLevel(kind)
+        if self.parent:
+            if last := getattr(self.parent, kind.value):
+                return last, self.parent
+            return self.parent.last_level_of_kind(kind)
+        return [], None
+
+    def string_repr(self, prefix: str = "") -> str:
+        out_text = ""
+        if self.text:
+            out_text += prefix + self.text + "\n"
+        for child in self.children:
+            out_text += child.string_repr(prefix + "  ")
+        return out_text
+
+    def __str__(self) -> str:
+        return self.string_repr()
+
+    def _build_doc_ref_to_parent_doc_ref(
+        self,
+    ) -> dict[Union[str, None], Union[str, None]]:
+        self._doc_ref_to_parent_doc_ref = {}
+        for child in self.children:
+            self._doc_ref_to_parent_doc_ref.update(
+                child._build_doc_ref_to_parent_doc_ref()
+            )
+        if self.parent is not None and self.parent.doc_ref is not None:
+            self._doc_ref_to_parent_doc_ref[self.doc_ref] = self.parent.doc_ref
+        else:
+            self._doc_ref_to_parent_doc_ref[self.doc_ref] = None
+        return self._doc_ref_to_parent_doc_ref
+
+    def get_parent_cid_of(self, doc_ref: str) -> Union[str, None]:
+        if self._doc_ref_to_parent_doc_ref is None:
+            self._build_doc_ref_to_parent_doc_ref()
+        return self._doc_ref_to_parent_doc_ref[doc_ref]