chunker.py

import os
from typing import Dict, List, Optional, Tuple

from docling.chunking import HybridChunker
from docling_core.transforms.chunker.hierarchical_chunker import (
    ChunkingDocSerializer,
    ChunkingSerializerProvider,
)
from docling_core.transforms.serializer.markdown import MarkdownParams
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc import PictureItem, TableItem

from models import BoundingBox, Chunk, FigureRef


class _ImgPlaceholderSerializerProvider(ChunkingSerializerProvider):
    """Replaces images with a placeholder comment to keep chunk text clean."""

    def get_serializer(self, doc: DoclingDocument) -> ChunkingDocSerializer:
        return ChunkingDocSerializer(
            doc=doc,
            params=MarkdownParams(image_placeholder="<!-- image -->"),
        )


def _extract_figure_ref(
    element, dl_doc: DoclingDocument
) -> Optional[FigureRef]:
    """Build a FigureRef from a PictureItem or TableItem if it has provenance."""
    if isinstance(element, PictureItem):
        kind = "picture"
    elif isinstance(element, TableItem):
        kind = "table"
    else:
        return None

    ref = element.self_ref
    page_number = None
    bbox = None
    caption = ""

    # Extract provenance (page + bbox)
    if hasattr(element, "prov") and element.prov:
        prov = element.prov[0]
        page_number = prov.page_no
        try:
            page = dl_doc.pages[prov.page_no]
            raw_bbox = prov.bbox.to_top_left_origin(page_height=page.size.height)
            raw_bbox = raw_bbox.normalized(page.size)
            bbox = BoundingBox(
                l=raw_bbox.l, r=raw_bbox.r,
                t=raw_bbox.t, b=raw_bbox.b,
                page_no=prov.page_no,
            )
        except (IndexError, AttributeError):
            pass

    # Extract caption if available
    if hasattr(element, "caption_text"):
        caption = element.caption_text(dl_doc) or ""
    elif hasattr(element, "captions") and element.captions:
        parts = []
        for cap in element.captions:
            if hasattr(cap, "text"):
                parts.append(cap.text)
        caption = " ".join(parts)

    return FigureRef(
        ref=ref,
        kind=kind,
        page_number=page_number,
        bbox=bbox,
        caption=caption,
    )


class Chunker:
    """Chunks a DoclingDocument into :class:`~models.Chunk` objects.

    Each chunk carries the original text (contextualized with surrounding
    headings), normalized bounding boxes for every covered page region,
    references to any figures (pictures/tables) that fall within the chunk's
    document range, and the document provenance needed for visual grounding.
    """

    def __init__(
        self,
        dl_doc: DoclingDocument,
        tokenizer: str = "sentence-transformers/all-MiniLM-L6-v2",
    ):
        self.dl_doc = dl_doc
        self.file_hash = str(dl_doc.origin.binary_hash) if dl_doc.origin else ""
        self.filename = dl_doc.origin.filename if dl_doc.origin else ""
        self.file_extension = os.path.splitext(self.filename)[1] if self.filename else ".pdf"

        self._chunker = HybridChunker(
            tokenizer=tokenizer,
            serializer_provider=_ImgPlaceholderSerializerProvider(),
        )

    def _build_figure_index(self) -> Tuple[Dict[str, int], List[Tuple[object, FigureRef]]]:
        """Build a position index of all document items, and collect figure refs.

        Returns:
            (ref_to_pos, figure_entries) where:
            - ref_to_pos maps each element's self_ref to its position in document order
            - figure_entries is a list of (element, FigureRef) for all pictures/tables
        """
        ref_to_pos: Dict[str, int] = {}
        figure_entries: List[Tuple[object, FigureRef]] = []

        for i, (element, _level) in enumerate(self.dl_doc.iterate_items()):
            ref = element.self_ref
            ref_to_pos[ref] = i

            fig_ref = _extract_figure_ref(element, self.dl_doc)
            if fig_ref is not None:
                figure_entries.append((element, fig_ref))
                # Also register the figure's ref in the position map
                if ref not in ref_to_pos:
                    ref_to_pos[ref] = i

        return ref_to_pos, figure_entries

    def chunk(self) -> List[Chunk]:
        """Split the document into chunks with text, headings, bboxes, and figure refs.

        Returns:
            Ordered list of :class:`~models.Chunk` objects.
        """
        raw_chunks = list(self._chunker.chunk(self.dl_doc))
        if not raw_chunks:
            return []

        # Build position index and figure list
        ref_to_pos, figure_entries = self._build_figure_index()

        chunks: List[Chunk] = []
        for raw in raw_chunks:
            index = len(chunks)
            doc_item_refs = [item.self_ref for item in raw.meta.doc_items]
            contextualized_text = self._chunker.contextualize(chunk=raw)

            bboxes: List[BoundingBox] = []
            for doc_item in raw.meta.doc_items:
                if doc_item.prov:
                    for prov in doc_item.prov:
                        try:
                            page = self.dl_doc.pages[prov.page_no]
                            bbox = prov.bbox.to_top_left_origin(page_height=page.size.height)
                            bbox = bbox.normalized(page.size)
                            bboxes.append(
                                BoundingBox(
                                    l=bbox.l, r=bbox.r,
                                    t=bbox.t, b=bbox.b,
                                    page_no=prov.page_no,
                                )
                            )
                        except (IndexError, AttributeError):
                            pass

            page_numbers = sorted(set(b.page_no for b in bboxes)) if bboxes else []

            # Find figures that fall within this chunk's document position range
            figures: List[FigureRef] = []
            chunk_positions = [ref_to_pos[r] for r in doc_item_refs if r in ref_to_pos]
            if chunk_positions:
                pos_start = min(chunk_positions)
                pos_end = max(chunk_positions)
                for _element, fig_ref in figure_entries:
                    fig_pos = ref_to_pos.get(fig_ref.ref)
                    if fig_pos is not None and pos_start <= fig_pos <= pos_end:
                        figures.append(fig_ref)

            # Also catch figures on the same pages that are near the chunk
            # but not strictly between its doc_items (e.g. a figure at the
            # bottom of a page whose text chunk ends mid-page).
            if page_numbers and not figures:
                page_set = set(page_numbers)
                for _element, fig_ref in figure_entries:
                    if fig_ref.page_number in page_set and fig_ref not in figures:
                        # Check if figure is already claimed by checking position
                        fig_pos = ref_to_pos.get(fig_ref.ref)
                        if fig_pos is not None:
                            # Only claim if within a small range beyond the chunk
                            if pos_start - 3 <= fig_pos <= pos_end + 3:
                                figures.append(fig_ref)

            chunks.append(
                Chunk(
                    id=f"{self.file_hash}_{index}",
                    index=index,
                    text=contextualized_text,
                    page_number=page_numbers[0] if page_numbers else None,
                    page_range=page_numbers,
                    headings=getattr(raw.meta, "headings", None) or [],
                    doc_items=doc_item_refs,
                    bboxes=bboxes,
                    figures=figures,
                    document_name=self.dl_doc.name,
                    file_hash=self.file_hash,
                    file_extension=self.file_extension,
                )
            )

        return chunks