| """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). 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 / 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 |
| """ |
|
|
| import os |
|
|
| import gradio as gr |
| from huggingface_hub import snapshot_download |
|
|
| 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 pipelines.ask import AskPipeline |
| from pipelines.ingest import IngestPipeline |
|
|
| |
| library = Store(PREINDEXED_DIR) |
| uploads = Store(UPLOADS_DIR) |
| embedder = ColEmbed() |
| ingest_pipeline = IngestPipeline(embedder, uploads) |
| ask_pipeline = AskPipeline() |
|
|
|
|
| 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: |
| 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 |
|
|
|
|
| def ask_upload(question, doc_id): |
| if not doc_id: |
| raise gr.Error("Upload a repair manual first.") |
| try: |
| 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" |
| "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." |
| ) |
| 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]) |
| up_ask_btn.click( |
| ask_upload, inputs=[up_question_in, doc_state], |
| outputs=[up_answer_out, up_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__": |
| demo.launch() |
|
|
