Sync from GitHub 7d9b758

README.md

@@ -1,6 +1,6 @@
 ---
 title: Repair Guy
-emoji: ⚡
+emoji: 🔧
 colorFrom: purple
 colorTo: red
 sdk: gradio
@@ -9,9 +9,28 @@ python_version: '3.12'
 app_file: app.py
 pinned: false
 preload_from_hub:
-  - nvidia/NVIDIA-Nemotron-Parse-v1.2
-  - nvidia/C-RADIOv2-H
+  - nvidia/nemotron-colembed-vl-4b-v2
 license: mit
 ---
 
+# Repair Guy — visual RAG over repair manuals
+
+No parsing, no chunking, no figure descriptions: every PDF page is embedded
+as an image with [Nemotron ColEmbed v2](https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2)
+(multi-vector, late interaction). At question time the query is embedded and
+scored against every page with MaxSim (batched torch matmuls on ZeroGPU,
+pages streamed from disk via numpy memmap), and the top-K page images are
+handed to a MiniCPM-V endpoint to produce a grounded answer.
+
+## Space setup
+
+- **Persistent storage** must be enabled (embeddings + PDFs live under
+  `/data/library`). Budget roughly 5–12 MB per page of float16 token
+  embeddings; a 300-page manual is ~2–3.5 GB.
+- **Secret `MINICPM_API_KEY`** — bearer token for the MiniCPM endpoint
+  (`MINICPM_BASE_URL` / `MINICPM_MODEL` are env-overridable).
+- Optional: `COLEMBED_MODEL_ID` (defaults to the 4B model),
+  `COLEMBED_ATTN` (defaults to `sdpa`; set `flash_attention_2` if flash-attn
+  is installed).
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,180 +1,95 @@
-"""Gradio + ZeroGPU Space for NVIDIA Nemotron Parse v1.2.
-
-Upload a PDF, pick a page, and get back the parsed markdown, a structured JSON of
-elements, and the page image annotated with bounding boxes.
-
-Runs on ZeroGPU: the model is loaded onto cuda at module level (ZeroGPU emulates
-CUDA at startup) and inference runs inside an @spaces.GPU-decorated function.
-
-This file targets the Space (cuda/bfloat16). For local CPU testing use
-parse_page.py in the repo root instead.
+"""Gradio + ZeroGPU Space: upload a repair manual, ask questions.
+
+Visual RAG with no parsing/chunking: every PDF page is embedded as an image
+with Nemotron ColEmbed v2 (multi-vector, late interaction) when the manual is
+uploaded. A question is answered by embedding the query, running MaxSim over
+the page embeddings streamed from disk, and handing the top pages to MiniCPM-V.
+
+Module layout:
+  models/colembed.py   ColEmbed     — embedding model + GPU embed/search
+  models/minicpm.py    MiniCPM      — remote VLM that answers over page images
+  core/store.py        Store        — on-disk per-page token embeddings
+  core/pdf.py          render_pages — PDF -> RGB page images (CPU)
+  pipelines/ingest.py  IngestPipeline — PDF -> embeddings -> store
+  pipelines/ask.py     AskPipeline    — question -> retrieve -> answer
+  app.py               this file: builds the objects + Gradio UI
 """
 
-import json
-import sys
+import os
 
-import fitz  # pymupdf
 import gradio as gr
-import spaces
-import torch
-from huggingface_hub import snapshot_download
-from PIL import Image, ImageDraw
-from transformers import AutoModel, AutoProcessor, GenerationConfig
-
-MODEL_ID = "nvidia/NVIDIA-Nemotron-Parse-v1.2"
-DEVICE = "cuda"
-DTYPE = torch.bfloat16
-MAX_PROMPT_DURATION = 120  # seconds of GPU time per page
-
-# ---------------------------------------------------------------------------
-# Load helpers + model once at module level (ZeroGPU loads cuda weights here).
-# ---------------------------------------------------------------------------
-
-
-def load_postprocessing():
-    """Download the repo's .py helpers and import postprocessing.
-
-    postprocessing.py imports sibling modules (latex2html, ...), so we pull all
-    top-level .py files into one dir and put it on sys.path before importing.
-    """
-    repo_dir = snapshot_download(repo_id=MODEL_ID, allow_patterns=["*.py"])
-    if repo_dir not in sys.path:
-        sys.path.insert(0, repo_dir)
-    import postprocessing  # noqa: E402  (resolved via sys.path above)
-
-    return postprocessing
-
-
-pp = load_postprocessing()
-
-# Every load passes trust_remote_code=True so the nested C-RADIO encoder code is
-# accepted non-interactively (no [y/N] prompt to hang the Space build).
-model = (
-    AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, dtype=DTYPE)
-    .to(DEVICE)
-    .eval()
-)
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-generation_config = GenerationConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-
-@spaces.GPU(duration=MAX_PROMPT_DURATION)
-def run_model(image: Image.Image, task_prompt: str) -> str:
-    """GPU-only step: preprocess + generate + decode. Returns raw model text."""
-    inputs = processor(
-        images=[image], text=task_prompt, return_tensors="pt", add_special_tokens=False
-    )
-    # Move to GPU; cast float tensors (pixel_values) to the model dtype.
-    inputs = {
-        k: (v.to(DEVICE, DTYPE) if torch.is_floating_point(v) else v.to(DEVICE))
-        for k, v in inputs.items()
-    }
-    with torch.no_grad():
-        outputs = model.generate(**inputs, generation_config=generation_config)
-    return processor.batch_decode(outputs, skip_special_tokens=True)[0]
-
-
-# ---------------------------------------------------------------------------
-# CPU-side orchestration: render page, call GPU, postprocess, annotate.
-# ---------------------------------------------------------------------------
 
-
-def render_page(pdf_path: str, page_num: int, dpi: int) -> Image.Image:
-    doc = fitz.open(pdf_path)
+from core.constants import DEFAULT_TOP_K
+from core.store import Store, slugify
+from models.colembed import ColEmbed
+from models.minicpm import MiniCPM
+from pipelines.ask import AskPipeline
+from pipelines.ingest import IngestPipeline
+
+# Construct once at startup (the model loads onto cuda here, in the main process).
+store = Store()
+embedder = ColEmbed()
+ingest_pipeline = IngestPipeline(embedder, store)
+ask_pipeline = AskPipeline(embedder, store, MiniCPM())
+
+
+def index_manual(pdf_file, progress=gr.Progress()):
+    """Runs on upload: embed the manual's pages (or reuse a previous index)."""
+    if not pdf_file:
+        return None, ""
+    name = os.path.splitext(os.path.basename(pdf_file))[0].replace("_", " ")
+    doc_id = slugify(name)
+    if store.exists(doc_id):
+        pages = len(store.meta(doc_id)["pages"])
+        return doc_id, f"**{name}** is already indexed ({pages} pages) — ask away."
     try:
