Sync from GitHub 4fcc227

README.md

@@ -24,13 +24,24 @@ pages streamed from disk via numpy memmap), and the top-K page images are
 read by MiniCPM-V 4.5 — also on ZeroGPU — to produce a grounded answer.
 Everything runs inside the Space; no external endpoints.
 
+## Two tabs
+
+- **Library** — large manuals indexed offline (`scripts/index_modal.py` in the
+  GitHub repo runs it on a Modal GPU; `index_local.py` if you have a CUDA box)
+  and pushed to the
+  [library dataset](https://huggingface.co/datasets/build-small-hackathon/repair-guy-library);
+  the Space syncs it to `/data/preindexed` at startup.
+- **Upload your own** — indexes on the Space's ZeroGPU, capped at
+  `MAX_UPLOAD_PAGES` (default 50) to protect quota.
+
 ## 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.
-- Optional env vars: `COLEMBED_MODEL_ID` (defaults to the 4B model),
-  `MINICPM_MODEL_ID` (defaults to `openbmb/MiniCPM-V-4_5`), `COLEMBED_ATTN`
-  (defaults to `sdpa`; set `flash_attention_2` if flash-attn is installed).
+- **Persistent storage** must be enabled (the library sync and uploads live
+  under `/data`). Budget roughly 5–12 MB per page of float16 token
+  embeddings; a 1000-page manual is ~6–12 GB — size the tier to the library.
+- Optional env vars: `LIBRARY_DATASET_ID`, `MAX_UPLOAD_PAGES`,
+  `COLEMBED_MODEL_ID` (defaults to the 4B model), `MINICPM_MODEL_ID`
+  (defaults to `openbmb/MiniCPM-V-4_5`), `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,96 +1,169 @@
-"""Gradio + ZeroGPU Space: upload a repair manual, ask questions.
+"""Gradio + ZeroGPU Space: ask questions over repair manuals.
 
 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.
+with Nemotron ColEmbed v2 (multi-vector, late interaction). A question is
+answered by embedding the query, running MaxSim over page embeddings streamed
+from disk, and handing the top pages to MiniCPM-V — all in one ZeroGPU call.
+
+Two tabs:
+  Library     — large manuals pre-indexed offline (scripts/index_local.py),
+                synced at startup from the library dataset on the Hub.
+  Upload      — index your own small PDF on ZeroGPU (page-capped to protect
+                the Space's GPU quota), then ask it questions.
 
 Module layout:
-  models/colembed.py   ColEmbed     — embedding model + GPU embed/search
-  models/minicpm.py    MiniCPM      — remote VLM that answers over page images
+  models/colembed.py   ColEmbed     — embedding model + GPU embed / MaxSim
+  models/minicpm.py    MiniCPM      — local 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 os
 
 import gradio as gr
+from huggingface_hub import snapshot_download
 
-from core.constants import DEFAULT_TOP_K
+from core.constants import (
+    DEFAULT_TOP_K,
+    LIBRARY_DATASET_ID,
+    MAX_UPLOAD_PAGES,
+    PREINDEXED_DIR,
+    UPLOADS_DIR,
+)
 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()
+# Construct once at startup (the models load onto cuda here, in the main process).
+library = Store(PREINDEXED_DIR)
+uploads = Store(UPLOADS_DIR)
 embedder = ColEmbed()
-ingest_pipeline = IngestPipeline(embedder, store)
-ask_pipeline = AskPipeline(embedder, store, MiniCPM())
+ingest_pipeline = IngestPipeline(embedder, uploads)
+ask_pipeline = AskPipeline()
 
 
-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."
+def sync_library() -> None:
+    """Pull pre-indexed manuals from the library dataset into /data.
+    A missing or empty dataset just means an empty library tab."""
     try:
-        doc = ingest_pipeline.run(
-            pdf_file, name, lambda frac, desc: progress(frac, desc=desc)
+        snapshot_download(
+            LIBRARY_DATASET_ID, repo_type="dataset", local_dir=library.root
         )
+    except Exception as e:
+        print(f"Library dataset not synced ({LIBRARY_DATASET_ID}): {e}")
+
+
+sync_library()
+
+
+def _choices(store: Store) -> list[tuple[str, str]]:
+    return [(d["name"], d["doc_id"]) for d in store.list_docs()]
+
+
+def ask_library(question, doc_id):
+    if not doc_id:
+        raise gr.Error("Pick a manual first.")
+    try:
+        return ask_pipeline.run(library, question, [doc_id], DEFAULT_TOP_K)
     except ValueError as e:
         raise gr.Error(str(e)) from e
-    return doc_id, f"Indexed **{doc['name']}** ({doc['pages']} pages) — ask away."
 
 
-def answer_question(question, doc_id):
+def ask_upload(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)
+        return ask_pipeline.run(uploads, question, [doc_id], DEFAULT_TOP_K)
     except ValueError as e:
         raise gr.Error(str(e)) from e
 
 
+def index_manual(pdf_file, progress=gr.Progress()):
+    """Runs on upload: embed the manual's pages (or reuse a previous index).
+    Generator — each yield streams (doc_state, status markdown) to the UI."""
+    if not pdf_file:
+        yield None, ""
+        return
+    name = os.path.splitext(os.path.basename(pdf_file))[0].replace("_", " ")
+    doc_id = slugify(name)
+    if uploads.exists(doc_id):
+        pages = len(uploads.meta(doc_id)["pages"])
+        yield doc_id, f"**{name}** is already indexed ({pages} pages) — ask away."
+        return
+    yield None, f"⏳ Indexing **{name}** — preparing…"
+    try:
+        for event in ingest_pipeline.run(pdf_file, name, max_pages=MAX_UPLOAD_PAGES):
+            if event[0] == "progress":
+                _, done, total = event
+                progress(done / total, desc=f"Embedding pages {done}/{total}")
+                yield None, f"⏳ Indexing **{name}** — {done}/{total} pages embedded…"
+            else:
+                doc = event[1]
+    except ValueError as e:
+        raise gr.Error(str(e)) from e
+    yield doc_id, f"✅ Indexed **{doc['name']}** ({doc['pages']} pages) — ask away."
+
+
 with gr.Blocks(title="Repair Guy") as demo:
     gr.Markdown(
         "# 🔧 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."
+        "Ask questions over repair manuals. Pages are retrieved visually with "
+        "[Nemotron ColEmbed v2](https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2) "
+        "(late interaction, no parsing) and answered by 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="Repair manual (PDF)", file_types=[".pdf"], type="filepath"
-            )
-            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?",
-            )
-            ask_btn = gr.Button("Ask", variant="primary")
-        with gr.Column(scale=2):
-            answer_out = gr.Markdown(label="Answer")
-            pages_out = gr.Gallery(label="Pages used", columns=3, height=420)
+    with gr.Tab("📚 Library"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                manual_in = gr.Dropdown(label="Manual", choices=[])
+                lib_question_in = gr.Textbox(
+                    label="Question",
+                    lines=2,
+                    placeholder="e.g. What is the tightening torque for the universal joint flange bolts?",
+                )
+                lib_ask_btn = gr.Button("Ask", variant="primary")
+            with gr.Column(scale=2):
+                lib_answer_out = gr.Markdown(label="Answer")
+                lib_pages_out = gr.Gallery(label="Pages used", columns=3, height=420)
+    with gr.Tab("📄 Upload your own"):
+        gr.Markdown(
+            f"Upload a PDF of up to **{MAX_UPLOAD_PAGES} pages** (indexing runs "
+            "on this Space's GPU quota — big manuals live in the Library tab)."
+        )
+        doc_state = gr.State(None)
+        with gr.Row():
+            with gr.Column(scale=1):
+                pdf_in = gr.File(
+                    label="Repair manual (PDF)", file_types=[".pdf"], type="filepath"
+                )
+                status_out = gr.Markdown()
+                up_question_in = gr.Textbox(label="Question", lines=2)
+                up_ask_btn = gr.Button("Ask", variant="primary")
+            with gr.Column(scale=2):
+                up_answer_out = gr.Markdown(label="Answer")
+                up_pages_out = gr.Gallery(label="Pages used", columns=3, height=420)
 
+    lib_ask_btn.click(
+        ask_library, inputs=[lib_question_in, manual_in],
+        outputs=[lib_answer_out, lib_pages_out],
+    )
+    lib_question_in.submit(
+        ask_library, inputs=[lib_question_in, manual_in],
+        outputs=[lib_answer_out, lib_pages_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]
+    up_ask_btn.click(
+        ask_upload, inputs=[up_question_in, doc_state],
+        outputs=[up_answer_out, up_pages_out],
     )
-    question_in.submit(
-        answer_question, inputs=[question_in, doc_state], outputs=[answer_out, pages_out]
+    up_question_in.submit(
+        ask_upload, inputs=[up_question_in, doc_state],
+        outputs=[up_answer_out, up_pages_out],
     )
+    demo.load(lambda: gr.update(choices=_choices(library)), outputs=[manual_in])
 
 
 if __name__ == "__main__":

core/constants.py

@@ -1,9 +1,14 @@
 import os
 
 # Embedding model (late-interaction / ColBERT-style page embeddings).
+# Revisions are pinned because both models load trust_remote_code; bump
+# deliberately after reviewing upstream changes.
 COLEMBED_MODEL_ID = os.environ.get(
     "COLEMBED_MODEL_ID", "nvidia/nemotron-colembed-vl-4b-v2"
 )
+COLEMBED_REVISION = os.environ.get(
+    "COLEMBED_REVISION", "0ed152d91f8ad4c5d48296b51c220f686641a398"
+)
 # 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")
@@ -13,27 +18,43 @@ 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
+# reported between chunks. Bigger chunks mean fewer ZeroGPU queue waits per
+# manual (~0.5s/page observed, so 64 pages ≈ 35-40s of a 240s budget — a
+# 1000-page manual is ~16 GPU calls).
+EMBED_PAGES_PER_CALL = 64
+EMBED_BATCH_SIZE = 8
+EMBED_GPU_DURATION = 240
 
 # 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
 
 # Answering model (runs locally on ZeroGPU).
 MINICPM_MODEL_ID = os.environ.get("MINICPM_MODEL_ID", "openbmb/MiniCPM-V-4_5")
-ANSWER_GPU_DURATION = 120
+MINICPM_REVISION = os.environ.get(
+    "MINICPM_REVISION", "fd3209b2e0580e346fc33d2c6f85b6e9332eecda"
+)
 ANSWER_MAX_NEW_TOKENS = 2048
 
-# 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"
+# One ZeroGPU call covers the whole question: query embedding + MaxSim +
+# page rendering + answer generation.
+ASK_GPU_DURATION = 120
+
+# Manual stores. HF Spaces persistent storage mounts at /data; fall back to a
+# local directory for development. Pre-indexed manuals (large PDFs embedded
+# offline, see scripts/index_local.py) are synced from the library dataset at
+# startup; in-app uploads are indexed on ZeroGPU and page-capped to protect
+# the Space's quota.
+_DATA_ROOT = os.environ.get("DATA_ROOT") or (
+    "/data"
     if os.path.isdir("/data")
-    else os.path.join(os.path.dirname(os.path.dirname(__file__)), "library")
+    else os.path.join(os.path.dirname(os.path.dirname(__file__)), "data")
+)
+PREINDEXED_DIR = os.path.join(_DATA_ROOT, "preindexed")
+UPLOADS_DIR = os.path.join(_DATA_ROOT, "uploads")
+LIBRARY_DATASET_ID = os.environ.get(
+    "LIBRARY_DATASET_ID", "build-small-hackathon/repair-guy-library"
 )
+MAX_UPLOAD_PAGES = int(os.environ.get("MAX_UPLOAD_PAGES", "50"))

core/store.py

@@ -1,6 +1,6 @@
 """On-disk store of per-page ColEmbed token embeddings.
 
-Layout: STORE_DIR/<doc_id>/
+Layout: <store root>/<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)
@@ -20,8 +20,6 @@ import shutil
 
 import numpy as np
 
-from core.constants import STORE_DIR
-
 EMB_DTYPE = np.float16
 
 
@@ -69,7 +67,7 @@ class DocWriter:
 
 
 class Store:
-    def __init__(self, root: str = STORE_DIR):
+    def __init__(self, root: str):
         self.root = os.path.abspath(root)
         os.makedirs(self.root, exist_ok=True)
 

models/colembed.py

@@ -1,14 +1,13 @@
 """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
-
 The model is a module-level global: ZeroGPU packs module-level CUDA tensors at
 startup and shares them with the GPU worker, whereas function arguments are
 pickled — and the trust_remote_code model class is not picklable.
 
+_embed_pages_on_gpu is a ZeroGPU entry point (used at index time);
+maxsim_search is a plain function so the ask pipeline can run it inside its
+own single GPU call together with answer generation.
+
 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.
@@ -22,20 +21,21 @@ import numpy as np
 import spaces
 import torch
 from PIL import Image
-from transformers import AutoModel
+from transformers import AutoModel, AutoProcessor
 
 from core.constants import (
     COLEMBED_ATTN,
     COLEMBED_MODEL_ID,
+    COLEMBED_REVISION,
     EMBED_BATCH_SIZE,
     EMBED_GPU_DURATION,
     SCORE_PAGES_PER_BATCH,
-    SEARCH_GPU_DURATION,
 )
 
 _MODEL = (
     AutoModel.from_pretrained(
         COLEMBED_MODEL_ID,
+        revision=COLEMBED_REVISION,
         trust_remote_code=True,
         dtype=torch.bfloat16,
         attn_implementation=COLEMBED_ATTN,
@@ -43,6 +43,11 @@ _MODEL = (
     .to("cuda")
     .eval()
 )
+# Pre-build the processor the remote code would otherwise lazily create per
+# GPU worker (it caches on this exact attribute, see _get_processor).
+_MODEL._processor = AutoProcessor.from_pretrained(
+    COLEMBED_MODEL_ID, revision=COLEMBED_REVISION, trust_remote_code=True
+)
 
 # The remote code's forward_documents hardcodes DataLoader(num_workers=8), but
 # the ZeroGPU worker is a daemonic process and may not spawn children
@@ -72,8 +77,11 @@ def _embed_pages_on_gpu(images: list[Image.Image]) -> list[np.ndarray]:
     return out
 
 
-@spaces.GPU(duration=SEARCH_GPU_DURATION)
-def _search_on_gpu(question: str, store, doc_ids, top_k: int):
+def maxsim_search(
+    question: str, store, doc_ids: list[str] | None, top_k: int
+) -> list[tuple[str, int, float]]:
+    """Top-K (doc_id, page_num, score) across docs. Must run on GPU (called
+    from within a @spaces.GPU context)."""
     results = []
     with torch.no_grad():
         q = _MODEL.forward_queries([question], batch_size=1)[0].to(torch.float16)
@@ -95,9 +103,3 @@ class ColEmbed:
     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(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(question, store, doc_ids, top_k)

models/minicpm.py

@@ -1,19 +1,20 @@
-"""MiniCPM-V on ZeroGPU: answers a question grounded in the retrieved
-repair-manual pages.
+"""MiniCPM-V: answers a question grounded in the retrieved repair-manual pages.
 
 The model and tokenizer are module-level globals: ZeroGPU packs module-level
 CUDA tensors at startup and shares them with the GPU worker, whereas function
 arguments are pickled — and trust_remote_code model classes are not picklable.
+
+generate_answer is a plain function: the ask pipeline calls it inside its own
+single @spaces.GPU call, right after retrieval.
 """
 
 from __future__ import annotations
 
-import spaces
 import torch
 from PIL import Image
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoProcessor, AutoTokenizer
 
-from core.constants import ANSWER_GPU_DURATION, ANSWER_MAX_NEW_TOKENS, MINICPM_MODEL_ID
+from core.constants import ANSWER_MAX_NEW_TOKENS, MINICPM_MODEL_ID, MINICPM_REVISION
 
 PROMPT = (
     "You are a repair-manual assistant. The images are the manual pages most "
@@ -38,6 +39,7 @@ PROMPT = (
 _MODEL = (
     AutoModel.from_pretrained(
         MINICPM_MODEL_ID,
+        revision=MINICPM_REVISION,
         trust_remote_code=True,
         dtype=torch.bfloat16,
         attn_implementation="sdpa",
@@ -45,11 +47,19 @@ _MODEL = (
     .to("cuda")
     .eval()
 )
-_TOKENIZER = AutoTokenizer.from_pretrained(MINICPM_MODEL_ID, trust_remote_code=True)
+_TOKENIZER = AutoTokenizer.from_pretrained(
+    MINICPM_MODEL_ID, revision=MINICPM_REVISION, trust_remote_code=True
+)
+# Pre-build the processor chat() would otherwise lazily create per GPU worker
+# (it caches on this exact attribute, see modeling_minicpmv.chat).
+_MODEL.processor = AutoProcessor.from_pretrained(
+    MINICPM_MODEL_ID, revision=MINICPM_REVISION, trust_remote_code=True
+)
 
 
-@spaces.GPU(duration=ANSWER_GPU_DURATION)
-def _answer_on_gpu(question: str, pages: list[tuple[str, Image.Image]]) -> str:
+def generate_answer(question: str, pages: list[tuple[str, Image.Image]]) -> str:
+    """pages: [(label, page image)] in retrieval order. Must run on GPU
+    (called from within a @spaces.GPU context)."""
     content = []
     for label, img in pages:  # chat() accepts interleaved strings and PIL images
         content.append(f"[{label}]")
@@ -63,9 +73,3 @@ def _answer_on_gpu(question: str, pages: list[tuple[str, Image.Image]]) -> str:
             max_new_tokens=ANSWER_MAX_NEW_TOKENS,
         )
     return str(answer).strip()
-
-
-class MiniCPM:
-    def answer(self, question: str, pages: list[tuple[str, Image.Image]]) -> str:
-        """pages: [(label, page image)] in retrieval order."""
-        return _answer_on_gpu(question, pages)

pipelines/ask.py

@@ -1,35 +1,51 @@
 """Ask pipeline: question -> MaxSim retrieval over the store -> top-K page
-images -> MiniCPM answer grounded in those pages."""
+images -> MiniCPM answer grounded in those pages.
+
+The whole question runs in ONE @spaces.GPU call (query embedding + MaxSim +
+page rendering + answer generation), so each question pays the ZeroGPU
+allocation wait once. Indexing has its own separate GPU entry point.
+"""
 
 from __future__ import annotations
 
+import spaces
+
+from core.constants import ASK_GPU_DURATION
 from core.pdf import render_page
 from core.store import Store
-from models.colembed import ColEmbed
-from models.minicpm import MiniCPM
+from models.colembed import maxsim_search
+from models.minicpm import generate_answer
+
+
+@spaces.GPU(duration=ASK_GPU_DURATION)
+def _ask_on_gpu(
+    question: str,
+    store: Store,
+    doc_ids: list[str] | None,
+    top_k: int,
+    names: dict[str, str],
+):
+    hits = maxsim_search(question, store, doc_ids, top_k)
+    pages = [
+        (f"{names[doc_id]} — p.{page}", render_page(store.pdf_path(doc_id), page), score)
+        for doc_id, page, score in hits
+    ]
+    answer = generate_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
 
 
 class AskPipeline:
-    def __init__(self, embedder: ColEmbed, store: Store, llm: MiniCPM):
-        self.embedder = embedder
-        self.store = store
-        self.llm = llm
+    """Stateless: the store is passed per call, so the same pipeline serves
+    both the pre-indexed library and user uploads."""
 
-    def run(self, question: str, doc_ids: list[str] | None, top_k: int):
+    def run(self, store: Store, 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()
+        docs = 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))
+            raise ValueError("No manuals indexed yet — upload one first.")
         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
+        return _ask_on_gpu(question, store, doc_ids or None, int(top_k), names)

pipelines/ingest.py

@@ -1,12 +1,13 @@
 """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.
+short; progress events are yielded between chunks so the UI can stream them.
 """
 
 from __future__ import annotations
 
 import os
+from collections.abc import Iterator
 
 from core.constants import EMBED_PAGES_PER_CALL, RENDER_DPI
 from core.pdf import page_count, render_pages
@@ -19,9 +20,12 @@ class IngestPipeline:
         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."""
+    def run(
+        self, pdf_path: str | None, doc_name: str = "", max_pages: int | None = None
+    ) -> Iterator[tuple]:
+        """Index one PDF. Generator yielding ("progress", pages_done, total)
+        after each embedded chunk, then ("done", doc summary dict) last.
+        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"):
@@ -32,6 +36,12 @@ class IngestPipeline:
         )
         doc_id = slugify(name)
         total = page_count(pdf_path)
+        if max_pages and total > max_pages:
+            raise ValueError(
+                f"This PDF has {total} pages — uploads are capped at {max_pages} "
+                "pages to conserve the Space's GPU quota. Large manuals belong "
+                "in the pre-indexed library."
+            )
 
         writer = self.store.create(doc_id, name, pdf_path, RENDER_DPI, self.embedder.MODEL_ID)
         try:
@@ -40,10 +50,9 @@ class IngestPipeline:
                 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")
+                yield ("progress", nums[-1], total)
             writer.finalize()
         except BaseException:
             writer.abort()
             raise
-        return {"doc_id": doc_id, "name": name, "pages": total}
+        yield ("done", {"doc_id": doc_id, "name": name, "pages": total})