pipelines/parsed_ingest.py

"""Parsed ingest pipeline: PDF -> Nemotron Parse -> MiniCPM descriptions ->
section chunks -> Nemotron Embed -> store.

Two stages that CANNOT share a process: Nemotron Parse requires
transformers==5.6.1 while MiniCPM/Nemotron Embed need <5. scripts/index_modal.py
runs each stage in its own Modal container (separate images); the parse
output (plain dicts) is what crosses the boundary, and it is also persisted
as parsed.json in the store so later re-chunking needs no re-parse.

Model imports happen lazily inside each stage's __init__ — importing this
module must not pull in either transformers world.
"""

from __future__ import annotations

import os
from collections.abc import Iterator

from core.chunking import (
    CAPTION_CLASS,
    FIGURE_CLASS,
    HEADING_CLASSES,
    SKIP_CLASSES,
    TABLE_CLASS,
    build_chunks,
)
from core.constants import (
    DESCRIBE_BATCH_SIZE,
    DESCRIBE_CONTEXT_MAX_CHARS,
    FIGURE_MIN_SIDE_PX,
    PARSE_BATCH_SIZE,
    RENDER_DPI,
)
from core.parsed_store import ParsedStore
from core.pdf import page_count, render_page, render_pages
from core.store import slugify


class ParseStage:
    """Stage 1 (transformers 5.x env): render pages and parse them into
    classified elements, PARSE_BATCH_SIZE pages per generate call."""

    def __init__(self):
        from models.nemotron_parse import NemotronParse

        self.parser = NemotronParse()

    def run(self, pdf_path: str) -> Iterator[tuple]:
        """Generator yielding ("progress", pages_done, total) per parsed
        batch, then ("done", pages) where pages = [{"page", "elements"}] in
        order."""
        total = page_count(pdf_path)
        pages = []
        for start in range(1, total + 1, PARSE_BATCH_SIZE):
            nums = list(range(start, min(start + PARSE_BATCH_SIZE, total + 1)))
            images = render_pages(pdf_path, nums, RENDER_DPI)
            for num, elements in zip(nums, self.parser.parse_pages(images)):
                pages.append({"page": num, "elements": elements})
            yield ("progress", nums[-1], total)
        yield ("done", pages)


class BuildStage:
    """Stage 2 (transformers 4.x env, same as the Space): describe figures and
    tables with MiniCPM, build section chunks, embed them, save to the store."""

    def __init__(self, store: ParsedStore):
        from models import minicpm
        from models.nemotron_embed import NemotronEmbed

        self.describer = minicpm
        self.embedder = NemotronEmbed()
        self.store = store

    @staticmethod
    def _element_context(doc_name: str, heading: str, elements: list[dict], i: int) -> str:
        """Document context for describing element i: the manual's name, the
        section heading in force, the adjacent Caption (Parse keeps captions
        next to their figure/table in reading order), and the page's other
        text. This is what lets MiniCPM use the manual's own terminology
        instead of guessing from pixels."""
        lines = [f"Manual: {doc_name}"]
        if heading:
            lines.append(f"Section: {heading}")
        for j in (i + 1, i - 1):
            if 0 <= j < len(elements) and elements[j]["class"] == CAPTION_CLASS:
                caption = (elements[j].get("text") or "").strip()
                if caption:
                    lines.append(f"Caption: {caption}")
                break
        page_text = " ".join(
            t
            for el in elements
            if el["class"] not in SKIP_CLASSES | {FIGURE_CLASS, TABLE_CLASS}
            for t in [(el.get("text") or "").strip()]
            if t
        )
        if page_text:
            lines.append(f"Text on the page: {page_text[:DESCRIBE_CONTEXT_MAX_CHARS]}")
        return "\n".join(lines)

    def _page_jobs(
        self, pdf_path: str, pg: dict, doc_name: str, heading: str
    ) -> tuple[str, list[tuple]]:
        """Collect (element, kind, payload, context) describe jobs for this
        page's Picture/Table elements. Tracks and returns the section heading
        in force so the next page's context carries it (sections span pages)."""
        image = None
        jobs = []
        elements = pg["elements"]
        prev_was_heading = False
        for i, el in enumerate(elements):
            if el["class"] in HEADING_CLASSES:
                text = (el.get("text") or "").strip()
                if text:
                    # consecutive headings combine into a breadcrumb, like the chunker
                    heading = f"{heading} — {text}" if prev_was_heading and heading else text
                    prev_was_heading = True
                continue
            prev_was_heading = False
            if el["class"] == FIGURE_CLASS:
                x1, y1, x2, y2 = el.get("bbox") or (0, 0, 0, 0)
                if x2 - x1 < FIGURE_MIN_SIDE_PX or y2 - y1 < FIGURE_MIN_SIDE_PX:
                    continue  # icons, bullets, print artifacts
                if image is None:
                    image = render_page(pdf_path, pg["page"], RENDER_DPI)
                crop = image.crop(
                    (max(x1, 0), max(y1, 0), min(x2, image.width), min(y2, image.height))
                )
                jobs.append(
                    (el, "figure", crop, self._element_context(doc_name, heading, elements, i))
                )
            elif el["class"] == TABLE_CLASS and el.get("text", "").strip():
                jobs.append(
                    (el, "table", el["text"], self._element_context(doc_name, heading, elements, i))
                )
        return heading, jobs

    def _describe(self, batch: list[tuple]) -> None:
        """Run one batched MiniCPM call and attach the descriptions to the
        batch's elements (in place)."""
        descriptions = self.describer.describe_batch(
            [(kind, payload, context) for _, kind, payload, context in batch]
        )
        for (el, _, _, _), description in zip(batch, descriptions):
            el["description"] = description

    def run(self, pdf_path: str, parsed_pages: list[dict], doc_name: str = "") -> Iterator[tuple]:
        """Generator yielding ("progress", pages_done, total) through the
        describe phase, then ("done", doc summary dict) after embed + save."""
        name = doc_name.strip() or (
            os.path.splitext(os.path.basename(pdf_path))[0].replace("_", " ")
        )
        doc_id = slugify(name)
        total = len(parsed_pages)

        # Jobs are collected page by page (headings carry across pages) and
        # flushed through MiniCPM in batches; pending never grows past one
        # batch plus one page's worth of elements.
        heading = ""
        pending = []
        for done, pg in enumerate(parsed_pages, start=1):
            heading, jobs = self._page_jobs(pdf_path, pg, name, heading)
            pending.extend(jobs)
            while len(pending) >= DESCRIBE_BATCH_SIZE:
                self._describe(pending[:DESCRIBE_BATCH_SIZE])
                del pending[:DESCRIBE_BATCH_SIZE]
            yield ("progress", done, total)
        if pending:
            self._describe(pending)

        chunks = build_chunks(parsed_pages)
        if not chunks:
            raise ValueError("Parsing produced no text — nothing to index.")
        embeddings = self.embedder.embed_texts([c["text"] for c in chunks])
        self.store.save(
            doc_id,
            name,
            pdf_path,
            chunks,
            embeddings,
            page_count=total,
            dpi=RENDER_DPI,
            model_id=self.embedder.MODEL_ID,
            parsed_pages=parsed_pages,
        )
        yield (
            "done",
            {"doc_id": doc_id, "name": name, "pages": total, "chunks": len(chunks)},
        )