-        if page_num < 1 or page_num > doc.page_count:
-            raise gr.Error(
-                f"Page {page_num} out of range — this PDF has {doc.page_count} pages."
-            )
-        pix = doc.load_page(page_num - 1).get_pixmap(dpi=dpi)
-        return Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
-    finally:
-        doc.close()
-
-
-def load_input(file_path: str, page_num: int, dpi: int) -> Image.Image:
-    """Return an RGB image from either a PDF page or an image file."""
-    if file_path.lower().endswith(".pdf"):
-        return render_page(file_path, page_num, dpi)
-    return Image.open(file_path).convert("RGB")
-
-
-def parse(input_file, page_num, dpi, text_in_pic, table_format):
-    if input_file is None:
-        raise gr.Error("Please upload a PDF or image first.")
-
-    image = load_input(input_file, int(page_num), int(dpi))
-
-    fourth = "<predict_text_in_pic>" if text_in_pic else "<predict_no_text_in_pic>"
-    task_prompt = f"</s><s><predict_bbox><predict_classes><output_markdown>{fourth}"
-
-    generated_text = run_model(image, task_prompt)
+        doc = ingest_pipeline.run(
+            pdf_file, name, lambda frac, desc: progress(frac, desc=desc)
+        )
+    except ValueError as e:
+        raise gr.Error(str(e)) from e
+    return doc_id, f"Indexed **{doc['name']}** ({doc['pages']} pages) — ask away."
 
