File size: 7,743 Bytes
2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b 51720d8 2ad630b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """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)},
)
|