-    classes, bboxes, texts = pp.extract_classes_bboxes(generated_text)
-    bboxes = [pp.transform_bbox_to_original(b, image.width, image.height) for b in bboxes]
-    texts = [
-        pp.postprocess_text(t, cls=c, table_format=table_format, text_format="markdown")
-        for t, c in zip(texts, classes)
-    ]
-
-    markdown = "\n\n".join(texts)
-    elements = [
-        {"class": c, "bbox": b, "text": t} for c, b, t in zip(classes, bboxes, texts)
-    ]
-
-    annotated = image.copy()
-    draw = ImageDraw.Draw(annotated)
-    for b in bboxes:
-        draw.rectangle((b[0], b[1], b[2], b[3]), outline="red", width=2)
-
-    return annotated, markdown, json.dumps(elements, indent=2)
 
+def answer_question(question, doc_id):
+    if not doc_id:
+        raise gr.Error("Upload a repair manual first.")
+    try:
+        return ask_pipeline.run(question, [doc_id], DEFAULT_TOP_K)
+    except ValueError as e:
+        raise gr.Error(str(e)) from e
 
-# ---------------------------------------------------------------------------
-# UI
-# ---------------------------------------------------------------------------
 
-with gr.Blocks(title="Nemotron Parse — Repair Manuals") as demo:
+with gr.Blocks(title="Repair Guy") as demo:
     gr.Markdown(
-        "# 🔧 Nemotron Parse v1.2 — Repair Manual Explorer\n"
-        "Upload a PDF (choose a page) or an image, and parse it with "
-        "[NVIDIA Nemotron Parse v1.2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2) "
-        "on ZeroGPU. Returns structured markdown, a JSON of elements, and an "
-        "annotated page image."
+        "# 🔧 Repair Guy\n"
+        "Upload a repair manual (PDF) and ask it questions. Pages are embedded "
+        "with [Nemotron ColEmbed v2](https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2) "
+        "on upload; answers come from MiniCPM-V reading the most relevant pages."
     )
+    doc_state = gr.State(None)
     with gr.Row():
         with gr.Column(scale=1):
             pdf_in = gr.File(
-                label="PDF or image",
-                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
-                type="filepath",
+                label="Repair manual (PDF)", file_types=[".pdf"], type="filepath"
             )
-            page_in = gr.Number(
-                label="Page (PDF only)", value=1, precision=0, minimum=1
+            status_out = gr.Markdown()
+            question_in = gr.Textbox(
+                label="Question",
+                lines=2,
+                placeholder="e.g. What is the tightening torque for the universal joint flange bolts?",
             )
-            dpi_in = gr.Slider(
-                label="Render DPI (PDF only)", minimum=72, maximum=300, value=150, step=10
-            )
-            text_in_pic_in = gr.Checkbox(
-                label="Extract text inside pictures/diagrams", value=False
-            )
-            table_format_in = gr.Dropdown(
-                label="Table format",
-                choices=["markdown", "latex", "HTML", "json", "csv"],
-                value="markdown",
-            )
-            run_btn = gr.Button("Parse page", variant="primary")
+            ask_btn = gr.Button("Ask", variant="primary")
         with gr.Column(scale=2):
-            img_out = gr.Image(label="Annotated page", type="pil")
-            with gr.Tab("Rendered markdown"):
-                md_out = gr.Markdown()
-            with gr.Tab("Structured JSON"):
-                json_out = gr.Code(language="json")
+            answer_out = gr.Markdown(label="Answer")
+            pages_out = gr.Gallery(label="Pages used", columns=3, height=420)
 
-    run_btn.click(
-        parse,
-        inputs=[pdf_in, page_in, dpi_in, text_in_pic_in, table_format_in],
-        outputs=[img_out, md_out, json_out],
+    pdf_in.upload(index_manual, inputs=[pdf_in], outputs=[doc_state, status_out])
+    ask_btn.click(
+        answer_question, inputs=[question_in, doc_state], outputs=[answer_out, pages_out]
+    )
+    question_in.submit(
+        answer_question, inputs=[question_in, doc_state], outputs=[answer_out, pages_out]
     )
 
 

core/constants.py

@@ -0,0 +1,34 @@
+import os
+
+# Embedding model (late-interaction / ColBERT-style page embeddings).
+COLEMBED_MODEL_ID = os.environ.get(
+    "COLEMBED_MODEL_ID", "nvidia/nemotron-colembed-vl-4b-v2"
+)
+# sdpa works for this model and needs no extra wheels on ZeroGPU; set
+# COLEMBED_ATTN=flash_attention_2 if flash-attn is installed.
+COLEMBED_ATTN = os.environ.get("COLEMBED_ATTN", "sdpa")
+
+# Page rendering (both for embedding at index time and for the answering model).
+RENDER_DPI = 150
+
+# Indexing: pages embedded per ZeroGPU call, and model batch size within a call.
+# Chunking keeps each GPU call well under its duration limit; progress is
+# reported between chunks.
+EMBED_PAGES_PER_CALL = 16
+EMBED_BATCH_SIZE = 4
+EMBED_GPU_DURATION = 120
+
+# Retrieval: MaxSim is computed on GPU over fixed-size batches of pages
+# streamed from the on-disk store.
+SCORE_PAGES_PER_BATCH = 32
+SEARCH_GPU_DURATION = 60
+DEFAULT_TOP_K = 3
+MAX_TOP_K = 5
+
+# Embedding store. HF Spaces persistent storage mounts at /data; fall back to a
+# local directory for development.
+STORE_DIR = os.environ.get("STORE_DIR") or (
+    "/data/library"
+    if os.path.isdir("/data")
+    else os.path.join(os.path.dirname(os.path.dirname(__file__)), "library")
+)

core/pdf.py

@@ -0,0 +1,35 @@
+import fitz
+from PIL import Image
+
+from core.constants import RENDER_DPI
+
+
+def page_count(pdf_path: str) -> int:
+    doc = fitz.open(pdf_path)
+    try:
+        return doc.page_count
+    finally:
+        doc.close()
+
+
+def render_pages(
+    pdf_path: str, page_nums: list[int], dpi: int = RENDER_DPI
+) -> list[Image.Image]:
+    """Render 1-based pages of a PDF to RGB images (document opened once)."""
+    doc = fitz.open(pdf_path)
+    try:
+        images = []
+        for num in page_nums:
+            if num < 1 or num > doc.page_count:
+                raise ValueError(
+                    f"Page {num} out of range — this PDF has {doc.page_count} pages."
+                )
+            pix = doc.load_page(num - 1).get_pixmap(dpi=dpi)
+            images.append(Image.frombytes("RGB", (pix.width, pix.height), pix.samples))
+        return images
+    finally:
+        doc.close()
+
+
+def render_page(pdf_path: str, page_num: int, dpi: int = RENDER_DPI) -> Image.Image:
+    return render_pages(pdf_path, [page_num], dpi)[0]

core/store.py

@@ -0,0 +1,142 @@
+"""On-disk store of per-page ColEmbed token embeddings.
+
+Layout: STORE_DIR/<doc_id>/
+    doc.pdf          original PDF, kept to re-render retrieved pages at answer time
+    embeddings.bin   raw float16 [total_tokens, dim], pages stored back to back
+    index.json       manual name, dim, dpi, model id, per-page (offset, count)
+
+embeddings.bin is read back as a numpy memmap, so scoring streams pages from
+disk in batches without ever loading a whole document into memory. index.json
+is written last, so a directory without it is an interrupted ingest and is
+ignored (and overwritten on the next attempt).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+
+import numpy as np
+
+from core.constants import STORE_DIR
+
+EMB_DTYPE = np.float16
+
+
+def slugify(name: str) -> str:
+    slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
+    return slug or "doc"
+
+
+class DocWriter:
+    """Streams one document's page embeddings to disk during ingest."""
+
+    def __init__(self, doc_dir: str, name: str, pdf_path: str, dpi: int, model_id: str):
+        self.doc_dir = doc_dir
+        os.makedirs(doc_dir, exist_ok=True)
+        shutil.copyfile(pdf_path, os.path.join(doc_dir, "doc.pdf"))
+        self._bin = open(os.path.join(doc_dir, "embeddings.bin"), "wb")
+        self._meta = {
+            "name": name,
+            "dpi": dpi,
+            "model_id": model_id,
+            "dim": None,
+            "pages": [],
+        }
+        self._offset = 0
+
+    def add_page(self, page_num: int, emb: np.ndarray) -> None:
+        """emb: [n_tokens, dim] for one page, padding rows already removed."""
+        emb = np.ascontiguousarray(emb, dtype=EMB_DTYPE)
+        if self._meta["dim"] is None:
+            self._meta["dim"] = int(emb.shape[1])
+        self._bin.write(emb.tobytes())
+        self._meta["pages"].append(
+            {"page": page_num, "offset": self._offset, "count": int(emb.shape[0])}
+        )
+        self._offset += int(emb.shape[0])
+
+    def finalize(self) -> None:
+        self._bin.close()
+        with open(os.path.join(self.doc_dir, "index.json"), "w") as f:
+            json.dump(self._meta, f)
+
+    def abort(self) -> None:
+        self._bin.close()
+        shutil.rmtree(self.doc_dir, ignore_errors=True)
+
+
+class Store:
+    def __init__(self, root: str = STORE_DIR):
+        self.root = os.path.abspath(root)
+        os.makedirs(self.root, exist_ok=True)
+
+    def _dir(self, doc_id: str) -> str:
+        return os.path.join(self.root, doc_id)
+
+    def exists(self, doc_id: str) -> bool:
+        return os.path.isfile(os.path.join(self._dir(doc_id), "index.json"))
+
+    def meta(self, doc_id: str) -> dict:
+        with open(os.path.join(self._dir(doc_id), "index.json")) as f:
+            return json.load(f)
+
+    def pdf_path(self, doc_id: str) -> str:
+        return os.path.join(self._dir(doc_id), "doc.pdf")
+
+    def list_docs(self) -> list[dict]:
+        docs = []
+        for doc_id in sorted(os.listdir(self.root)):
+            if not self.exists(doc_id):
+                continue
+            meta = self.meta(doc_id)
+            size = os.path.getsize(os.path.join(self._dir(doc_id), "embeddings.bin"))
+            docs.append(
+                {
+                    "doc_id": doc_id,
+                    "name": meta["name"],
+                    "pages": len(meta["pages"]),
+                    "size_mb": size / 1e6,
+                }
+            )
+        return docs
+
+    def create(self, doc_id: str, name: str, pdf_path: str, dpi: int, model_id: str) -> DocWriter:
+        self.delete(doc_id)
+        return DocWriter(self._dir(doc_id), name, pdf_path, dpi, model_id)
+
+    def delete(self, doc_id: str) -> None:
+        shutil.rmtree(self._dir(doc_id), ignore_errors=True)
+
+    def iter_page_batches(self, doc_ids: list[str] | None = None, pages_per_batch: int = 32):
+        """Yield (refs, embs): refs is [(doc_id, page_num)] and embs is a
+        zero-padded float16 array [batch, max_tokens, dim].
+
+        Zero rows are inert under MaxSim (the model zeroes padding before
+        L2-normalizing real tokens), matching the reference scorer.
+        """
+        if doc_ids is None:
+            doc_ids = [d["doc_id"] for d in self.list_docs()]
+        for doc_id in doc_ids:
+            if not self.exists(doc_id):  # e.g. deleted while still selected in the UI
+                continue
+            meta = self.meta(doc_id)
+            pages, dim = meta["pages"], meta["dim"]
+            if not pages:
+                continue
+            total = pages[-1]["offset"] + pages[-1]["count"]
+            mm = np.memmap(
+                os.path.join(self._dir(doc_id), "embeddings.bin"),
+                dtype=EMB_DTYPE,
+                mode="r",
+                shape=(total, dim),
+            )
+            for i in range(0, len(pages), pages_per_batch):
+                chunk = pages[i : i + pages_per_batch]
+                t_max = max(p["count"] for p in chunk)
+                out = np.zeros((len(chunk), t_max, dim), dtype=EMB_DTYPE)
+                for j, p in enumerate(chunk):
+                    out[j, : p["count"]] = mm[p["offset"] : p["offset"] + p["count"]]
+                yield [(doc_id, p["page"]) for p in chunk], out

models/colembed.py

@@ -0,0 +1,82 @@
+"""Nemotron ColEmbed v2: late-interaction page embeddings + MaxSim retrieval.
+
+Two ZeroGPU entry points:
+  _embed_pages_on_gpu   page images -> per-page token embeddings (index time)
+  _search_on_gpu        question -> top-K (doc, page, score) via MaxSim over
+                        batches of page embeddings streamed from the store
+
+forward_images/forward_queries return zero-padded [batch, tokens, dim] tensors
+with real tokens L2-normalized, so padding rows are exactly zero. We strip them
+before storing and rely on the same property when scoring zero-padded batches.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import spaces
+import torch
+from PIL import Image
+from transformers import AutoModel
+
+from core.constants import (
+    COLEMBED_ATTN,
+    COLEMBED_MODEL_ID,
+    EMBED_BATCH_SIZE,
+    EMBED_GPU_DURATION,
+    SCORE_PAGES_PER_BATCH,
+    SEARCH_GPU_DURATION,
+)
+
+
+@spaces.GPU(duration=EMBED_GPU_DURATION)
+def _embed_pages_on_gpu(model, images: list[Image.Image]) -> list[np.ndarray]:
+    with torch.no_grad():
+        embs = model.forward_images(images, batch_size=EMBED_BATCH_SIZE)
+    out = []
+    for emb in embs:  # [tokens, dim]; zero rows are padding
+        mask = emb.abs().sum(dim=-1) > 0
+        out.append(emb[mask].to(torch.float16).cpu().numpy())
+    return out
+
+
+@spaces.GPU(duration=SEARCH_GPU_DURATION)
+def _search_on_gpu(model, question: str, store, doc_ids, top_k: int):
+    results = []
+    with torch.no_grad():
+        q = model.forward_queries([question], batch_size=1)[0].to(torch.float16)
+        for refs, batch in store.iter_page_batches(doc_ids, SCORE_PAGES_PER_BATCH):
+            emb = torch.from_numpy(batch).to(q.device)  # [B, T, D] float16
+            sim = torch.einsum("qd,btd->bqt", q, emb).float()
+            scores = sim.amax(dim=2).sum(dim=1)  # MaxSim: max over doc tokens, sum over query tokens
+            results.extend(
+                (doc_id, page, s)
+                for (doc_id, page), s in zip(refs, scores.tolist())
+            )
+    results.sort(key=lambda r: r[2], reverse=True)
+    return results[:top_k]
+
+
+class ColEmbed:
+    MODEL_ID = COLEMBED_MODEL_ID
+
+    def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
+        self.model = (
+            AutoModel.from_pretrained(
+                self.MODEL_ID,
+                trust_remote_code=True,
+                torch_dtype=dtype,
+                attn_implementation=COLEMBED_ATTN,
+            )
+            .to(device)
+            .eval()
+        )
+
+    def embed_pages(self, images: list[Image.Image]) -> list[np.ndarray]:
+        """Embed page images -> list of [n_tokens, dim] float16 arrays."""
+        return _embed_pages_on_gpu(self.model, images)
+
+    def search(
+        self, question: str, store, doc_ids: list[str] | None, top_k: int
+    ) -> list[tuple[str, int, float]]:
+        """Return the top_k (doc_id, page_num, score) across the given docs."""
+        return _search_on_gpu(self.model, question, store, doc_ids, top_k)

models/minicpm.py

@@ -0,0 +1,97 @@
+"""Client for a MiniCPM-V OpenAI-compatible vision endpoint: answers a question
+grounded in the retrieved repair-manual pages."""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import os
+import time
+import urllib.error
+import urllib.request
+
+from PIL import Image
+
+BASE_URL = os.environ.get("MINICPM_BASE_URL", "http://35.203.155.71:8003").rstrip("/")
+MODEL = os.environ.get("MINICPM_MODEL", "MiniCPM-V-4.6")
+API_KEY = os.environ.get("MINICPM_API_KEY", "")
+MAX_EDGE = 1024  # downscale images; endpoint max_model_len is only 8192 and load-sensitive
+RETRIES = 3
+
+PROMPT = (
+    "You are a repair-manual assistant. The images are the manual pages most "
+    "relevant to the user's question, each preceded by its label (manual name "
+    "and page number).\n\n"
+    "Answer the question using ONLY these pages. Quote exact values (torques, "
+    "clearances, part numbers, capacities) as printed, and cite the page label "
+    "for each fact. If the pages do not contain the answer, say so instead of "
+    "guessing.\n\nQuestion: {question}"
+)
+
+
+class MiniCPM:
+    def __init__(
+        self,
+        base_url: str = BASE_URL,
+        model: str = MODEL,
+        api_key: str = API_KEY,
+        max_edge: int = MAX_EDGE,
+        retries: int = RETRIES,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+        self.api_key = api_key
+        self.max_edge = max_edge
+        self.retries = retries
+
+    def answer(self, question: str, pages: list[tuple[str, Image.Image]]) -> str:
+        """pages: [(label, page image)] in retrieval order."""
+        if not self.api_key:
+            raise ValueError(
+                "MINICPM_API_KEY is not set — add it as a secret on the Space."
+            )
+        content = [{"type": "text", "text": PROMPT.format(question=question)}]
+        for label, img in pages:
+            content.append({"type": "text", "text": f"\n[{label}]"})
+            content.append(
+                {"type": "image_url", "image_url": {"url": self._data_uri(img)}}
+            )
+        return self._chat(content, max_tokens=900).strip()
+
+    def _data_uri(self, img: Image.Image) -> str:
+        im = img.convert("RGB")
+        w, h = im.size
+        if max(w, h) > self.max_edge:
+            s = self.max_edge / max(w, h)
+            im = im.resize((int(w * s), int(h * s)))
+        buf = io.BytesIO()
+        im.save(buf, format="JPEG", quality=88)
+        return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
+
+    def _chat(self, content: list[dict], max_tokens: int = 512) -> str:
+        body = json.dumps(
+            {
+                "model": self.model,
+                "temperature": 0,
+                "max_tokens": max_tokens,
+                "messages": [{"role": "user", "content": content}],
+            }
+        ).encode()
+        req = urllib.request.Request(
+            f"{self.base_url}/v1/chat/completions",
+            data=body,
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+        )
+        last = None
+        for attempt in range(self.retries):
+            try:
+                with urllib.request.urlopen(req, timeout=120) as r:
+                    return json.loads(r.read())["choices"][0]["message"]["content"]
+            except (urllib.error.URLError, TimeoutError, OSError) as ex:
+                last = ex
+                time.sleep(2 * (attempt + 1))  # 2s, 4s backoff between retries
+        raise last

pipelines/ask.py

@@ -0,0 +1,35 @@
+"""Ask pipeline: question -> MaxSim retrieval over the store -> top-K page
+images -> MiniCPM answer grounded in those pages."""
+
+from __future__ import annotations
+
+from core.pdf import render_page
+from core.store import Store
+from models.colembed import ColEmbed
+from models.minicpm import MiniCPM
+
+
+class AskPipeline:
+    def __init__(self, embedder: ColEmbed, store: Store, llm: MiniCPM):
+        self.embedder = embedder
+        self.store = store
+        self.llm = llm
+
+    def run(self, question: str, doc_ids: list[str] | None, top_k: int):
+        """Return (answer markdown, gallery items [(image, caption)])."""
+        question = (question or "").strip()
+        if not question:
+            raise ValueError("Please enter a question.")
+        docs = self.store.list_docs()
+        if not docs:
+            raise ValueError("No manuals indexed yet — add one in the Library tab.")
+
+        hits = self.embedder.search(question, self.store, doc_ids or None, int(top_k))
+        names = {d["doc_id"]: d["name"] for d in docs}
+        pages = [
+            (f"{names[doc_id]} — p.{page}", render_page(self.store.pdf_path(doc_id), page), score)
+            for doc_id, page, score in hits
+        ]
+        answer = self.llm.answer(question, [(label, img) for label, img, _ in pages])
+        gallery = [(img, f"{label} (score {score:.1f})") for label, img, score in pages]
+        return answer, gallery

pipelines/ingest.py

@@ -0,0 +1,49 @@
+"""Ingest pipeline: PDF -> page images -> ColEmbed embeddings -> on-disk store.
+
+Pages are embedded in chunks of EMBED_PAGES_PER_CALL so each ZeroGPU call stays
+short; progress is reported between chunks.
+"""
+
+from __future__ import annotations
+
+import os
+
+from core.constants import EMBED_PAGES_PER_CALL, RENDER_DPI
+from core.pdf import page_count, render_pages
+from core.store import Store, slugify
+from models.colembed import ColEmbed
+
+
+class IngestPipeline:
+    def __init__(self, embedder: ColEmbed, store: Store):
+        self.embedder = embedder
+        self.store = store
+
+    def run(self, pdf_path: str | None, doc_name: str = "", progress=None) -> dict:
+        """Index one PDF; returns the stored doc's summary. Re-indexing a manual
+        with the same name overwrites it."""
+        if not pdf_path:
+            raise ValueError("Please upload a PDF first.")
+        if not pdf_path.lower().endswith(".pdf"):
+            raise ValueError("Only PDFs can be indexed.")
+
+        name = doc_name.strip() or (
+            os.path.splitext(os.path.basename(pdf_path))[0].replace("_", " ")
+        )
+        doc_id = slugify(name)
+        total = page_count(pdf_path)
+
+        writer = self.store.create(doc_id, name, pdf_path, RENDER_DPI, self.embedder.MODEL_ID)
+        try:
+            for start in range(1, total + 1, EMBED_PAGES_PER_CALL):
+                nums = list(range(start, min(start + EMBED_PAGES_PER_CALL, total + 1)))
+                images = render_pages(pdf_path, nums)
+                for num, emb in zip(nums, self.embedder.embed_pages(images)):
+                    writer.add_page(num, emb)
+                if progress:
+                    progress(nums[-1] / total, f"Embedded {nums[-1]}/{total} pages")
+            writer.finalize()
+        except BaseException:
+            writer.abort()
+            raise
+        return {"doc_id": doc_id, "name": name, "pages": total}

requirements.txt

@@ -1,12 +1,8 @@
 spaces
 gradio
-transformers==5.6.1
+transformers>=4.57.2,<5
 accelerate
-albumentations
-timm
-open_clip_torch
-einops
-beautifulsoup4
-lxml
+torchvision
 pymupdf
 pillow
+numpy