diff --git a/docs/docs/extraction/releasenotes.md b/docs/docs/extraction/releasenotes.md index a38d0384c4..8f5ef23848 100644 --- a/docs/docs/extraction/releasenotes.md +++ b/docs/docs/extraction/releasenotes.md @@ -19,14 +19,14 @@ Highlights for the 26.05 release include: ### Pipeline and ingestion -- Legacy `nv-ingest` code paths removed; `graph_pipeline` and the graph stage registry are the canonical ingestion path +- Legacy `nv-ingest` and compatibility pipeline CLI code paths removed; `retriever ingest` and the graph stage registry are the canonical ingestion paths - Manifest-based ingest routing replaces input-type routing; `retriever ingest` is input-aware for PDF, image, audio, video, text, HTML, DOCX/PPTX, SVG, and related types - `allow_no_gpu` option to skip GPU requirement during ingest for CPU-only experimentation ### CLI - Root CLI adds first-class `retriever ingest` and `retriever query` commands with NIM URL flags, batch tuning, and LanceDB overwrite/append controls -- For product ingest and retrieval, prefer `retriever ingest` and `retriever query`; `retriever pipeline run` remains available for compatibility and development workflows. Other top-level subcommands—including `pdf`, `html`, `eval`, `benchmark`, `harness`, `online`, `compare`, `image`, and `skill-eval`—are development and experimental +- `retriever ingest` and `retriever query` replace the retired compatibility pipeline command. Other top-level subcommands—including `eval`, `benchmark`, `harness`, and `skill-eval`—are development and experimental ### Retriever Service and deployment diff --git a/evaluation/bo767_recall.ipynb b/evaluation/bo767_recall.ipynb deleted file mode 100644 index fc89808d35..0000000000 --- a/evaluation/bo767_recall.ipynb +++ /dev/null @@ -1,129 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "3c4c7d5f-51fb-4879-8fd3-d304165ffd38", - "metadata": {}, - "source": [ - "# Evaluate bo767 retrieval recall accuracy with NeMo Retriever" - ] - }, - { - "cell_type": "markdown", - "id": "f6f80e87", - "metadata": {}, - "source": [ - "To download the bo767 PDF corpus, please refer to [digital_corpora_download.ipynb](https://github.com/NVIDIA/NeMo-Retriever/blob/main/evaluation/digital_corpora_download.ipynb)" - ] - }, - { - "cell_type": "markdown", - "id": "0d3116aa-2992-4798-bae6-42a4e3cac58f", - "metadata": {}, - "source": [ - "### CLI" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "00aeacea-2eb7-45a8-8e62-edf52fdc9d9e", - "metadata": {}, - "outputs": [], - "source": [ - "%%bash\n", - "retriever pipeline run path/to/bo767/pdfs \\\n", - " --vdb-kwargs-json '{\"uri\":\"../lancedb\",\"table_name\":\"bo767\"}' \\\n", - " --evaluation-mode beir \\\n", - " --beir-dataset-name bo767 \\\n", - " --quiet" - ] - }, - { - "cell_type": "markdown", - "id": "e25582b6-005b-47d2-8b47-b0823422bda9", - "metadata": {}, - "source": [ - "### Python" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9f1dba08-c468-425f-9eb3-48fe568b67c7", - "metadata": {}, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "\n", - "from nemo_retriever import create_ingestor\n", - "from nemo_retriever.model import VL_EMBED_MODEL\n", - "from nemo_retriever.params import EmbedParams, VdbUploadParams\n", - "from nemo_retriever.recall.beir import BeirConfig, evaluate_lancedb_beir, resolve_beir_dataset_options\n", - "\n", - "input_path = str(Path(\"path/to/bo767/pdfs\").resolve())\n", - "lancedb_uri = str(Path(\"../lancedb\").resolve())\n", - "table_name = \"bo767\"\n", - "\n", - "result = (\n", - " create_ingestor(run_mode=\"batch\")\n", - " .files(input_path)\n", - " .extract()\n", - " .embed(EmbedParams(model_name=VL_EMBED_MODEL))\n", - " .vdb_upload(\n", - " VdbUploadParams(\n", - " vdb_op=\"lancedb\",\n", - " vdb_kwargs={\"uri\": lancedb_uri, \"table_name\": table_name, \"overwrite\": True},\n", - " )\n", - " )\n", - " .ingest()\n", - ")\n", - "\n", - "beir = resolve_beir_dataset_options(dataset_name=\"bo767\")\n", - "\n", - "dataset, raw_hits, run, metrics = evaluate_lancedb_beir(\n", - " BeirConfig(\n", - " lancedb_uri=lancedb_uri,\n", - " lancedb_table=table_name,\n", - " embedding_model=VL_EMBED_MODEL,\n", - " loader=beir.loader,\n", - " dataset_name=beir.dataset_name,\n", - " doc_id_field=beir.doc_id_field,\n", - " ks=beir.ks,\n", - " )\n", - ")\n", - "\n", - "metrics" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "249a4852", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.9" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/evaluation/bo767_recall.md b/evaluation/bo767_recall.md new file mode 100644 index 0000000000..2545f6c4b4 --- /dev/null +++ b/evaluation/bo767_recall.md @@ -0,0 +1,70 @@ +# Evaluate BO767 Retrieval with the Retriever Harness + +Use the Retriever harness for the canonical BO767 end-to-end ingest, query, +and BEIR evaluation. The checked-in runfile selects batch execution and carries +the benchmark's worker tuning and required file, page, and query counts. + +Run these commands from the repository root. + +## Configure the Dataset Paths + +Copy the example path map to an untracked location: + +```bash +cp nemo_retriever/harness/dataset_paths.example.yaml /tmp/retriever-dataset-paths.yaml +``` + +Edit the `bo767` entry so `path` points to the directory containing the 767 +documents and `query_file` points to the BO767 query/qrels CSV available on the +machine running the benchmark. + +## Validate the Resolved Run + +Resolve the run without launching ingestion or evaluation: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name bo767_beir_check \ + --output-dir /tmp/retriever-harness-bo767-check \ + --dataset-paths /tmp/retriever-dataset-paths.yaml \ + --dry-run \ + --json \ + nemo_retriever/harness/runfiles/bo767_beir.json +``` + +Inspect `session_summary.json`, `expanded_runs.json`, and the child run's +`resolved_benchmark.json` before starting the full run. + +## Run the Evaluation + +Remove `--dry-run` and choose a durable artifact directory: + +```bash +uv run --project nemo_retriever retriever harness run-files \ + --session-name bo767_beir \ + --output-dir /local/path/to/retriever-artifacts/bo767-beir \ + --dataset-paths /tmp/retriever-dataset-paths.yaml \ + --json \ + nemo_retriever/harness/runfiles/bo767_beir.json +``` + +BO767 is a large batch benchmark. Keep the terminal process alive until the +harness reaches a terminal state; model startup and ingestion can be quiet for +extended periods. + +## Read the Results + +The harness artifacts are the evaluation contract: + +- `status.json` reports the current phase and concise failure state. +- `results.json` is the authoritative terminal result and summary metrics. +- `session_summary.json` is the terminal result for the runfile session. +- `beir_metrics.json` contains the complete BEIR metric family. +- `query_results.jsonl` contains per-query latency and ranked hits. +- `environment.json` records the commit and runtime context. + +Refer to +[`nemo_retriever/harness/EXPECTED_RESULTS.md`](../nemo_retriever/harness/EXPECTED_RESULTS.md#bo767) +for the expected BO767 counts and current reference metrics. For the complete +harness contract and troubleshooting guidance, refer to +[`nemo_retriever/harness/README.md`](../nemo_retriever/harness/README.md). diff --git a/examples/README.md b/examples/README.md index 883b482fd8..79117823ea 100644 --- a/examples/README.md +++ b/examples/README.md @@ -14,10 +14,10 @@ Start with these guides and notebooks: - [Workflow: Ingest documents](https://docs.nvidia.com/nemo/retriever/latest/extraction/workflow-document-ingestion/) - [Adding Custom Metadata for Filtered Search/Retrieval](nemo_retriever_retriever_query_metadata_filter.ipynb) — also summarized on [Vector databases — Metadata and filtering](https://docs.nvidia.com/nemo/retriever/latest/extraction/vdbs/#metadata-and-filtering) -For advanced scenarios, use these notebooks: +For advanced scenarios, use these guides and notebooks: - [Build a Custom Vector Database Operator](building_vdb_operator.ipynb) - [Try Enterprise RAG Blueprint](https://build.nvidia.com/nvidia/multimodal-pdf-data-extraction-for-enterprise-rag) -- [Evaluate bo767 retrieval recall accuracy with NeMo Retriever Library](https://github.com/NVIDIA/NeMo-Retriever/blob/main/evaluation/bo767_recall.ipynb) +- [Evaluate BO767 retrieval with the Retriever harness](https://github.com/NVIDIA/NeMo-Retriever/blob/main/evaluation/bo767_recall.md) - [Multimodal RAG with LangChain](langchain_multimodal_rag.ipynb) - [Multimodal RAG with LlamaIndex](llama_index_multimodal_rag.ipynb) diff --git a/nemo_retriever/README.md b/nemo_retriever/README.md index 857581ecb4..e4b783bdaa 100644 --- a/nemo_retriever/README.md +++ b/nemo_retriever/README.md @@ -97,7 +97,6 @@ The [test PDF](../data/multimodal_test.pdf) contains text, tables, charts, and i > **Note:** `retriever ingest` defaults to local, in-process execution. Use `retriever ingest batch ...` for Ray Data scale-out on larger workloads. > File formats and internal extraction stages are not separate root commands; configure supported behavior through `retriever ingest`. -> `retriever pipeline run` remains callable for compatibility while existing callers migrate, but it is hidden from root help. The examples below use default local GPU inference (no `invoke_url` specified) and require the `[local]` extra and the CUDA 13 torch override from the setup steps above. For remote NIM inference without a local GPU, refer to [Run with remote inference](#run-with-remote-inference-no-local-gpu-required). @@ -147,22 +146,22 @@ chunks = ingestor.ingest() # pandas.DataFrame (batch and inprocess) ### Ingest a test corpus (CLI) -`graph_pipeline` is the canonical ingestion script for building a multi-document -LanceDB corpus. Point it at a **directory** of PDFs to produce a ready-to-query table. +Point `retriever ingest` at a **directory** of PDFs to produce a ready-to-query +LanceDB table. > **Corpus size matters.** LanceDB's default IVF index needs at least 16 > chunks to train its 16 k-means partitions. Single-PDF ingestion will fail -> at the indexing step; point `graph_pipeline` at a directory with enough +> at the indexing step; point `retriever ingest` at a directory with enough > documents to clear that threshold. Replace `/your-example-dir` below with > the path to your own corpus. ```bash -python -m nemo_retriever.examples.graph_pipeline \ - /your-example-dir \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' +retriever ingest /your-example-dir \ + --lancedb-uri lancedb \ + --table-name nemo-retriever ``` -Chunks land at `./lancedb/nemo-retriever`, which matches the `vdb_kwargs` +Chunks land at `./lancedb/nemo-retriever`, which matches the storage settings used in [Run a recall query](#run-a-recall-query) below. With the `[local]` extra installed (see setup), defaults point at local-GPU extraction and embedding. Use enough documents in the directory to clear the LanceDB IVF @@ -174,9 +173,9 @@ through [build.nvidia.com](https://build.nvidia.com/) NIMs instead: ```bash export NVIDIA_API_KEY=nvapi-... -python -m nemo_retriever.examples.graph_pipeline \ - /your-example-dir \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' \ +retriever ingest /your-example-dir \ + --lancedb-uri lancedb \ + --table-name nemo-retriever \ --page-elements-invoke-url https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-page-elements-v3 \ --ocr-invoke-url https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v2 \ --table-structure-invoke-url https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-table-structure-v1 \ @@ -228,7 +227,7 @@ Since the ingestion job automatically populated a lancedb table with all these c from nemo_retriever.graph.retriever import Retriever retriever = Retriever( - # values used by the graph_pipeline example above + # values used by the retriever ingest example above vdb_kwargs={"uri": "lancedb", "table_name": "nemo-retriever"}, top_k=5, rerank=False @@ -454,8 +453,8 @@ print(len(result.chunks), "chunks from", {m.get("source") for m in result.metada print(f"{result.latency_s:.2f}s on {result.model}") ``` -Local-GPU shortcut: if you ingested with default `graph_pipeline` flags -(`--embed` omitted, `[local]` extra installed), drop `embed_kwargs` to reuse +Local-GPU shortcut: if you ingested with default `retriever ingest` flags +(`[local]` extra installed), drop `embed_kwargs` to reuse the bundled `VL_EMBED_MODEL`. Live RAG with scoring and an LLM judge (requires a ground-truth `reference`): diff --git a/nemo_retriever/developer_docs/README.md b/nemo_retriever/developer_docs/README.md index 64ed705d72..3f06086398 100644 --- a/nemo_retriever/developer_docs/README.md +++ b/nemo_retriever/developer_docs/README.md @@ -12,4 +12,3 @@ architecture, subsystems, and developer-facing tools. | [Retriever Harness README](../harness/README.md) | Operator and agent instructions for the artifact-first Retriever benchmark harness. | | [Retriever Harness PRD](harness_retriever_ingest_query_prd.md) | Product requirements for the artifact-first Retriever ingest/query benchmark harness revamp. | | [Root Ingest CLI Design](root_ingest_cli_design.md) | Reviewer guide for the `retriever ingest` local, batch, and service CLI ownership split. | -| [VDB Retrieval Refactor Scope](vdb_retrieval_refactor.md) | Motivation, in-scope paths, and ownership boundaries for graph-pipeline VDB-agnostic retrieval. | diff --git a/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md b/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md index d8b7fd5c58..6ced880c91 100644 --- a/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md +++ b/nemo_retriever/developer_docs/harness_retriever_ingest_query_prd.md @@ -23,10 +23,10 @@ engineers. The harness should run the new library direction end to end: 4. Emit stable `summary_metrics` and machine-readable artifacts for humans, agents, and downstream reporting. -This is a total revamp, not a compatibility wrapper around the current harness. -Assume `retriever pipeline run` is deleted in the next release. The current -`sweep`, `compare`, and legacy graph-pipeline command builder can be rewritten -or removed if they get in the way. +This is a total revamp, not a compatibility wrapper around the old harness. The +retired pipeline CLI is not part of the design. The old `sweep`, `compare`, and +graph-pipeline command builder can be rewritten or removed if they get in the +way. The most important design choice: use a small typed benchmark registry in code, not a sprawling YAML configuration system. YAML/runfiles can exist as an escape @@ -745,8 +745,7 @@ validation through the CLI and artifact contract. - The harness has no user-facing `--engine` flag. - The harness does not duplicate Typer options from `retriever ingest` or `retriever query`. -- No phase-one code path invokes `retriever pipeline run` or - `nemo_retriever.examples.graph_pipeline`. +- No phase-one code path invokes a retired CLI adapter. - CLI text formatting is explicitly non-contractual; machine consumers use artifact files or `--json` read-only commands. - Validation combines focused contract tests with functional, artifact-driven diff --git a/nemo_retriever/developer_docs/root_ingest_cli_design.md b/nemo_retriever/developer_docs/root_ingest_cli_design.md index 00ec3bbc89..db00913b88 100644 --- a/nemo_retriever/developer_docs/root_ingest_cli_design.md +++ b/nemo_retriever/developer_docs/root_ingest_cli_design.md @@ -215,20 +215,9 @@ image-store URI, dry-run, and quiet output. Service-backed query support belongs in the query CLI/service boundary, not in the ingest CLI. -## Pipeline Compatibility - -`retriever pipeline run` is not the future public ingest interface. It remains -callable, but hidden from root help, for compatibility and development behavior -such as: - -- intermediate Parquet artifacts -- pipeline reports and runtime metrics -- eval, recall, harness, BEIR/QA workflows -- legacy callers not yet migrated to root ingest/query - -For graph ingest paths, pipeline compatibility should continue to reuse the -canonical ingest plan/execution layer instead of shelling out to root CLI -commands. +The supported command-line seam is `retriever ingest` for ingestion and +`retriever query` for retrieval. Graph ingest paths reuse the canonical ingest +plan and execution modules rather than duplicating CLI option handling. ## Adding Or Changing A Flag diff --git a/nemo_retriever/developer_docs/vdb_retrieval_refactor.md b/nemo_retriever/developer_docs/vdb_retrieval_refactor.md deleted file mode 100644 index f7ea5adff7..0000000000 --- a/nemo_retriever/developer_docs/vdb_retrieval_refactor.md +++ /dev/null @@ -1,319 +0,0 @@ -# VDB Retrieval Refactor Scope and Current Contract - -**Module area:** `nemo_retriever.common.vdb`, `nemo_retriever.graph.retriever`, graph pipeline recall - -## Purpose - -This PR is about making graph-pipeline upload and retrieval work across VDB -backends without turning Retriever into a new retrieval framework. - -The intended mental model is: - -> Old Retriever, but the LanceDB-only search/write assumptions become VDB -> operator boundaries that can support LanceDB and other client `VDB` backends. - -The immediate target is graph pipeline recall over embedded document chunks -stored in LanceDB (or another configured client `VDB`). The implementation should stay small, reviewable, -and centered on the graph pipeline path. - -## Why This Exists - -Historically, much of the system was LanceDB-first: - -- pipeline CLI options and runtime summaries talked directly about LanceDB; -- graph ingestion wrote rows through LanceDB-specific utilities; -- recall assumed a LanceDB table; -- the older client VDB retrieval path owned query embedding through an endpoint - or client-side retrieval helper. - -That was reasonable before NeMo Retriever had a clear query-embedding owner. -The newer Retriever path is different: the graph pipeline creates document -embeddings with its configured embedding path, and retrieval should create -matching query embeddings before calling VDB search. - -The goal is not to discard `nv_ingest_client.util.vdb.VDB`. That client VDB API -is the backend boundary we want to satisfy. For this PR, that means: - -- `VDB.run(records)` uploads already-embedded graph records. -- `VDB.retrieval(vectors, **kwargs)` searches using already-embedded query - vectors. - -## Current Contract - -For graph recall, the desired contract is: - -1. The graph pipeline embeds document chunks with local HF by default, or with a - configured embedding endpoint. -2. The pipeline CLI materializes the graph result once. -3. `IngestVdbOperator` uploads those already-embedded records through - `nv_ingest_client.util.vdb.VDB.run(...)`. -4. `Retriever` embeds query strings with the matching embedding model/path. -5. `RetrieveVdbOperator` passes those precomputed query vectors to - `nv_ingest_client.util.vdb.VDB.retrieval(...)`. -6. Recall scoring consumes normalized Retriever hits. - -For this PR, VDB-agnostic means backend-agnostic storage and vector search. It -does not mean embedding-owner-agnostic. Retriever remains the owner of query -embedding; local HF is the default, and a configured embedding endpoint is used -only when one is explicitly provided. - -## In-Scope Flow: Graph Ingestion and VDB Upload - -The graph upload path is currently CLI orchestration, not a graph node: - -```text -pipeline run - -> _build_ingestor(...) - -> GraphIngestor.ingest() - -> build_graph(...) - -> RayDataExecutor/InprocessExecutor - -> raw graph result - -> _collect_results(run_mode, raw_result) - -> _upload_vdb_records(records, vdb_op, vdb_kwargs) - -> IngestVdbOperator(records) - -> IngestVdbOperator.process(...) - -> graph rows -> ingestion Python client VDB record shape - -> nv_ingest_client.util.vdb..run(records) -``` - -Responsibilities: - -- `pipeline/__main__.py` resolves `vdb_op` and opaque `vdb_kwargs`, owns result - materialization, and invokes VDB upload. -- `GraphIngestor` only builds and executes the graph, then returns the graph - result. -- `_collect_results(...)` is the single place that calls `take_all()` in batch - mode. -- `_upload_vdb_records(...)` is the CLI handoff from materialized graph records - to VDB upload. -- `IngestVdbOperator` is the boundary to ingestion Python client VDB writers. -- Record conversion exists only because `VDB.run(...)` expects the - ingestion Python client nested record shape, while graph output rows are ordinary - extraction/embed rows. - -This path should not contain LanceDB-specific logic except for the default -`vdb_op="lancedb"` and whatever opaque kwargs are passed to the LanceDB client -VDB implementation. - -### Why Upload Is Not a Graph Step Yet - -`vdb_upload` is intentionally not recorded in `GraphIngestor._stage_order` -today. The previous post-graph hook made upload look like a graph stage while it -actually ran after graph execution and cached materialized Ray records for the -CLI. The current shape is more explicit: - -```text -GraphIngestor - -> graph construction/execution only - -pipeline CLI - -> materialize once - -> upload once - -> save/evaluate/summarize -``` - -When VDB upload is made part of the graph, it should be added deliberately as a -real graph sink/stage, likely by wiring `IngestVdbOperator` through graph -construction rather than reintroducing a hidden post-execution hook. - -## In-Scope Flow: Recall Against a Populated VDB - -The graph recall path is: - -```text -pipeline run --query-csv ... - -> _run_evaluation(...) - -> RecallConfig(vdb_op=..., vdb_kwargs=...) - -> retrieve_and_score(...) - -> Retriever(vdb=..., vdb_kwargs=...) - -> Retriever.queries(...) - -> local HF or endpoint query embedding - -> RetrieveVdbOperator.process(query_vectors, ...) - -> VDB.retrieval(query_vectors, ...) - -> normalized Retriever hits - -> recall@k scoring -``` - -Responsibilities: - -- `RecallConfig` carries the selected VDB backend and kwargs into recall. -- `retrieve_and_score()` constructs `Retriever`. -- `Retriever.queries()` owns query text normalization, top-k handling, query - embedding, optional reranking, and returning Retriever-shaped hits. -- Backend-specific VDB search should live behind the client VDB retrieval - boundary, not inline in `Retriever`. - -The important invariant is that document vectors and query vectors are produced -by the same embedding model/path. In `pipeline run`, recall receives the same -embedding model and endpoint/API-key settings used by document embedding. When -no endpoint is configured, both sides default to local HF. - -## Retrieval Contract Change - -Older client VDB retrieval implementations treated `VDB.retrieval(...)` as a -query-string API: - -```text -query strings -> endpoint/service embedding -> vector search -``` - -That contract does not match graph recall, where Retriever must use the same -query/document embedding configuration before VDB search. - -The approved direction for this PR is to redefine retrieval VDBs toward: - -```text -query vectors -> vector search -``` - -That lets `Retriever.queries()` stay close to the upstream flow: - -```text -resolve query strings -resolve top_k -embed query strings locally or through the configured endpoint -call VDB retrieval with query vectors -rerank if configured -``` - -## Records and Hit Normalization - -There are two separate conversion concerns that should not be confused: - -1. **Upload conversion** - - graph rows -> ingestion Python client nested records - - required while `VDB.run(records)` is the upload contract -2. **Retrieval normalization** - - backend/client hits -> Retriever hit dictionaries - - required while different client VDB backends can return different hit shapes - -The upload conversion can be aggressively trimmed if we only support canonical -graph output rows. It does not need to support every historical input shape. - -The retrieval normalization can also be narrowed, but any legacy client -`Hit.to_dict()`-style handling must survive where a backend returns mapping-like -hits instead of plain dicts, so recall does not silently drop to zero. - -## Out of Scope - -These are intentionally out of scope for this PR unless explicitly reopened: - -- changing ingestion Python client VDB behavior beyond the retrieval vector-in - contract; -- moving query embedding into a broad new adapter framework; -- making all legacy BEIR, harness, or outdated example pipelines fully - VDB-agnostic beyond the compatibility fixes needed to keep current callers - constructible; -- renaming/generalizing the legacy `evaluate_lancedb_beir(...)` helper and - `BeirConfig` LanceDB fields. That API should become VDB-agnostic in a - follow-up, but this PR only keeps the existing BEIR compatibility path wired - through `Retriever`; -- preserving deprecated LanceDB public shims in the graph pipeline CLI; -- adding a broad new retrieval framework. - -## Review Objectives - -The implementation should converge on: - -- no first-class LanceDB concepts in the generic graph pipeline path; -- `vdb_op` plus opaque `vdb_kwargs` for backend selection/configuration; -- generic runtime language such as "VDB upload" instead of "LanceDB write"; -- minimal record conversion needed by the graph pipeline upload path; -- one graph-result materialization point in the CLI before upload/evaluation; -- query embedding remaining in Retriever, with local HF as the default and - endpoint embedding used only when configured; -- backend-specific vector search hidden behind VDB code, not spread through - Retriever; -- ingestion Python client VDB upload behavior unchanged. - -## Validation - -The PR is considered behaviorally sound when graph pipeline recall runs end-to-end -against LanceDB (upload, retrieve, score): - -```text -ingest jp20 - -> upload 3147 embedded records - -> retrieve 115 query-csv queries - -> recall@1/5/10 within expected tolerances for the dataset -``` - -Two embedding modes were validated: - -- **Local HF default:** no embedding endpoint is configured. The graph embeds - document chunks locally and `Retriever` embeds queries locally with the same - model family. -- **Hosted embedding endpoint:** `pipeline run` is passed - `--embed-invoke-url http://localhost:8012/v1` and - `--embed-model-name nvidia/llama-nemotron-embed-1b-v2`. The same endpoint and - model are propagated into recall, so document and query embeddings are - produced by the hosted NIM. - -The hosted service was checked with: - -```text -curl http://localhost:8012/v1/models -``` - -It advertised: - -```text -nvidia/llama-nemotron-embed-1b-v2 -``` - -Direct `/v1/embeddings` probes with both `input_type=query` and -`input_type=passage` succeeded. A failed full run with the pipeline's default -`nvidia/llama-nemotron-embed-vl-1b-v2` model produced `/v1/embeddings` 404s; -that was a hosted-model name mismatch, not an embedding concurrency issue. - -### Local HF JP20 - -Artifacts: - -- `nemo_retriever/artifacts/jp20_vdb_e2e_lancedb_vector_in_retrieval/` - -Results: - -```text -LanceDB: - rows collected: 3154 - rows uploaded: 3147 - queries: 115 - recall@1/5/10: 0.6261 / 0.9043 / 0.9391 - total_secs: 84.26 -``` - -### Hosted Embedding Endpoint JP20 - -Artifacts: - -- `nemo_retriever/artifacts/jp20_vdb_e2e_lancedb_remote_1b_embed_retrieval_20260428_1922/` -- `nemo_retriever/artifacts/jp20_vdb_e2e_lancedb_remote_1b_embed_default_embed_flags_20260428_1936/` - -Results: - -```text -LanceDB, conservative embed tuning: - rows collected: 3154 - rows uploaded: 3147 - queries: 115 - recall@1/5/10: 0.6261 / 0.8783 / 0.9391 - total_secs: 121.79 - -LanceDB, default embed tuning: - rows collected: 3154 - rows uploaded: 3147 - queries: 115 - recall@1/5/10: 0.6348 / 0.8696 / 0.9304 - total_secs: 73.20 -``` - -The default-tuning LanceDB run used no `--embed-actors` or -`--embed-batch-size` overrides. It used the normal pipeline embed settings and -confirmed that explicit hosted model alignment is sufficient for the endpoint -path. - -Some non-embedded client backends can show slower uploads for this local -deployment when they rely on object-store bulk import plus collection -load/refresh for datasets above the streaming threshold. That behavior is -backend-specific. diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 9806e0d877..8e7f39ab75 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -15,10 +15,6 @@ Format names and internal stages are not root commands. Use `retriever ingest` for PDF, HTML, TXT, image, Office, audio, and video inputs; it owns extraction, embedding, and index creation as one workflow. -`retriever pipeline run` remains callable as hidden compatibility while existing -development callers migrate. It is not shown in root help and is not the -preferred product ingest path. - ## Public ingest shape `retriever ingest` defaults to local, in-process ingest: @@ -349,16 +345,3 @@ Ingested 20 file(s) -> 1940 row(s) through retriever service http://localhost:76 Use `--dry-run` on any ingest mode to inspect the resolved request without creating an ingestor or contacting the service. - -## Development / compatibility command - -`retriever pipeline run` remains available, but hidden from root help, for -pipeline-specific behavior such as: - -- `--save-intermediate` Parquet artifacts. -- runtime metrics and pipeline reports. -- eval, recall, harness, or BEIR/QA workflows. -- legacy compatibility while callers migrate to `retriever ingest` and - `retriever query`. - -Run `retriever pipeline run --help` for the compatibility command flag list. diff --git a/nemo_retriever/docs/cli/benchmarking.md b/nemo_retriever/docs/cli/benchmarking.md index 72a9f69f3b..dd401c0f9b 100644 --- a/nemo_retriever/docs/cli/benchmarking.md +++ b/nemo_retriever/docs/cli/benchmarking.md @@ -56,7 +56,6 @@ enabled in the resolved benchmark query config. Set `query.agentic: true` in a code-owned benchmark or runfile, or use repeatable `--set` overrides on the CLI. The agentic harness path runs the same ReAct retrieval graph used by root query, but only after ingest and only for BEIR evaluation (`evaluation.mode: beir`). -`retriever pipeline run` does not expose agentic evaluation flags. Minimal BEIR override example: @@ -89,10 +88,8 @@ Useful agentic query overrides: ### Image storage For normal ingest, configure image persistence on `retriever ingest` with -`--store-images-uri ` (local path or fsspec URI). The harness does not -configure store directly; `retriever pipeline run --store-images-uri ` -remains available for pipeline-specific compatibility workflows. Stored assets -follow `--embed-granularity` (page vs element images). +`--store-images-uri ` (local path or fsspec URI). Stored assets follow +`--embed-granularity` (page vs element images). ## Per-stage micro-benchmarks diff --git a/nemo_retriever/harness/HANDOFF.md b/nemo_retriever/harness/HANDOFF.md index 87d454053c..9746127a10 100644 --- a/nemo_retriever/harness/HANDOFF.md +++ b/nemo_retriever/harness/HANDOFF.md @@ -25,8 +25,7 @@ the harness. - `service` is a system-under-test mode that uses an endpoint supplied by the caller; Helm is only an optional outer provisioning mechanism. -The harness must not route these commands through `retriever pipeline run` or -`nemo_retriever.examples.graph_pipeline`. +The harness calls the shared ingest and query workflow modules directly. ## Implementation Map diff --git a/nemo_retriever/harness/README.md b/nemo_retriever/harness/README.md index 73ed9bb412..04fb9df1a7 100644 --- a/nemo_retriever/harness/README.md +++ b/nemo_retriever/harness/README.md @@ -80,7 +80,7 @@ uses the product service APIs for ingest and query while preserving the same - `post-slack`: preview or post existing artifacts; it never executes a run. - `diff`: compare two run artifact directories by `results.json` summary metrics. -Legacy graph-pipeline, sweep, recurring-job, runner, reporting-UI, and portal +Legacy sweep, recurring-job, runner, reporting-UI, and portal commands are not part of this CLI surface. Scheduling and deployment belong to separate infrastructure, not the benchmark harness. @@ -369,9 +369,8 @@ plans before launching an expensive run. ## Implementation Boundary -The harness does not shell out to `retriever ingest`, `retriever query`, or -`retriever pipeline run`. It calls the same Python workflow/planning APIs used -by the CLI: +The harness does not shell out to `retriever ingest` or `retriever query`. It +calls the same Python workflow/planning modules used by the CLI: - ingest: `resolve_ingest_plan(...)` and `run_ingest_workflow(...)` - query: `resolve_query_plan(...)` and shared query workflow objects diff --git a/nemo_retriever/pyproject.toml b/nemo_retriever/pyproject.toml index 5248693165..543f07f7ef 100644 --- a/nemo_retriever/pyproject.toml +++ b/nemo_retriever/pyproject.toml @@ -88,7 +88,7 @@ dependencies = [ [project.optional-dependencies] # Core ``pip install nemo-retriever`` (no extras) supports remote HTTP NIM and the -# ``retriever pipeline`` CLI without ``tritonclient``. Triton gRPC is imported only +# ``retriever ingest`` / ``retriever query`` CLIs without ``tritonclient``. Triton gRPC is imported only # when a gRPC client is constructed; install ``nemo-retriever[local]`` for on-box # Triton + GPU stacks (``tritonclient`` is listed under ``local``). diff --git a/nemo_retriever/src/nemo_retriever/cli/main.py b/nemo_retriever/src/nemo_retriever/cli/main.py index 2c8488ecfa..6b9b233cc9 100644 --- a/nemo_retriever/src/nemo_retriever/cli/main.py +++ b/nemo_retriever/src/nemo_retriever/cli/main.py @@ -38,7 +38,6 @@ ("benchmark", "nemo_retriever.tools.benchmark", "app", True), ("recall", "nemo_retriever.tools.recall", "app", True), ("skill-eval", "nemo_retriever.tools.skill_eval", "app", True), - ("pipeline", "nemo_retriever.cli.pipeline.__main__", "app", True), ] for _name, _module, _attr, _hidden in _LAZY_SUBAPPS: diff --git a/nemo_retriever/src/nemo_retriever/cli/pipeline/__init__.py b/nemo_retriever/src/nemo_retriever/cli/pipeline/__init__.py deleted file mode 100644 index 549f5ce291..0000000000 --- a/nemo_retriever/src/nemo_retriever/cli/pipeline/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Legacy end-to-end ingestion pipeline subcommand for the ``retriever`` CLI. - -This package owns the ``retriever pipeline`` Typer application. Local pipeline -runs delegate ingest graph construction to the core ingest package; service -runs still build the service client path directly. - -It is registered on the ``retriever`` CLI as the ``pipeline`` subcommand:: - - retriever pipeline run [OPTIONS] - -The implementation historically lived in -``nemo_retriever/examples/graph_pipeline.py``; that module is now a thin -backward-compat shim that re-exports the same Typer app from -:mod:`nemo_retriever.pipeline.__main__`. - -``app`` and ``run`` are exposed via lazy attribute access so that -``python -m nemo_retriever.pipeline`` can import the ``__main__`` module -cleanly (without a re-import warning). -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -__all__ = ["app", "run"] - - -if TYPE_CHECKING: - from nemo_retriever.cli.pipeline.__main__ import app, run # noqa: F401 - - -def __getattr__(name: str) -> Any: - if name in {"app", "run"}: - from nemo_retriever.cli.pipeline import __main__ as _main - - return getattr(_main, name) - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/nemo_retriever/src/nemo_retriever/cli/pipeline/__main__.py b/nemo_retriever/src/nemo_retriever/cli/pipeline/__main__.py deleted file mode 100644 index cb776473fb..0000000000 --- a/nemo_retriever/src/nemo_retriever/cli/pipeline/__main__.py +++ /dev/null @@ -1,1765 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Development Typer CLI for the end-to-end graph ingestion pipeline. - -Registered on the ``retriever`` CLI as the ``pipeline`` subcommand. -For user-facing workflows, prefer ``retriever ingest`` and ``retriever query``. - -Examples:: - - # Batch mode (Ray) with PDF extraction + embedding - retriever pipeline run /data/pdfs \\ - --run-mode batch \\ - --embed-invoke-url http://localhost:8000/v1 - - # In-process mode (no Ray) for quick local testing - retriever pipeline run /data/pdfs \\ - --run-mode inprocess \\ - --ocr-invoke-url http://localhost:9000/v1 - - # Service mode (delegate to a running retriever service) - retriever pipeline run /data/pdfs \\ - --run-mode service \\ - --service-url http://localhost:7670 - - # Save extraction Parquet for full-page markdown (page index / export) - retriever pipeline run /data/pdfs \\ - --save-intermediate /path/to/extracted_parquet_dir - - # Override the default VDB backend/configuration - retriever pipeline run /data/pdfs \\ - --vdb-op \\ - --vdb-kwargs-json '' - - # Extract + embed only (skip in-graph VDB upload) - retriever pipeline run /data/pdfs \\ - --no-vdb - - # Sidecar metadata (merged into each chunk's content_metadata, same triplet as nv-ingest-client) - retriever pipeline run /data/pdfs \\ - --meta-dataframe ./meta.csv \\ - --meta-source-field source \\ - --meta-fields meta_a,meta_b -""" - -from __future__ import annotations - -import glob as _glob -import json -import logging -import os -import time -from dataclasses import replace -from pathlib import Path -from typing import Any, Optional, TextIO - -import pandas as pd -import typer - -from nemo_retriever.ingest import service as ingest_service -from nemo_retriever.ingest.execution import execute_ingest_plan -from nemo_retriever.ingest.plan import ( - IngestCaptionOptions, - IngestChunkOptions, - IngestDedupOptions, - IngestEmbedBatchOptions, - IngestEmbedOptions, - IngestExtractBatchOptions, - IngestExtractOptions, - IngestImageStoreOptions, - IngestMediaOptions, - IngestPlanRequest, - IngestRuntimeOptions, - IngestSourceOptions, - IngestStorageOptions, - resolve_ingest_plan, -) -from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL -from nemo_retriever.common.input_files import resolve_input_patterns -from nemo_retriever.common.modality.caption.model_profiles import DEFAULT_LOCAL_CAPTION_MODEL_ID -from nemo_retriever.common.params import ( - CaptionParams, - DedupParams, - EmbedParams, - ExtractParams, - StoreParams, - TextChunkParams, - VdbUploadParams, -) -from nemo_retriever.common.params.models import BatchTuningParams -from nemo_retriever.common.remote_auth import resolve_remote_api_key - -logger = logging.getLogger(__name__) - -app = typer.Typer( - help=( - "Development graph ingestion pipeline (compatibility wrapper around " - "retriever ingest/query paths; prefer retriever ingest and retriever query)." - ) -) - -DEFAULT_VDB_OP = "lancedb" - -# Help panel labels (keep stable so --help groupings read consistently). -_PANEL_IO = "I/O and Execution" -_PANEL_EXTRACT = "PDF / Document Extraction" -_PANEL_REMOTE = "Remote NIM Endpoints" -_PANEL_EMBED = "Embedding" -_PANEL_DEDUP_CAPTION = "Dedup and Caption" -_PANEL_STORE_CHUNK = "Storage and Text Chunking" -_PANEL_AUDIO = "Audio" -_PANEL_VIDEO = "Video" -_PANEL_RAY = "Ray / Batch Tuning" -_PANEL_VDB = "VDB and Outputs" -_PANEL_EVAL = "Evaluation (Recall / BEIR)" -_PANEL_OBS = "Observability" -_PANEL_SERVICE = "Service Mode" - - -# CLI flags that have no effect in --run-mode=service: either silently -# overridden by retriever-service.yaml (server-owned endpoints / models), -# bound to local execution (Ray actors, GPU placement), or never wired -# through the service ingestor (VDB upload is handled server-side; audio -# and video extract paths still run locally). Flags wired into the -# service ``PipelineSpec`` by ``ingest.service.build_service_ingestor`` — extract knobs, embed -# granularity / modality, dedup threshold, caption behaviour, text chunk -# config, ``--store-images-uri`` — are intentionally NOT in this list and -# pass through to ``ServiceIngestor``; the server's -# ``_DEFAULT_ALLOWED_*_KEYS`` allowlists are the final authority on which -# keys survive. -_SERVICE_INCOMPATIBLE_FLAGS: tuple[tuple[str, str], ...] = ( - # Remote NIM endpoints + model names — server-owned via retriever-service.yaml - ("--page-elements-invoke-url", "page_elements_invoke_url"), - ("--ocr-invoke-url", "ocr_invoke_url"), - ("--ocr-lang", "ocr_lang"), - ("--table-structure-invoke-url", "table_structure_invoke_url"), - ("--caption-invoke-url", "caption_invoke_url"), - ("--caption-model-name", "caption_model_name"), - # Local-execution knobs (no in-cluster equivalent) - ("--local-ingest-embed-backend", "local_ingest_embed_backend"), - ("--caption-device", "caption_device"), - ("--caption-gpu-memory-utilization", "caption_gpu_memory_utilization"), - ("--caption-gpus-per-actor", "caption_gpus_per_actor"), - # Audio (service path is pdf-only today) - ("--segment-audio/--no-segment-audio", "segment_audio"), - ("--audio-split-type", "audio_split_type"), - ("--audio-split-interval", "audio_split_interval"), - # Video (service path is pdf-only today) - ("--video-extract-audio/--no-video-extract-audio", "video_extract_audio"), - ("--video-extract-frames/--no-video-extract-frames", "video_extract_frames"), - ("--video-frame-fps", "video_frame_fps"), - ("--video-frame-dedup/--no-video-frame-dedup", "video_frame_dedup"), - ("--video-frame-text-dedup/--no-video-frame-text-dedup", "video_frame_text_dedup"), - ("--video-frame-text-dedup-max-dropped-frames", "video_frame_text_dedup_max_dropped_frames"), - ("--video-av-fuse/--no-video-av-fuse", "video_av_fuse"), - # Ray / batch tuning — no analog when the worker is a service pod - ("--ray-address", "ray_address"), - ("--ray-log-to-driver/--no-ray-log-to-driver", "ray_log_to_driver"), - ("--ocr-actors", "ocr_actors"), - ("--ocr-batch-size", "ocr_batch_size"), - ("--ocr-cpus-per-actor", "ocr_cpus_per_actor"), - ("--ocr-gpus-per-actor", "ocr_gpus_per_actor"), - ("--page-elements-actors", "page_elements_actors"), - ("--page-elements-batch-size", "page_elements_batch_size"), - ("--page-elements-cpus-per-actor", "page_elements_cpus_per_actor"), - ("--page-elements-gpus-per-actor", "page_elements_gpus_per_actor"), - ("--embed-actors", "embed_actors"), - ("--embed-batch-size", "embed_batch_size"), - ("--embed-cpus-per-actor", "embed_cpus_per_actor"), - ("--embed-gpus-per-actor", "embed_gpus_per_actor"), - ("--store-actors", "store_actors"), - ("--pdf-split-batch-size", "pdf_split_batch_size"), - ("--pdf-extract-batch-size", "pdf_extract_batch_size"), - ("--pdf-extract-tasks", "pdf_extract_tasks"), - ("--pdf-extract-cpus-per-task", "pdf_extract_cpus_per_task"), - ("--nemotron-parse-actors", "nemotron_parse_actors"), - ("--nemotron-parse-gpus-per-actor", "nemotron_parse_gpus_per_actor"), - ("--nemotron-parse-batch-size", "nemotron_parse_batch_size"), - # In-graph VDB / sidecar metadata — service mode does VDB writes - # server-side via LanceDBWriteOperator and never wires these through - # the service ingestor (see ``enable_in_graph_vdb_upload`` gate). - ("--no-vdb", "no_vdb"), - ("--vdb-op", "vdb_op"), - ("--vdb-kwargs-json", "vdb_kwargs_json"), - ("--vdb-overwrite/--vdb-append", "vdb_overwrite"), - ("--meta-dataframe", "meta_dataframe"), - ("--meta-source-field", "meta_source_field"), - ("--meta-fields", "meta_fields"), - ("--meta-join-key", "meta_join_key"), -) - - -def _reject_service_incompatible_flags(ctx: typer.Context) -> None: - user_set: list[str] = [] - for cli_flag, param_name in _SERVICE_INCOMPATIBLE_FLAGS: - source = ctx.get_parameter_source(param_name) - if getattr(source, "name", None) in {"COMMANDLINE", "ENVIRONMENT"}: - user_set.append(cli_flag) - if not user_set: - return - raise typer.BadParameter( - "--run-mode=service delegates pipeline configuration to the " - "retriever service; the following flag(s) cannot be set on the " - "client and would be silently dropped: " + ", ".join(user_set) + ". " - "Remove them, or use --run-mode batch/inprocess to apply them locally. " - "Server-side pipeline configuration lives in retriever-service.yaml." - ) - - -# --------------------------------------------------------------------------- -# Logging helpers -# --------------------------------------------------------------------------- - - -class _TeeStream: - """Mirror stdout/stderr writes into a second stream (e.g. a log file).""" - - def __init__(self, primary: TextIO, mirror: TextIO) -> None: - self._primary = primary - self._mirror = mirror - - def write(self, data: str) -> int: - self._primary.write(data) - self._mirror.write(data) - return len(data) - - def flush(self) -> None: - self._primary.flush() - self._mirror.flush() - - def isatty(self) -> bool: - return bool(getattr(self._primary, "isatty", lambda: False)()) - - def fileno(self) -> int: - return int(getattr(self._primary, "fileno")()) - - def writable(self) -> bool: - return bool(getattr(self._primary, "writable", lambda: True)()) - - @property - def encoding(self) -> str: - return str(getattr(self._primary, "encoding", "utf-8")) - - -def _configure_logging(log_file: Optional[Path], *, debug: bool = False) -> tuple[Optional[TextIO], TextIO, TextIO]: - original_stdout = os.sys.stdout - original_stderr = os.sys.stderr - log_level = logging.DEBUG if debug else logging.INFO - if log_file is None: - logging.basicConfig( - level=log_level, - format="%(asctime)s %(levelname)s %(name)s: %(message)s", - force=True, - ) - return None, original_stdout, original_stderr - - target = Path(log_file).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - fh = open(target, "a", encoding="utf-8", buffering=1) - os.sys.stdout = _TeeStream(os.sys.__stdout__, fh) - os.sys.stderr = _TeeStream(os.sys.__stderr__, fh) - logging.basicConfig( - level=log_level, - format="%(asctime)s %(levelname)s %(name)s: %(message)s", - handlers=[logging.StreamHandler(os.sys.stdout)], - force=True, - ) - logger.info("Writing combined pipeline logs to %s", str(target)) - return fh, original_stdout, original_stderr - - -# --------------------------------------------------------------------------- -# Small utilities (summaries, file patterns) -# --------------------------------------------------------------------------- - - -def _write_runtime_summary( - runtime_metrics_dir: Optional[Path], - runtime_metrics_prefix: Optional[str], - payload: dict[str, object], -) -> None: - if runtime_metrics_dir is None and not runtime_metrics_prefix: - return - - target_dir = Path(runtime_metrics_dir or Path.cwd()).expanduser().resolve() - target_dir.mkdir(parents=True, exist_ok=True) - prefix = (runtime_metrics_prefix or "run").strip() or "run" - target = target_dir / f"{prefix}.runtime.summary.json" - target.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") - - -def _count_input_units(result_df) -> int: - if "source_id" in result_df.columns: - return int(result_df["source_id"].nunique()) - if "source_path" in result_df.columns: - return int(result_df["source_path"].nunique()) - return int(len(result_df.index)) - - -def _format_vdb_target(vdb_op: str, vdb_kwargs: Optional[dict[str, Any]]) -> str: - """``/`` for a vdb destination, mirroring the lancedb-specific - fallbacks used downstream when those keys are absent from ``vdb_kwargs``.""" - kw = vdb_kwargs or {} - uri = kw.get("uri") or kw.get("lancedb_uri") or ("lancedb" if vdb_op == "lancedb" else "?") - table = kw.get("table_name") or kw.get("lancedb_table") or ("nv-ingest" if vdb_op == "lancedb" else "?") - return f"{uri}/{table}" - - -def _resolve_file_patterns(input_path: Path, input_type: str) -> list[str]: - """Resolve input paths to glob patterns, recursing into subdirectories. - - Uses :func:`~nemo_retriever.common.input_files.resolve_input_patterns` (``**`` - segments) and keeps only patterns that match at least one file, matching the - historical ``graph_pipeline`` / main-branch behavior. - """ - - input_path = Path(input_path) - if input_path.is_file(): - return [str(input_path)] - if not input_path.is_dir(): - raise typer.BadParameter(f"Path does not exist: {input_path}") - - if input_type not in {"pdf", "doc", "txt", "html", "image", "audio", "video"}: - raise typer.BadParameter(f"Unsupported --input-type: {input_type!r}") - - patterns = resolve_input_patterns(input_path, input_type) - matched = [p for p in patterns if _glob.glob(p, recursive=True)] - if not matched: - raise typer.BadParameter(f"No files found for input_type={input_type!r} in {input_path}") - logger.debug("Using recursive input globs: %s", matched) - return matched - - -# --------------------------------------------------------------------------- -# Parameter builders (split out from the old monolithic main()) -# --------------------------------------------------------------------------- - - -def _build_extract_params( - *, - method: Optional[str], - dpi: Optional[int], - extract_text: Optional[bool], - extract_tables: Optional[bool], - extract_charts: Optional[bool], - extract_infographics: Optional[bool], - extract_page_as_image: Optional[bool], - use_page_elements: Optional[bool], - use_table_structure: Optional[bool], - table_output_format: Optional[str], - extract_remote_api_key: Optional[str], - page_elements_invoke_url: Optional[str], - ocr_invoke_url: Optional[str], - ocr_version: str, - ocr_lang: Optional[str], - table_structure_invoke_url: Optional[str], - pdf_split_batch_size: int, - pdf_extract_batch_size: Optional[int], - pdf_extract_tasks: Optional[int], - pdf_extract_cpus_per_task: Optional[float], - page_elements_actors: Optional[int], - page_elements_batch_size: Optional[int], - page_elements_cpus_per_actor: Optional[float], - page_elements_gpus_per_actor: Optional[float], - ocr_actors: Optional[int], - ocr_batch_size: Optional[int], - ocr_cpus_per_actor: Optional[float], - ocr_gpus_per_actor: Optional[float], - nemotron_parse_actors: Optional[int], - nemotron_parse_batch_size: Optional[int], - nemotron_parse_gpus_per_actor: Optional[float], -) -> ExtractParams: - """Assemble :class:`ExtractParams` plus its :class:`BatchTuningParams`.""" - - extract_batch_tuning = BatchTuningParams( - **{ - k: v - for k, v in { - "pdf_split_batch_size": pdf_split_batch_size, - "pdf_extract_batch_size": pdf_extract_batch_size or None, - "pdf_extract_workers": pdf_extract_tasks or None, - "pdf_extract_num_cpus": pdf_extract_cpus_per_task or None, - "page_elements_batch_size": page_elements_batch_size or None, - "page_elements_workers": page_elements_actors or None, - "page_elements_cpus_per_actor": page_elements_cpus_per_actor or None, - "gpu_page_elements": ( - 0.0 - if page_elements_invoke_url - else (page_elements_gpus_per_actor if page_elements_gpus_per_actor is not None else None) - ), - "ocr_inference_batch_size": ocr_batch_size or None, - "ocr_workers": ocr_actors or None, - "ocr_cpus_per_actor": ocr_cpus_per_actor or None, - "gpu_ocr": ( - 0.0 if ocr_invoke_url else (ocr_gpus_per_actor if ocr_gpus_per_actor is not None else None) - ), - "nemotron_parse_batch_size": nemotron_parse_batch_size or None, - "nemotron_parse_workers": nemotron_parse_actors or None, - "gpu_nemotron_parse": ( - nemotron_parse_gpus_per_actor if nemotron_parse_gpus_per_actor is not None else None - ), - }.items() - if v is not None - } - ) - return ExtractParams( - **{ - k: v - for k, v in { - "method": method, - "dpi": int(dpi) if dpi is not None else None, - "extract_text": extract_text, - "extract_tables": extract_tables, - "extract_charts": extract_charts, - "extract_infographics": extract_infographics, - "extract_page_as_image": extract_page_as_image, - "use_page_elements": use_page_elements, - "api_key": extract_remote_api_key, - "page_elements_invoke_url": page_elements_invoke_url, - "ocr_invoke_url": ocr_invoke_url, - "ocr_version": ocr_version, - "ocr_lang": ocr_lang, - "table_structure_invoke_url": table_structure_invoke_url, - "use_table_structure": use_table_structure, - "table_output_format": table_output_format, - "inference_batch_size": page_elements_batch_size or None, - "batch_tuning": extract_batch_tuning, - }.items() - if v is not None - } - ) - - -def _build_embed_params( - *, - embed_model_name: str, - embed_invoke_url: Optional[str], - embed_remote_api_key: Optional[str], - embed_modality: str, - text_elements_modality: Optional[str], - structured_elements_modality: Optional[str], - embed_granularity: str, - embed_actors: Optional[int], - embed_batch_size: Optional[int], - embed_cpus_per_actor: Optional[float], - embed_gpus_per_actor: Optional[float], - local_ingest_embed_backend: str = "vllm", -) -> EmbedParams: - """Assemble :class:`EmbedParams` plus its :class:`BatchTuningParams`.""" - - embed_batch_tuning = BatchTuningParams( - **{ - k: v - for k, v in { - "embed_batch_size": embed_batch_size or None, - "embed_workers": embed_actors or None, - "embed_cpus_per_actor": embed_cpus_per_actor or None, - "gpu_embed": ( - 0.0 if embed_invoke_url else (embed_gpus_per_actor if embed_gpus_per_actor is not None else None) - ), - }.items() - if v is not None - } - ) - return EmbedParams( - **{ - k: v - for k, v in { - "model_name": embed_model_name, - "embed_invoke_url": embed_invoke_url, - "api_key": embed_remote_api_key, - "embed_modality": embed_modality, - "text_elements_modality": text_elements_modality, - "structured_elements_modality": structured_elements_modality, - "embed_granularity": embed_granularity, - "local_ingest_embed_backend": local_ingest_embed_backend, - "batch_tuning": embed_batch_tuning, - "inference_batch_size": embed_batch_size or None, - }.items() - if v is not None - } - ) - - -def _parse_vdb_kwargs_json(vdb_kwargs_json: Optional[str]) -> dict[str, Any]: - """Parse opaque nv-ingest-client VDB constructor kwargs from CLI JSON.""" - if vdb_kwargs_json: - try: - parsed = json.loads(vdb_kwargs_json) - except json.JSONDecodeError as exc: - raise typer.BadParameter(f"--vdb-kwargs-json must be valid JSON: {exc}") from exc - if not isinstance(parsed, dict): - raise typer.BadParameter("--vdb-kwargs-json must decode to a JSON object.") - return parsed - return {} - - -def _collect_results(run_mode: str, result: Any) -> tuple[list[dict[str, Any]], Any, float, int]: - """Materialize the graph result into a list of records + DataFrame. - - Ingest may return a ``pandas.DataFrame`` (in-process and batch graph paths, - where batch materializes via ``ds.to_pandas()`` in the executor) or a - :class:`~nemo_retriever.service.service_ingestor.ServiceIngestResult` (service mode); - normalize to a consistent ``(records, DataFrame, secs, units)`` tuple. - - Returns ``(records, result_df, ray_download_secs, num_input_units)``. - """ - - if run_mode == "service": - records = list(result) - result_df = pd.DataFrame(records) if records else pd.DataFrame() - num_units = getattr(result, "total_pages", 0) or len(records) - return records, result_df, 0.0, num_units - - if isinstance(result, pd.DataFrame): - result_df = result - else: - result_df = result.to_pandas() - records = result_df.to_dict("records") - ray_download_time = 0.0 - - return records, result_df, float(ray_download_time), _count_input_units(result_df) - - -def _count_uploadable_vdb_records(records: list[dict[str, Any]]) -> int: - """Count records that will survive conversion into the client VDB record contract.""" - - from nemo_retriever.common.vdb.records import to_client_vdb_records - - return sum(len(batch) for batch in to_client_vdb_records(records)) - - -def _run_evaluation( - *, - evaluation_mode: str, - vdb_op: str, - vdb_kwargs: dict[str, Any], - embed_model_name: str, - embed_invoke_url: Optional[str], - embed_remote_api_key: Optional[str], - embed_modality: str, - query_csv: Path, - recall_match_mode: str, - audio_match_tolerance_secs: float, - reranker: Optional[bool], - reranker_model_name: str, - reranker_invoke_url: Optional[str], - reranker_api_key: str, - local_reranker_backend: str, - local_hf_batch_size: int, - local_query_max_length: int, - beir_loader: Optional[str], - beir_dataset_name: Optional[str], - beir_split: str, - beir_query_language: Optional[str], - beir_doc_id_field: Optional[str], - beir_k: list[int], - local_query_embed_backend: str = "hf", - run_mode: str = "inprocess", - service_url: Optional[str] = None, - service_api_token: Optional[str] = None, -) -> tuple[str, float, dict[str, float], Optional[int], bool]: - """Run audio recall or BEIR evaluation. - - Returns ``(label, elapsed_secs, metrics, query_count, ran)``. When the - query CSV is missing in audio recall mode, ``ran`` is ``False`` and the - caller should skip metric recording. - """ - - if evaluation_mode == "none": - return "None", 0.0, {}, None, False - - from nemo_retriever.models import resolve_embed_model - - embed_model = resolve_embed_model(str(embed_model_name)) - eval_vdb_kwargs = dict(vdb_kwargs or {}) - - if evaluation_mode == "beir": - from nemo_retriever.tools.recall.beir import BeirConfig, resolve_beir_dataset_options - - beir_options = resolve_beir_dataset_options( - dataset_name=beir_dataset_name, - loader=beir_loader, - doc_id_field=beir_doc_id_field, - ks=beir_k, - ) - if not beir_options.loader: - raise ValueError("--beir-loader is required when --evaluation-mode=beir") - if not beir_options.dataset_name: - raise ValueError("--beir-dataset-name is required when --evaluation-mode=beir") - - lancedb_uri = str(eval_vdb_kwargs.get("uri") or eval_vdb_kwargs.get("lancedb_uri") or "lancedb") - lancedb_table = str(eval_vdb_kwargs.get("table_name") or eval_vdb_kwargs.get("lancedb_table") or "nv-ingest") - - cfg = BeirConfig( - lancedb_uri=lancedb_uri, - lancedb_table=lancedb_table, - embedding_model=embed_model, - loader=str(beir_options.loader), - dataset_name=str(beir_options.dataset_name), - split=str(beir_split), - query_language=beir_query_language, - doc_id_field=str(beir_options.doc_id_field), - ks=beir_options.ks, - embedding_http_endpoint=embed_invoke_url, - embedding_api_key=embed_remote_api_key or "", - hybrid=bool(eval_vdb_kwargs.get("hybrid", False)), - nprobes=int(eval_vdb_kwargs.get("nprobes", 0) or 0), - refine_factor=int(eval_vdb_kwargs.get("refine_factor", 10) or 10), - reranker=bool(reranker), - reranker_model_name=str(reranker_model_name), - reranker_endpoint=reranker_invoke_url, - reranker_api_key=reranker_api_key, - local_reranker_backend=local_reranker_backend, - local_hf_batch_size=int(local_hf_batch_size), - local_query_max_length=int(local_query_max_length), - local_query_embed_backend=local_query_embed_backend, - service_url=service_url if run_mode == "service" else None, - service_api_token=service_api_token, - ) - - evaluation_start = time.perf_counter() - if run_mode == "service" and service_url: - from nemo_retriever.tools.recall.beir import evaluate_service_beir - - beir_dataset, _raw_hits, _run, metrics = evaluate_service_beir(cfg) - else: - if str(vdb_op).strip().lower() != "lancedb": - raise ValueError("--evaluation-mode=beir currently requires --vdb-op=lancedb") - from nemo_retriever.tools.recall.beir import evaluate_lancedb_beir - - beir_dataset, _raw_hits, _run, metrics = evaluate_lancedb_beir(cfg) - return "BEIR", time.perf_counter() - evaluation_start, metrics, len(beir_dataset.query_ids), True - - if evaluation_mode != "audio_recall": - raise ValueError(f"Unsupported --evaluation-mode: {evaluation_mode!r}") - - if recall_match_mode != "audio_segment": - raise ValueError("Audio recall evaluation is only supported for audio_segment matching") - - # Legacy scorer is retained for audio segment evaluation only. - query_csv_path = Path(query_csv) - if not query_csv_path.exists(): - logger.warning("Query CSV not found at %s; skipping audio recall evaluation.", query_csv_path) - return "Audio Recall", 0.0, {}, None, False - - from nemo_retriever.tools.recall.core import RecallConfig, retrieve_and_score - - recall_cfg = RecallConfig( - vdb_op=str(vdb_op), - vdb_kwargs=eval_vdb_kwargs, - query_embedder=embed_model, - embedding_endpoint=embed_invoke_url, - embedding_api_key=embed_remote_api_key or "", - embedding_use_grpc=False if embed_invoke_url else None, - top_k=10, - ks=(1, 5, 10), - match_mode=recall_match_mode, - audio_match_tolerance_secs=float(audio_match_tolerance_secs), - reranker=reranker_model_name if reranker else None, - reranker_endpoint=reranker_invoke_url, - reranker_api_key=reranker_api_key, - local_reranker_backend=local_reranker_backend, - local_hf_batch_size=int(local_hf_batch_size), - local_query_max_length=int(local_query_max_length), - embed_modality=embed_modality, - local_query_embed_backend=local_query_embed_backend, - ) - evaluation_start = time.perf_counter() - df_query, _gold, _raw_hits, _retrieved_keys, metrics = retrieve_and_score(query_csv=query_csv_path, cfg=recall_cfg) - return "Audio Recall", time.perf_counter() - evaluation_start, metrics, len(df_query.index), True - - -# --------------------------------------------------------------------------- -# Typer command: `retriever pipeline run` -# --------------------------------------------------------------------------- - - -@app.command("run") -def run( - ctx: typer.Context, - input_path: Path = typer.Argument( - ..., - help="File or directory of documents to ingest.", - path_type=Path, - ), - # --- I/O and execution ------------------------------------------------ - # Default to in-process execution for small-document initial testing; batch - # remains explicit for Ray Data throughput runs. - run_mode: str = typer.Option( - "inprocess", - "--run-mode", - help=( - "Execution mode: 'inprocess' (pandas, no Ray), 'batch' (Ray Data), " - "or 'service' (remote retriever service)." - ), - rich_help_panel=_PANEL_IO, - ), - input_type: str = typer.Option( - "pdf", - "--input-type", - help="Input type: 'pdf', 'doc', 'txt', 'html', 'image', or 'audio'.", - rich_help_panel=_PANEL_IO, - ), - debug: bool = typer.Option( - False, "--debug/--no-debug", help="Enable debug-level logging.", rich_help_panel=_PANEL_IO - ), - log_file: Optional[Path] = typer.Option( - None, "--log-file", path_type=Path, dir_okay=False, rich_help_panel=_PANEL_IO - ), - quiet: bool = typer.Option( - False, - "--quiet", - help=( - "Suppress verbose output on success: progress bars, HuggingFace " - "downloads, vLLM init logs, Ray worker stdout, and INFO-level " - "pipeline status lines. On error, the fd-captured output is " - "flushed to stderr for debugging. Implies --no-ray-log-to-driver." - ), - rich_help_panel=_PANEL_IO, - ), - # --- PDF / document extraction --------------------------------------- - method: Optional[str] = typer.Option( - None, "--method", help="PDF text extraction method.", rich_help_panel=_PANEL_EXTRACT - ), - dpi: Optional[int] = typer.Option( - None, "--dpi", min=72, help="Render DPI for PDF page images.", rich_help_panel=_PANEL_EXTRACT - ), - extract_text: Optional[bool] = typer.Option( - None, "--extract-text/--no-extract-text", rich_help_panel=_PANEL_EXTRACT - ), - extract_tables: Optional[bool] = typer.Option( - None, "--extract-tables/--no-extract-tables", rich_help_panel=_PANEL_EXTRACT - ), - extract_charts: Optional[bool] = typer.Option( - None, "--extract-charts/--no-extract-charts", rich_help_panel=_PANEL_EXTRACT - ), - extract_infographics: Optional[bool] = typer.Option( - None, "--extract-infographics/--no-extract-infographics", rich_help_panel=_PANEL_EXTRACT - ), - extract_page_as_image: Optional[bool] = typer.Option( - None, - "--extract-page-as-image/--no-extract-page-as-image", - rich_help_panel=_PANEL_EXTRACT, - ), - use_page_elements: Optional[bool] = typer.Option( - None, - "--use-page-elements/--no-use-page-elements", - rich_help_panel=_PANEL_EXTRACT, - help=( - "Run PageElementDetection (layout/yolox). Auto-skipped when no downstream stage " - "(TableStructure or OCR) consumes its output. Pass --no-use-page-elements " - "to force-skip for a faster text-only ingest." - ), - ), - use_table_structure: Optional[bool] = typer.Option(None, "--use-table-structure", rich_help_panel=_PANEL_EXTRACT), - table_output_format: Optional[str] = typer.Option(None, "--table-output-format", rich_help_panel=_PANEL_EXTRACT), - # --- Remote NIM endpoints -------------------------------------------- - api_key: Optional[str] = typer.Option( - None, - "--api-key", - help="Bearer token for remote NIM endpoints.", - rich_help_panel=_PANEL_REMOTE, - ), - page_elements_invoke_url: Optional[str] = typer.Option( - None, "--page-elements-invoke-url", rich_help_panel=_PANEL_REMOTE - ), - ocr_invoke_url: Optional[str] = typer.Option(None, "--ocr-invoke-url", rich_help_panel=_PANEL_REMOTE), - ocr_version: str = typer.Option( - "v2", - "--ocr-version", - help="OCR engine: 'v2' (default, multilingual, higher throughput) or 'v1' (legacy, English-only).", - rich_help_panel=_PANEL_REMOTE, - ), - ocr_lang: Optional[str] = typer.Option( - None, - "--ocr-lang", - help="OCR language selector for v2: 'multi' (default) or 'english'. Not valid with --ocr-version v1.", - rich_help_panel=_PANEL_REMOTE, - ), - table_structure_invoke_url: Optional[str] = typer.Option( - None, "--table-structure-invoke-url", rich_help_panel=_PANEL_REMOTE - ), - embed_invoke_url: Optional[str] = typer.Option(None, "--embed-invoke-url", rich_help_panel=_PANEL_REMOTE), - # --- Embedding -------------------------------------------------------- - embed_model_name: str = typer.Option(VL_EMBED_MODEL, "--embed-model-name", rich_help_panel=_PANEL_EMBED), - embed_modality: str = typer.Option("text", "--embed-modality", rich_help_panel=_PANEL_EMBED), - embed_granularity: str = typer.Option("element", "--embed-granularity", rich_help_panel=_PANEL_EMBED), - local_ingest_embed_backend: str = typer.Option( - "vllm", - "--local-ingest-embed-backend", - help="Local ingest-time text embedder when --embed-invoke-url is unset: vllm or hf. VL models always use hf.", - rich_help_panel=_PANEL_EMBED, - ), - text_elements_modality: Optional[str] = typer.Option( - None, "--text-elements-modality", rich_help_panel=_PANEL_EMBED - ), - structured_elements_modality: Optional[str] = typer.Option( - None, "--structured-elements-modality", rich_help_panel=_PANEL_EMBED - ), - # --- Dedup / caption ------------------------------------------------- - dedup: Optional[bool] = typer.Option(None, "--dedup/--no-dedup", rich_help_panel=_PANEL_DEDUP_CAPTION), - dedup_iou_threshold: float = typer.Option(0.45, "--dedup-iou-threshold", rich_help_panel=_PANEL_DEDUP_CAPTION), - caption: bool = typer.Option(False, "--caption/--no-caption", rich_help_panel=_PANEL_DEDUP_CAPTION), - caption_invoke_url: Optional[str] = typer.Option( - None, "--caption-invoke-url", rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_model_name: str = typer.Option( - DEFAULT_LOCAL_CAPTION_MODEL_ID, - "--caption-model-name", - help=( - "VLM caption model. Defaults to Nemotron 3 Nano Omni 30B BF16 for local vLLM execution; " - "use an API model ID when --caption-invoke-url targets a hosted or deployed endpoint." - ), - rich_help_panel=_PANEL_DEDUP_CAPTION, - ), - caption_device: Optional[str] = typer.Option(None, "--caption-device", rich_help_panel=_PANEL_DEDUP_CAPTION), - caption_context_text_max_chars: int = typer.Option( - 0, "--caption-context-text-max-chars", rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_gpu_memory_utilization: float = typer.Option( - 0.5, "--caption-gpu-memory-utilization", rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_gpus_per_actor: Optional[float] = typer.Option( - None, "--caption-gpus-per-actor", max=1.0, rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_temperature: float = typer.Option( - 1.0, "--caption-temperature", min=0.0, max=2.0, rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_top_p: Optional[float] = typer.Option( - None, "--caption-top-p", min=0.0, max=1.0, rich_help_panel=_PANEL_DEDUP_CAPTION - ), - caption_max_tokens: int = typer.Option(1024, "--caption-max-tokens", min=1, rich_help_panel=_PANEL_DEDUP_CAPTION), - # --- Storage and text chunking -------------------------------------- - store_images_uri: Optional[str] = typer.Option( - None, - "--store-images-uri", - help="Store extracted images to this URI.", - rich_help_panel=_PANEL_STORE_CHUNK, - ), - text_chunk: bool = typer.Option(False, "--text-chunk", rich_help_panel=_PANEL_STORE_CHUNK), - text_chunk_max_tokens: Optional[int] = typer.Option( - None, "--text-chunk-max-tokens", rich_help_panel=_PANEL_STORE_CHUNK - ), - text_chunk_overlap_tokens: Optional[int] = typer.Option( - None, "--text-chunk-overlap-tokens", rich_help_panel=_PANEL_STORE_CHUNK - ), - # --- Ray / batch tuning --------------------------------------------- - # *_gpus_per_actor defaults are None (not 0.0) so we can distinguish - # "not set -> use heuristic" from "explicitly 0 -> no GPU". Other tuning - # defaults use 0/0.0 because those values are never valid explicit choices. - ray_address: Optional[str] = typer.Option(None, "--ray-address", rich_help_panel=_PANEL_RAY), - ray_log_to_driver: bool = typer.Option( - True, "--ray-log-to-driver/--no-ray-log-to-driver", rich_help_panel=_PANEL_RAY - ), - ocr_actors: Optional[int] = typer.Option(0, "--ocr-actors", rich_help_panel=_PANEL_RAY), - ocr_batch_size: Optional[int] = typer.Option(0, "--ocr-batch-size", rich_help_panel=_PANEL_RAY), - ocr_cpus_per_actor: Optional[float] = typer.Option(0.0, "--ocr-cpus-per-actor", rich_help_panel=_PANEL_RAY), - ocr_gpus_per_actor: Optional[float] = typer.Option( - None, "--ocr-gpus-per-actor", max=1.0, rich_help_panel=_PANEL_RAY - ), - page_elements_actors: Optional[int] = typer.Option(0, "--page-elements-actors", rich_help_panel=_PANEL_RAY), - page_elements_batch_size: Optional[int] = typer.Option(0, "--page-elements-batch-size", rich_help_panel=_PANEL_RAY), - page_elements_cpus_per_actor: Optional[float] = typer.Option( - 0.0, "--page-elements-cpus-per-actor", rich_help_panel=_PANEL_RAY - ), - page_elements_gpus_per_actor: Optional[float] = typer.Option( - None, "--page-elements-gpus-per-actor", max=1.0, rich_help_panel=_PANEL_RAY - ), - embed_actors: Optional[int] = typer.Option(0, "--embed-actors", rich_help_panel=_PANEL_RAY), - embed_batch_size: Optional[int] = typer.Option(0, "--embed-batch-size", rich_help_panel=_PANEL_RAY), - embed_cpus_per_actor: Optional[float] = typer.Option(0.0, "--embed-cpus-per-actor", rich_help_panel=_PANEL_RAY), - embed_gpus_per_actor: Optional[float] = typer.Option( - None, "--embed-gpus-per-actor", max=1.0, rich_help_panel=_PANEL_RAY - ), - store_actors: Optional[int] = typer.Option( - 0, - "--store-actors", - min=0, - help=( - "Maximum StoreOperator Ray actors. Store sinks autoscale from one actor to this cap; " - "0 uses the default cap." - ), - rich_help_panel=_PANEL_RAY, - ), - pdf_split_batch_size: int = typer.Option(1, "--pdf-split-batch-size", min=1, rich_help_panel=_PANEL_RAY), - pdf_extract_batch_size: Optional[int] = typer.Option(0, "--pdf-extract-batch-size", rich_help_panel=_PANEL_RAY), - pdf_extract_tasks: Optional[int] = typer.Option(0, "--pdf-extract-tasks", rich_help_panel=_PANEL_RAY), - pdf_extract_cpus_per_task: Optional[float] = typer.Option( - 0.0, "--pdf-extract-cpus-per-task", rich_help_panel=_PANEL_RAY - ), - nemotron_parse_actors: Optional[int] = typer.Option(0, "--nemotron-parse-actors", rich_help_panel=_PANEL_RAY), - nemotron_parse_gpus_per_actor: Optional[float] = typer.Option( - None, - "--nemotron-parse-gpus-per-actor", - min=0.0, - max=1.0, - rich_help_panel=_PANEL_RAY, - ), - nemotron_parse_batch_size: Optional[int] = typer.Option( - 0, "--nemotron-parse-batch-size", rich_help_panel=_PANEL_RAY - ), - # --- Audio ---------------------------------------------------------- - segment_audio: bool = typer.Option(False, "--segment-audio/--no-segment-audio", rich_help_panel=_PANEL_AUDIO), - audio_split_type: str = typer.Option("size", "--audio-split-type", rich_help_panel=_PANEL_AUDIO), - audio_split_interval: int = typer.Option(500000, "--audio-split-interval", min=1, rich_help_panel=_PANEL_AUDIO), - audio_match_tolerance_secs: float = typer.Option( - 2.0, "--audio-match-tolerance-secs", min=0.0, rich_help_panel=_PANEL_AUDIO - ), - # --- Video ---------------------------------------------------------- - video_extract_audio: bool = typer.Option( - True, - "--video-extract-audio/--no-video-extract-audio", - help=( - "Extract the video's audio track and run ASR. Disable to " - "produce frame-OCR rows only (no audio, no fusion)." - ), - rich_help_panel=_PANEL_VIDEO, - ), - video_extract_frames: bool = typer.Option( - True, - "--video-extract-frames/--no-video-extract-frames", - help=( - "Extract video frames and run frame OCR. Disable to produce " - "audio-only rows from video input (no frames, no OCR, no fusion)." - ), - rich_help_panel=_PANEL_VIDEO, - ), - video_frame_fps: float = typer.Option( - 0.5, - "--video-frame-fps", - min=0.001, - help="Frames per second to extract from videos (input_type=video).", - rich_help_panel=_PANEL_VIDEO, - ), - video_frame_dedup: bool = typer.Option( - True, - "--video-frame-dedup/--no-video-frame-dedup", - help="Drop content-hash-duplicate frames before OCR.", - rich_help_panel=_PANEL_VIDEO, - ), - video_frame_text_dedup: bool = typer.Option( - True, - "--video-frame-text-dedup/--no-video-frame-text-dedup", - help=( - "Merge consecutive frame OCR rows whose text is identical into " - "a single row spanning their combined time window." - ), - rich_help_panel=_PANEL_VIDEO, - ), - video_frame_text_dedup_max_dropped_frames: int = typer.Option( - 2, - "--video-frame-text-dedup-max-dropped-frames", - min=0, - help=( - "Tolerated dropped-frame count between same-text frames before they are " - "treated as separate runs. Converted to seconds at runtime via " - "max_gap_seconds = max_dropped_frames / fps." - ), - rich_help_panel=_PANEL_VIDEO, - ), - video_av_fuse: bool = typer.Option( - True, - "--video-av-fuse/--no-video-av-fuse", - help="Emit fused per-utterance rows (audio transcript + concurrent OCR).", - rich_help_panel=_PANEL_VIDEO, - ), - # --- Service mode --------------------------------------------------- - service_url: str = typer.Option( - "http://localhost:7670", - "--service-url", - help="Base URL of the retriever service (used only when --run-mode=service).", - rich_help_panel=_PANEL_SERVICE, - ), - service_concurrency: int = typer.Option( - 8, - "--service-concurrency", - min=1, - help="Maximum concurrent page uploads to the service (used only when --run-mode=service).", - rich_help_panel=_PANEL_SERVICE, - ), - service_api_token: Optional[str] = typer.Option( - None, - "--service-api-token", - help=( - "Bearer token for authenticating with the retriever service " - "(used only when --run-mode=service). " - "Falls back to $NEMO_RETRIEVER_API_TOKEN." - ), - envvar="NEMO_RETRIEVER_API_TOKEN", - rich_help_panel=_PANEL_SERVICE, - ), - # --- VDB / outputs -------------------------------------------------- - vdb_op: str = typer.Option( - DEFAULT_VDB_OP, - "--vdb-op", - help="nv-ingest-client VDB operator key for in-graph upload after embed/store (skipped with --no-vdb).", - rich_help_panel=_PANEL_VDB, - ), - vdb_kwargs_json: Optional[str] = typer.Option( - None, - "--vdb-kwargs-json", - help=( - "JSON object forwarded as constructor kwargs to the selected VDB operator " - "(optional; backends such as LanceDB use sensible defaults when omitted)." - ), - rich_help_panel=_PANEL_VDB, - ), - vdb_overwrite: Optional[bool] = typer.Option( - None, - "--vdb-overwrite/--vdb-append", - help=( - "Overwrite the target VDB table by default. Use --vdb-append to add rows to an existing " - "table without duplicate checks; rerunning the same inputs in append mode creates duplicates." - ), - rich_help_panel=_PANEL_VDB, - ), - no_vdb: bool = typer.Option( - False, - "--no-vdb", - help="Skip in-graph vector DB upload (extract+embed only).", - rich_help_panel=_PANEL_VDB, - ), - meta_dataframe: Optional[Path] = typer.Option( - None, - "--meta-dataframe", - help="CSV/JSON/Parquet sidecar metadata (requires --meta-source-field and --meta-fields).", - path_type=Path, - exists=True, - dir_okay=False, - file_okay=True, - rich_help_panel=_PANEL_VDB, - ), - meta_source_field: Optional[str] = typer.Option( - None, - "--meta-source-field", - help="Column in the metadata file that matches document path (same as nv-ingest-client).", - rich_help_panel=_PANEL_VDB, - ), - meta_fields: Optional[str] = typer.Option( - None, - "--meta-fields", - help="Comma-separated metadata columns to copy onto each chunk's content_metadata.", - rich_help_panel=_PANEL_VDB, - ), - meta_join_key: str = typer.Option( - "auto", - "--meta-join-key", - help="Document match key: auto (try source_id then source_name), source_id, or source_name.", - rich_help_panel=_PANEL_VDB, - ), - save_intermediate: Optional[Path] = typer.Option( - None, - "--save-intermediate", - help="Directory to write extraction results as Parquet (for full-page markdown / page index).", - path_type=Path, - file_okay=False, - dir_okay=True, - rich_help_panel=_PANEL_VDB, - ), - detection_summary_file: Optional[Path] = typer.Option( - None, "--detection-summary-file", path_type=Path, rich_help_panel=_PANEL_VDB - ), - runtime_metrics_dir: Optional[Path] = typer.Option( - None, "--runtime-metrics-dir", path_type=Path, rich_help_panel=_PANEL_OBS - ), - runtime_metrics_prefix: Optional[str] = typer.Option(None, "--runtime-metrics-prefix", rich_help_panel=_PANEL_OBS), - # --- Evaluation ----------------------------------------------------- - evaluation_mode: str = typer.Option( - "none", - "--evaluation-mode", - help="Post-ingest evaluation: none (default), audio_recall, beir, or qa.", - rich_help_panel=_PANEL_EVAL, - ), - query_csv: Path = typer.Option( - "./data/bo767_query_gt.csv", - "--query-csv", - path_type=Path, - rich_help_panel=_PANEL_EVAL, - ), - recall_match_mode: str = typer.Option("audio_segment", "--recall-match-mode", rich_help_panel=_PANEL_EVAL), - recall_details: bool = typer.Option(True, "--recall-details/--no-recall-details", rich_help_panel=_PANEL_EVAL), - local_query_embed_backend: str = typer.Option( - "hf", - "--local-query-embed-backend", - help="Local query embedding backend when --embed-invoke-url is unset: hf (default) or vllm.", - rich_help_panel=_PANEL_EVAL, - ), - reranker: Optional[bool] = typer.Option(False, "--reranker/--no-reranker", rich_help_panel=_PANEL_EVAL), - reranker_model_name: str = typer.Option(VL_RERANK_MODEL, "--reranker-model-name", rich_help_panel=_PANEL_EVAL), - reranker_invoke_url: Optional[str] = typer.Option( - None, - "--reranker-invoke-url", - help="OpenAI-compatible reranker NIM HTTP endpoint (recall and BEIR evaluation).", - rich_help_panel=_PANEL_EVAL, - ), - reranker_api_key: Optional[str] = typer.Option( - None, - "--reranker-api-key", - help="Bearer token for the reranker NIM; defaults to --api-key / NVIDIA_API_KEY when omitted.", - rich_help_panel=_PANEL_EVAL, - ), - local_reranker_backend: str = typer.Option( - "vllm", - "--local-reranker-backend", - help="Local reranker backend: 'vllm' (default) or 'hf'.", - rich_help_panel=_PANEL_EVAL, - ), - local_hf_batch_size: int = typer.Option( - 32, - "--local-hf-batch-size", - min=1, - help="Batch size for local HF query embedding during retrieval/reranking.", - rich_help_panel=_PANEL_EVAL, - ), - local_query_max_length: int = typer.Option( - 128, - "--local-query-max-length", - min=1, - help="Fixed token length for local HF query embeddings; longer queries are truncated.", - rich_help_panel=_PANEL_EVAL, - ), - beir_loader: Optional[str] = typer.Option(None, "--beir-loader", rich_help_panel=_PANEL_EVAL), - beir_dataset_name: Optional[str] = typer.Option(None, "--beir-dataset-name", rich_help_panel=_PANEL_EVAL), - beir_split: str = typer.Option("test", "--beir-split", rich_help_panel=_PANEL_EVAL), - beir_query_language: Optional[str] = typer.Option(None, "--beir-query-language", rich_help_panel=_PANEL_EVAL), - beir_doc_id_field: Optional[str] = typer.Option( - None, - "--beir-doc-id-field", - help="BEIR document ID field. Defaults to the known dataset setting, or pdf_basename for custom datasets.", - rich_help_panel=_PANEL_EVAL, - ), - beir_k: list[int] = typer.Option([], "--beir-k", rich_help_panel=_PANEL_EVAL), - eval_config: Optional[Path] = typer.Option( - None, - "--eval-config", - help="Path to QA sweep YAML/JSON (required when --evaluation-mode=qa; same as `retriever eval run --config`).", - path_type=Path, - dir_okay=False, - rich_help_panel=_PANEL_EVAL, - ), - retrieval_save_path: Optional[Path] = typer.Option( - None, - "--retrieval-save-path", - help="Override retrieval.save_path in the QA config (page-index / export JSON, optional).", - path_type=Path, - rich_help_panel=_PANEL_EVAL, - ), - eval_page_index: Optional[Path] = typer.Option( - None, - "--page-index", - help="Override retrieval.page_index in the QA config (optional).", - path_type=Path, - dir_okay=False, - file_okay=True, - exists=True, - rich_help_panel=_PANEL_EVAL, - ), -) -> None: - """Run the end-to-end graph ingestion pipeline against ``INPUT_PATH``.""" - - if quiet: - # Imported lazily to avoid a cycle (main.py lazy-imports this module). - from nemo_retriever.cli.shared import silence_noisy_libraries - - silence_noisy_libraries() - log_handle, original_stdout, original_stderr = _configure_logging(log_file, debug=bool(debug)) - if quiet: - # Hide INFO-level "Building graph pipeline...", "Starting ingestion...", - # etc. Errors still propagate. - logging.getLogger("nemo_retriever").setLevel(logging.WARNING) - try: - if run_mode not in {"batch", "inprocess", "service"}: - raise ValueError(f"Unsupported --run-mode: {run_mode!r}") - if run_mode == "service": - _reject_service_incompatible_flags(ctx) - if audio_split_type not in {"size", "time", "frame"}: - raise ValueError(f"Unsupported --audio-split-type: {audio_split_type!r}") - if evaluation_mode not in {"none", "audio_recall", "beir", "qa"}: - raise ValueError(f"Unsupported --evaluation-mode: {evaluation_mode!r}") - if evaluation_mode == "audio_recall": - if input_type != "audio": - raise ValueError("--evaluation-mode=audio_recall is only supported with --input-type=audio") - if recall_match_mode != "audio_segment": - raise ValueError("--evaluation-mode=audio_recall requires --recall-match-mode=audio_segment") - if evaluation_mode == "qa" and eval_config is None: - raise typer.BadParameter( - "--evaluation-mode=qa requires --eval-config (QA sweep YAML/JSON). " - "Use the same file format as `retriever eval run --config` (dataset, retrieval, models, ...)." - ) - - if run_mode == "batch": - # --quiet implies --no-ray-log-to-driver: Ray flushes worker stdout - # (HF download lines, etc.) to the driver asynchronously, often - # after ingestor.ingest() returns and the fd capture has exited. - # The env var is cached at ray import time so we also override the - # variable that ultimately reaches ray.init(log_to_driver=...). - if quiet: - ray_log_to_driver = False - os.environ["RAY_LOG_TO_DRIVER"] = "1" if ray_log_to_driver else "0" - - resolved_vdb_op = str(vdb_op or DEFAULT_VDB_OP) - resolved_vdb_kwargs = _parse_vdb_kwargs_json(vdb_kwargs_json) - if vdb_overwrite is None: - resolved_vdb_kwargs.setdefault("overwrite", True) - else: - resolved_vdb_kwargs["overwrite"] = bool(vdb_overwrite) - - _sidecar_n = sum(1 for x in (meta_dataframe, meta_source_field, meta_fields) if x is not None) - if _sidecar_n not in (0, 3): - raise typer.BadParameter( - "Sidecar metadata: pass all of --meta-dataframe, --meta-source-field, and --meta-fields, or omit all." - ) - if _sidecar_n == 3: - assert meta_dataframe is not None and meta_source_field is not None and meta_fields is not None - cols = [c.strip() for c in meta_fields.split(",") if c.strip()] - if not cols: - raise typer.BadParameter("--meta-fields must list at least one column name.") - if meta_join_key not in ("auto", "source_id", "source_name"): - raise typer.BadParameter("--meta-join-key must be one of: auto, source_id, source_name.") - resolved_vdb_kwargs = { - **resolved_vdb_kwargs, - "meta_dataframe": str(meta_dataframe.expanduser().resolve()), - "meta_source_field": meta_source_field.strip(), - "meta_fields": cols, - "meta_join_key": meta_join_key, - } - - remote_api_key = resolve_remote_api_key(api_key) - extract_remote_api_key = remote_api_key - embed_remote_api_key = remote_api_key - caption_remote_api_key = remote_api_key - reranker_bearer = ( - resolve_remote_api_key(reranker_api_key) if reranker_api_key is not None else remote_api_key - ) or "" - - if ( - any( - ( - page_elements_invoke_url, - ocr_invoke_url, - table_structure_invoke_url, - embed_invoke_url, - ) - ) - and remote_api_key is None - ): - logger.warning("Remote endpoint URL(s) were configured without an API key.") - if reranker_invoke_url and not reranker_bearer.strip(): - logger.warning( - "Reranker invoke URL is set but no bearer token was resolved; " - "set --reranker-api-key or --api-key / NVIDIA_API_KEY." - ) - - # Zero out GPU fractions when a remote URL replaces the local model. - if page_elements_invoke_url and float(page_elements_gpus_per_actor or 0.0) != 0.0: - logger.warning("Forcing page-elements GPUs to 0.0 because --page-elements-invoke-url is set.") - page_elements_gpus_per_actor = 0.0 - if ocr_invoke_url and float(ocr_gpus_per_actor or 0.0) != 0.0: - logger.warning("Forcing OCR GPUs to 0.0 because --ocr-invoke-url is set.") - ocr_gpus_per_actor = 0.0 - if embed_invoke_url and float(embed_gpus_per_actor or 0.0) != 0.0: - logger.warning("Forcing embed GPUs to 0.0 because --embed-invoke-url is set.") - embed_gpus_per_actor = 0.0 - - text_chunk_params = TextChunkParams( - max_tokens=text_chunk_max_tokens or 1024, - overlap_tokens=text_chunk_overlap_tokens if text_chunk_overlap_tokens is not None else 150, - ) - - enable_text_chunk = text_chunk or text_chunk_max_tokens is not None or text_chunk_overlap_tokens is not None - enable_caption = caption or caption_invoke_url is not None - enable_dedup = dedup if dedup is not None else enable_caption - - # In-graph VDB upload is enabled by default; opt out with --no-vdb. - enable_in_graph_vdb_upload = run_mode != "service" and not no_vdb - pipeline_vdb_upload: Optional[VdbUploadParams] = None - if enable_in_graph_vdb_upload: - pipeline_vdb_upload = VdbUploadParams(vdb_op=resolved_vdb_op, vdb_kwargs=resolved_vdb_kwargs) - - logger.info("Building graph pipeline (run_mode=%s) for %s ...", run_mode, input_path) - service_request = None - ingest_plan = None - local_execute_kwargs: dict[str, Any] = {} - if run_mode == "service": - file_patterns = _resolve_file_patterns(Path(input_path), input_type) - extract_params = _build_extract_params( - method=method, - dpi=dpi, - extract_text=extract_text, - extract_tables=extract_tables, - extract_charts=extract_charts, - extract_infographics=extract_infographics, - extract_page_as_image=extract_page_as_image, - use_page_elements=use_page_elements, - use_table_structure=use_table_structure, - table_output_format=table_output_format, - extract_remote_api_key=extract_remote_api_key, - page_elements_invoke_url=page_elements_invoke_url, - ocr_invoke_url=ocr_invoke_url, - ocr_version=ocr_version, - ocr_lang=ocr_lang, - table_structure_invoke_url=table_structure_invoke_url, - pdf_split_batch_size=pdf_split_batch_size, - pdf_extract_batch_size=pdf_extract_batch_size, - pdf_extract_tasks=pdf_extract_tasks, - pdf_extract_cpus_per_task=pdf_extract_cpus_per_task, - page_elements_actors=page_elements_actors, - page_elements_batch_size=page_elements_batch_size, - page_elements_cpus_per_actor=page_elements_cpus_per_actor, - page_elements_gpus_per_actor=page_elements_gpus_per_actor, - ocr_actors=ocr_actors, - ocr_batch_size=ocr_batch_size, - ocr_cpus_per_actor=ocr_cpus_per_actor, - ocr_gpus_per_actor=ocr_gpus_per_actor, - nemotron_parse_actors=nemotron_parse_actors, - nemotron_parse_batch_size=nemotron_parse_batch_size, - nemotron_parse_gpus_per_actor=nemotron_parse_gpus_per_actor, - ) - embed_params = _build_embed_params( - embed_model_name=embed_model_name, - embed_invoke_url=embed_invoke_url, - embed_remote_api_key=embed_remote_api_key, - embed_modality=embed_modality, - text_elements_modality=text_elements_modality, - structured_elements_modality=structured_elements_modality, - embed_granularity=embed_granularity, - embed_actors=embed_actors, - embed_batch_size=embed_batch_size, - embed_cpus_per_actor=embed_cpus_per_actor, - embed_gpus_per_actor=embed_gpus_per_actor, - local_ingest_embed_backend=local_ingest_embed_backend, - ) - service_request = ingest_service.ServiceIngestRequest( - documents=file_patterns, - input_type=input_type, - extract_params=extract_params, - embed_params=embed_params, - text_chunk_params=text_chunk_params, - enable_text_chunk=enable_text_chunk, - dedup_params=DedupParams(iou_threshold=dedup_iou_threshold) if enable_dedup else None, - caption_params=( - CaptionParams( - context_text_max_chars=caption_context_text_max_chars, - temperature=caption_temperature, - top_p=caption_top_p, - max_tokens=caption_max_tokens, - ) - if enable_caption - else None - ), - store_params=StoreParams(storage_uri=store_images_uri) if store_images_uri is not None else None, - connection=ingest_service.ServiceIngestConnectionOptions( - service_url=service_url, - service_concurrency=service_concurrency, - service_api_token=service_api_token, - ), - ) - else: - if store_actors and store_images_uri is None: - logger.warning("Ignoring --store-actors because --store-images-uri was not provided.") - - plan_lancedb_uri = str( - resolved_vdb_kwargs.get("uri") or resolved_vdb_kwargs.get("lancedb_uri") or "lancedb" - ) - plan_table_name = str( - resolved_vdb_kwargs.get("table_name") or resolved_vdb_kwargs.get("lancedb_table") or "nv-ingest" - ) - ingest_plan = resolve_ingest_plan( - IngestPlanRequest( - source=IngestSourceOptions( - documents=[str(input_path)], - input_type=input_type, - ), - runtime=IngestRuntimeOptions( - run_mode=run_mode, - ray_address=ray_address, - ray_log_to_driver=ray_log_to_driver, - ), - extract=IngestExtractOptions( - method=method, - dpi=dpi, - extract_text=extract_text, - extract_tables=extract_tables, - extract_charts=extract_charts, - extract_infographics=extract_infographics, - extract_page_as_image=extract_page_as_image, - use_page_elements=use_page_elements, - use_table_structure=use_table_structure, - page_elements_invoke_url=page_elements_invoke_url, - ocr_invoke_url=ocr_invoke_url, - ocr_version=ocr_version, - ocr_lang=ocr_lang, - table_structure_invoke_url=table_structure_invoke_url, - table_output_format=table_output_format, - extract_api_key=extract_remote_api_key, - batch=IngestExtractBatchOptions( - pdf_split_batch_size=pdf_split_batch_size, - pdf_extract_workers=pdf_extract_tasks or None, - pdf_extract_batch_size=pdf_extract_batch_size or None, - pdf_extract_cpus_per_task=pdf_extract_cpus_per_task or None, - page_elements_workers=page_elements_actors or None, - page_elements_batch_size=page_elements_batch_size or None, - page_elements_cpus_per_actor=page_elements_cpus_per_actor or None, - page_elements_gpus_per_actor=page_elements_gpus_per_actor, - ocr_workers=ocr_actors or None, - ocr_batch_size=ocr_batch_size or None, - ocr_cpus_per_actor=ocr_cpus_per_actor or None, - ocr_gpus_per_actor=ocr_gpus_per_actor, - nemotron_parse_workers=nemotron_parse_actors or None, - nemotron_parse_batch_size=nemotron_parse_batch_size or None, - nemotron_parse_gpus_per_actor=nemotron_parse_gpus_per_actor, - ), - ), - media=IngestMediaOptions( - segment_audio=segment_audio, - audio_split_type=audio_split_type, - audio_split_interval=audio_split_interval, - video_extract_audio=video_extract_audio, - video_extract_frames=video_extract_frames, - video_frame_fps=video_frame_fps, - video_frame_dedup=video_frame_dedup, - video_frame_text_dedup=video_frame_text_dedup, - video_frame_text_dedup_max_dropped_frames=video_frame_text_dedup_max_dropped_frames, - video_av_fuse=video_av_fuse, - ), - caption=IngestCaptionOptions( - enabled=enable_caption, - caption_invoke_url=caption_invoke_url if enable_caption else None, - caption_api_key=caption_remote_api_key if enable_caption else None, - caption_model_name=caption_model_name if enable_caption else None, - caption_device=caption_device if enable_caption else None, - caption_context_text_max_chars=caption_context_text_max_chars if enable_caption else None, - caption_gpu_memory_utilization=(caption_gpu_memory_utilization if enable_caption else None), - caption_temperature=caption_temperature if enable_caption else None, - caption_top_p=caption_top_p if enable_caption else None, - caption_max_tokens=caption_max_tokens if enable_caption else None, - ), - dedup=IngestDedupOptions( - enabled=enable_dedup, - iou_threshold=dedup_iou_threshold if enable_dedup else None, - ), - chunk=IngestChunkOptions( - enabled=enable_text_chunk, - text_chunk_max_tokens=text_chunk_params.max_tokens if enable_text_chunk else None, - text_chunk_overlap_tokens=text_chunk_params.overlap_tokens if enable_text_chunk else None, - ), - embed=IngestEmbedOptions( - embed_invoke_url=embed_invoke_url, - embed_model_name=embed_model_name, - embed_api_key=embed_remote_api_key, - local_ingest_embed_backend=local_ingest_embed_backend, - embed_modality=embed_modality, - text_elements_modality=text_elements_modality, - structured_elements_modality=structured_elements_modality, - embed_granularity=embed_granularity, - batch=IngestEmbedBatchOptions( - embed_workers=embed_actors or None, - embed_batch_size=embed_batch_size or None, - embed_cpus_per_actor=embed_cpus_per_actor or None, - embed_gpus_per_actor=embed_gpus_per_actor, - ), - ), - image_store=IngestImageStoreOptions( - images_uri=store_images_uri, - workers=store_actors, - ), - storage=IngestStorageOptions( - lancedb_uri=plan_lancedb_uri, - table_name=plan_table_name, - overwrite=bool(resolved_vdb_kwargs.get("overwrite", True)), - ), - ) - ) - execution_create_kwargs = dict(ingest_plan.create_kwargs) - if caption_gpus_per_actor is not None: - execution_create_kwargs["node_overrides"] = {"CaptionActor": {"num_gpus": caption_gpus_per_actor}} - ingest_plan = replace( - ingest_plan, - create_kwargs=execution_create_kwargs, - vdb_params=pipeline_vdb_upload, - ) - local_execute_kwargs = { - "verify_rows": False, - "raise_on_empty": False, - } - - # --- Execute --------------------------------------------------- - logger.info("Starting ingestion of %s ...", input_path) - ingest_start = time.perf_counter() - - def _run_ingest() -> Any: - if run_mode == "service": - if service_request is None: - raise RuntimeError("service_request must be resolved before execution in service mode") - return ingest_service.execute_service_ingest_request(service_request).result - if ingest_plan is None: - raise RuntimeError("ingest_plan must be resolved before execution in non-service mode") - return execute_ingest_plan(ingest_plan, **local_execute_kwargs).result - - if quiet: - from nemo_retriever.cli.shared import quiet_capture - - with quiet_capture(): - raw_result = _run_ingest() - else: - raw_result = _run_ingest() - ingestion_only_total_time = time.perf_counter() - ingest_start - ingest_local_results, result_df, ray_download_time, num_rows = _collect_results(run_mode, raw_result) - - if run_mode == "service": - # The service writes embeddings to LanceDB server-side during - # processing (via LanceDBWriteOperator); embedding vectors are - # stripped from SSE results to keep payloads small. Client-side - # VDB upload is therefore skipped. - logger.info( - "Service-mode ingestion complete (%d results from %d input(s), %.1fs). " - "VDB writes are handled server-side.", - len(ingest_local_results), - num_rows, - ingestion_only_total_time, - ) - uploadable_vdb_records = len(ingest_local_results) - vdb_upload_time = 0.0 - else: - uploadable_vdb_records = _count_uploadable_vdb_records(ingest_local_results) - vdb_upload_time = 0.0 - if uploadable_vdb_records == 0: - logger.warning( - "No uploadable VDB records produced; skipping %s evaluation.", - evaluation_mode, - ) - elif enable_in_graph_vdb_upload: - logger.info( - "Prepared %s uploadable VDB records (%s graph rows) for in-graph upload to %s " - "(row conversion count, not backend-confirmed writes; see VDB/operator logs for persistence).", - uploadable_vdb_records, - len(ingest_local_results), - resolved_vdb_op, - ) - - if save_intermediate is not None: - out_dir = Path(save_intermediate).expanduser().resolve() - out_dir.mkdir(parents=True, exist_ok=True) - out_path = out_dir / "extraction.parquet" - result_df.to_parquet(out_path, index=False) - logger.info("Wrote extraction Parquet for intermediate use: %s", out_path) - - if detection_summary_file is not None: - from nemo_retriever.common.detection_summary import ( - collect_detection_summary_from_df, - write_detection_summary, - ) - - write_detection_summary( - Path(detection_summary_file), - collect_detection_summary_from_df(result_df), - ) - - if uploadable_vdb_records == 0 and run_mode != "service": - if run_mode == "batch": - import ray - - ray.shutdown() - typer.echo(f"Pipeline complete: 0 uploadable records from {input_path}.") - return - - if evaluation_mode == "qa": - from nemo_retriever.tools.evaluation.cli import run_qa_sweep_from_config_dict - from nemo_retriever.tools.evaluation.config import load_eval_config - - assert eval_config is not None - cfg = load_eval_config(str(eval_config)) - r = cfg.setdefault("retrieval", {}) - if retrieval_save_path is not None: - r["save_path"] = str(Path(retrieval_save_path).resolve()) - if eval_page_index is not None: - r["page_index"] = str(Path(eval_page_index).resolve()) - - qa_t0 = time.perf_counter() - qa_code = run_qa_sweep_from_config_dict(cfg) - evaluation_total_time = time.perf_counter() - qa_t0 - total_time = time.perf_counter() - ingest_start - - _write_runtime_summary( - runtime_metrics_dir, - runtime_metrics_prefix, - { - "run_mode": run_mode, - "input_path": str(Path(input_path).resolve()), - "input_pages": int(num_rows), - "num_pages": int(num_rows), - "num_rows": int(len(result_df.index)), - "ingestion_only_secs": float(ingestion_only_total_time), - "ray_download_secs": float(ray_download_time), - "vdb_upload_secs": float(vdb_upload_time), - "evaluation_secs": float(evaluation_total_time), - "total_secs": float(total_time), - "evaluation_mode": "qa", - "evaluation_metrics": {}, - "evaluation_count": None, - "recall_details": bool(recall_details), - "vdb_op": str(resolved_vdb_op), - "qa_sweep_exit_code": qa_code, - }, - ) - if run_mode == "batch": - import ray - - ray.shutdown() - - from nemo_retriever.common.detection_summary import print_run_summary - - print_run_summary( - processed_pages=num_rows, - input_path=Path(input_path), - vdb_op=str(resolved_vdb_op), - vdb_kwargs=resolved_vdb_kwargs, - total_time=total_time, - ingest_only_total_time=ingestion_only_total_time, - ray_dataset_download_total_time=ray_download_time, - vdb_upload_total_time=vdb_upload_time, - evaluation_total_time=evaluation_total_time, - evaluation_metrics={}, - evaluation_label="QA", - evaluation_count=None, - ) - typer.echo( - f"Pipeline complete (QA): {num_rows} page(s) → {resolved_vdb_op} " - f"{_format_vdb_target(resolved_vdb_op, resolved_vdb_kwargs)} ({total_time:.1f}s)." - ) - if qa_code != 0: - raise typer.Exit(code=qa_code) - return - - evaluation_label, evaluation_total_time, evaluation_metrics, evaluation_query_count, ran = _run_evaluation( - evaluation_mode=evaluation_mode, - vdb_op=resolved_vdb_op, - vdb_kwargs=resolved_vdb_kwargs, - embed_model_name=embed_model_name, - embed_invoke_url=embed_invoke_url, - embed_remote_api_key=embed_remote_api_key, - embed_modality=embed_modality, - query_csv=query_csv, - recall_match_mode=recall_match_mode, - audio_match_tolerance_secs=audio_match_tolerance_secs, - reranker=reranker, - reranker_model_name=reranker_model_name, - reranker_invoke_url=reranker_invoke_url, - reranker_api_key=reranker_bearer, - local_reranker_backend=local_reranker_backend, - local_hf_batch_size=local_hf_batch_size, - local_query_max_length=local_query_max_length, - beir_loader=beir_loader, - beir_dataset_name=beir_dataset_name, - beir_split=beir_split, - beir_query_language=beir_query_language, - beir_doc_id_field=beir_doc_id_field, - beir_k=beir_k, - local_query_embed_backend=local_query_embed_backend, - run_mode=run_mode, - service_url=service_url, - service_api_token=service_api_token, - ) - - if not ran: - no_eval_total_time = time.perf_counter() - ingest_start - _write_runtime_summary( - runtime_metrics_dir, - runtime_metrics_prefix, - { - "run_mode": run_mode, - "input_path": str(Path(input_path).resolve()), - "input_pages": int(num_rows), - "num_pages": int(num_rows), - "num_rows": int(len(result_df.index)), - "ingestion_only_secs": float(ingestion_only_total_time), - "ray_download_secs": float(ray_download_time), - "vdb_upload_secs": float(vdb_upload_time), - "evaluation_secs": 0.0, - "total_secs": float(no_eval_total_time), - "evaluation_mode": evaluation_mode, - "evaluation_metrics": {}, - "recall_details": bool(recall_details), - "vdb_op": str(resolved_vdb_op), - }, - ) - if run_mode == "batch": - import ray - - ray.shutdown() - typer.echo( - f"Pipeline complete: {num_rows} page(s) → {resolved_vdb_op} " - f"{_format_vdb_target(resolved_vdb_op, resolved_vdb_kwargs)} ({no_eval_total_time:.1f}s)." - ) - return - - total_time = time.perf_counter() - ingest_start - - _write_runtime_summary( - runtime_metrics_dir, - runtime_metrics_prefix, - { - "run_mode": run_mode, - "input_path": str(Path(input_path).resolve()), - "input_pages": int(num_rows), - "num_pages": int(num_rows), - "num_rows": int(len(result_df.index)), - "ingestion_only_secs": float(ingestion_only_total_time), - "ray_download_secs": float(ray_download_time), - "vdb_upload_secs": float(vdb_upload_time), - "evaluation_secs": float(evaluation_total_time), - "total_secs": float(total_time), - "evaluation_mode": evaluation_mode, - "evaluation_metrics": dict(evaluation_metrics), - "evaluation_count": evaluation_query_count, - "recall_details": bool(recall_details), - "vdb_op": str(resolved_vdb_op), - }, - ) - - if run_mode == "batch": - import ray - - ray.shutdown() - - from nemo_retriever.common.detection_summary import print_run_summary - - print_run_summary( - processed_pages=num_rows, - input_path=Path(input_path), - vdb_op=str(resolved_vdb_op), - vdb_kwargs=resolved_vdb_kwargs, - total_time=total_time, - ingest_only_total_time=ingestion_only_total_time, - ray_dataset_download_total_time=ray_download_time, - vdb_upload_total_time=vdb_upload_time, - evaluation_total_time=evaluation_total_time, - evaluation_metrics=evaluation_metrics, - evaluation_label=evaluation_label, - evaluation_count=evaluation_query_count, - ) - - # Final one-line success report (mirrors `retriever ingest`). Important - # for --quiet where print_run_summary's output may otherwise be the - # only signal of completion. - typer.echo( - f"Pipeline complete: {num_rows} page(s) → {resolved_vdb_op} " - f"{_format_vdb_target(resolved_vdb_op, resolved_vdb_kwargs)} ({total_time:.1f}s)." - ) - finally: - os.sys.stdout = original_stdout - os.sys.stderr = original_stderr - if log_handle is not None: - log_handle.close() - - -def main() -> None: - """Entrypoint for ``python -m nemo_retriever.pipeline``.""" - app() - - -if __name__ == "__main__": - main() diff --git a/nemo_retriever/src/nemo_retriever/common/detection_summary.py b/nemo_retriever/src/nemo_retriever/common/detection_summary.py deleted file mode 100644 index 56c81165b3..0000000000 --- a/nemo_retriever/src/nemo_retriever/common/detection_summary.py +++ /dev/null @@ -1,292 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Shared detection summary logic. - -Provides a single function that accumulates per-page detection counters from -an iterable of ``(page_key, metadata_dict, row_dict)`` tuples. Both the -batch pipeline (reading from LanceDB) and inprocess pipeline (reading from -a DataFrame) can produce these tuples, allowing the summary computation to -be shared. -""" - -from __future__ import annotations - -from datetime import datetime -import json -from collections import defaultdict -from pathlib import Path -from typing import Any, Dict, Iterable, Optional, Tuple - - -def _safe_int(value: object, default: int = 0) -> int: - try: - if value is None: - return default - return int(value) - except Exception: - return default - - -def compute_detection_summary( - rows: Iterable[Tuple[Any, Dict[str, Any], Dict[str, Any]]], -) -> Dict[str, Any]: - """Compute deduped detection totals from an iterable of page data. - - Each element is ``(page_key, metadata_dict, row_dict)`` where: - - - *page_key* is a hashable value used to deduplicate exploded content rows - (e.g. ``(source_id, page_number)``). - - *metadata_dict* is the parsed JSON metadata (may contain counters from the - LanceDB metadata column or from direct DataFrame columns). - - *row_dict* is the raw row dict, used as fallback for counters stored as - top-level DataFrame columns (e.g. ``table``, ``chart`` lists). - """ - per_page: dict[Any, dict] = {} - - for page_key, meta, raw_row in rows: - entry = per_page.setdefault( - page_key, - { - "pe": 0, - "ocr_table": 0, - "ocr_chart": 0, - "ocr_infographic": 0, - "pe_by_label": defaultdict(int), - }, - ) - - pe = _safe_int(meta.get("page_elements_v3_num_detections") or raw_row.get("page_elements_v3_num_detections")) - entry["pe"] = max(entry["pe"], pe) - - for field, meta_key, col_key in [ - ("ocr_table", "ocr_table_detections", "table"), - ("ocr_chart", "ocr_chart_detections", "chart"), - ("ocr_infographic", "ocr_infographic_detections", "infographic"), - ]: - val = _safe_int(meta.get(meta_key)) - if val == 0: - col_val = raw_row.get(col_key) - if isinstance(col_val, list): - val = len(col_val) - entry[field] = max(entry[field], val) - - label_counts = meta.get("page_elements_v3_counts_by_label") or raw_row.get("page_elements_v3_counts_by_label") - if isinstance(label_counts, dict): - for label, count in label_counts.items(): - entry["pe_by_label"][str(label)] = max( - entry["pe_by_label"][str(label)], - _safe_int(count), - ) - - pe_by_label_totals: dict[str, int] = defaultdict(int) - pe_total = ocr_table_total = ocr_chart_total = ocr_infographic_total = 0 - for e in per_page.values(): - pe_total += e["pe"] - ocr_table_total += e["ocr_table"] - ocr_chart_total += e["ocr_chart"] - ocr_infographic_total += e["ocr_infographic"] - for label, count in e["pe_by_label"].items(): - pe_by_label_totals[label] += count - - return { - "pages_seen": len(per_page), - "page_elements_v3_total_detections": pe_total, - "page_elements_v3_counts_by_label": dict(sorted(pe_by_label_totals.items())), - "ocr_table_total_detections": ocr_table_total, - "ocr_chart_total_detections": ocr_chart_total, - "ocr_infographic_total_detections": ocr_infographic_total, - } - - -def iter_dataframe_rows(df): - """Yield ``(page_key, meta, row_dict)`` tuples from a pandas DataFrame.""" - for _, row in df.iterrows(): - row_dict = row.to_dict() - path = str(row_dict.get("path") or row_dict.get("source_id") or "") - page_number = _safe_int(row_dict.get("page_number", -1), default=-1) - - meta = row_dict.get("metadata") - if isinstance(meta, str): - try: - meta = json.loads(meta) - except Exception: - meta = {} - if not isinstance(meta, dict): - meta = {} - - yield (path, page_number), meta, row_dict - - -def collect_detection_summary_from_lancedb(uri: str, table_name: str) -> Optional[Dict[str, Any]]: - """Collect detection summary from a LanceDB table.""" - try: - from nemo_retriever.common.vdb.lancedb_read import iter_lancedb_rows - - return compute_detection_summary(iter_lancedb_rows(uri, table_name)) - except Exception: - return None - - -def collect_detection_summary_from_df(df) -> Dict[str, Any]: - """Collect detection summary from a pandas DataFrame.""" - return compute_detection_summary(iter_dataframe_rows(df)) - - -def print_detection_summary(summary: Optional[Dict[str, Any]]) -> None: - """Print a detection summary to stdout.""" - if summary is None: - print("Detection summary: unavailable (could not read metadata).") - return - print("\nDetection summary (deduped by source_id/page_number):") - print(f" Pages seen: {summary['pages_seen']}") - print(f" PageElements v3 total detections: {summary['page_elements_v3_total_detections']}") - print(f" OCR table detections: {summary['ocr_table_total_detections']}") - print(f" OCR chart detections: {summary['ocr_chart_total_detections']}") - print(f" OCR infographic detections: {summary['ocr_infographic_total_detections']}") - print(" PageElements v3 counts by label:") - by_label = summary.get("page_elements_v3_counts_by_label") or {} - if not by_label: - print(" (none)") - else: - for label, count in by_label.items(): - print(f" {label}: {count}") - - -def write_detection_summary(path: Path, summary: Optional[Dict[str, Any]]) -> None: - """Write a detection summary dict to a JSON file.""" - target = Path(path).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - payload = summary if summary is not None else {"error": "Detection summary unavailable."} - target.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") - - -def print_pages_per_second(processed_pages: Optional[int], ingest_elapsed_s: float) -> None: - """Print pages-per-second throughput to stdout.""" - if ingest_elapsed_s <= 0: - print("Pages/sec: unavailable (ingest elapsed time was non-positive).") - return - if processed_pages is None: - print("Pages/sec: unavailable (could not estimate processed pages). " f"Ingest time: {ingest_elapsed_s:.2f}s") - return - - pps = processed_pages / ingest_elapsed_s - print(f"Pages processed: {processed_pages}") - print(f"Pages/sec (ingest only; excludes Ray startup and recall): {pps:.2f}") - - -def _fmt_time(seconds: float) -> str: - """Format *seconds* as ``raw / H:MM:SS.mmm``.""" - ms = int(round(seconds * 1000)) - h, remainder = divmod(ms, 3_600_000) - m, remainder = divmod(remainder, 60_000) - s, millis = divmod(remainder, 1000) - return f"{seconds:.2f}s / {h}:{m:02d}:{s:02d}.{millis:03d}" - - -def _evaluation_metric_sort_key(item: tuple[str, float]) -> tuple[str, int, str]: - """Sort metrics like ndcg@1, ndcg@3, ..., recall@1, recall@3, ... .""" - key, _value = item - metric_name, sep, suffix = str(key).partition("@") - if sep: - try: - return metric_name, int(suffix), str(key) - except ValueError: - pass - return metric_name, 0, str(key) - - -def print_run_summary( - processed_pages: Optional[int], - input_path: Path, - vdb_op: str, - vdb_kwargs: Optional[Dict[str, Any]], - total_time: float, - ingest_only_total_time: float, - ray_dataset_download_total_time: float, - vdb_upload_total_time: float, - evaluation_total_time: float = 0.0, - evaluation_metrics: Optional[Dict[str, float]] = None, - recall_total_time: float = 0.0, - recall_metrics: Optional[Dict[str, float]] = None, - processed_files: Optional[int] = None, - evaluation_label: str = "Recall", - evaluation_count: Optional[int] = None, -) -> Dict[str, Any]: - """Print a human-readable run summary and return all metrics as a dict. - - The returned dict is the authoritative structured representation of every - metric collected during the run. Callers should persist it to a JSON file - so that the harness can read it directly instead of parsing stdout. - """ - if recall_metrics is None: - recall_metrics = {} - if evaluation_metrics is None: - evaluation_metrics = {} - pages = processed_pages if processed_pages is not None else 0 - - ingest_only_pps = pages / ingest_only_total_time if ingest_only_total_time > 0 else 0 - ingest_write_denom = ingest_only_total_time + vdb_upload_total_time - ingest_and_vdb_upload_pps = pages / ingest_write_denom if ingest_write_denom > 0 else 0 - recall_qps = pages / recall_total_time if recall_total_time > 0 else 0 - total_pps = pages / total_time if total_time > 0 else 0 - vdb_kwargs = dict(vdb_kwargs or {}) - - print(f"===== Run Summary - {datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')} UTC =====") - - print("Run Configuration:") - print(f"\tInput path: {input_path}") - print(f"\tVDB op: {vdb_op}") - if vdb_kwargs: - print(f"\tVDB kwargs: {json.dumps(vdb_kwargs, default=str, sort_keys=True)}") - - print("Runtimes:") - if processed_files is not None: - print(f"\tTotal files processed: {processed_files}") - print(f"\tTotal pages processed: {pages} from {input_path}") - print(f"\tIngestion only time: {_fmt_time(ingest_only_total_time)}") - print(f"\tRay dataset download time: {_fmt_time(ray_dataset_download_total_time)}") - print(f"\tVDB upload time: {_fmt_time(vdb_upload_total_time)}") - if recall_total_time > 0: - print(f"\tRecall time: {_fmt_time(recall_total_time)}") - if evaluation_total_time > 0: - print(f"\t{evaluation_label} time: {_fmt_time(evaluation_total_time)}") - - print("PPS:") - print(f"\tIngestion only PPS: {ingest_only_pps:.2f}") - print(f"\tIngestion + VDB upload PPS: {ingest_and_vdb_upload_pps:.2f}") - if recall_total_time > 0: - print(f"\tRecall QPS: {recall_qps:.2f}") - print(f"\tTotal - Processed: {pages} pages in {_fmt_time(total_time)} @ {total_pps:.2f} PPS") - - if recall_metrics: - print("Recall metrics:") - for k, v in sorted(recall_metrics.items(), key=_evaluation_metric_sort_key): - print(f" {k}: {v:.4f}") - elif not evaluation_metrics: - print("Recall metrics: skipped (no query CSV configured)") - - if evaluation_metrics: - print(f"{evaluation_label} metrics:") - for k, v in sorted(evaluation_metrics.items(), key=_evaluation_metric_sort_key): - print(f" {k}: {v:.4f}") - - return { - "pages": pages, - "files": processed_files, - "ingest_secs": round(ingest_only_total_time, 4), - "pages_per_sec_ingest": round(ingest_only_pps, 4), - "total_time_secs": round(total_time, 4), - "total_pps": round(total_pps, 4), - "ray_dataset_download_secs": round(ray_dataset_download_total_time, 4), - "vdb_op": str(vdb_op), - "vdb_kwargs": vdb_kwargs, - "vdb_upload_secs": round(vdb_upload_total_time, 4), - "recall_time_secs": round(recall_total_time, 4), - "evaluation_time_secs": round(evaluation_total_time, 4), - "evaluation_label": evaluation_label, - "evaluation_count": evaluation_count, - "recall_metrics": recall_metrics, - "evaluation_metrics": evaluation_metrics, - } diff --git a/nemo_retriever/src/nemo_retriever/common/parquet_to_lancedb.py b/nemo_retriever/src/nemo_retriever/common/parquet_to_lancedb.py deleted file mode 100644 index c95678248c..0000000000 --- a/nemo_retriever/src/nemo_retriever/common/parquet_to_lancedb.py +++ /dev/null @@ -1,69 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Reload extraction Parquet into LanceDB (recovery helper). - -Use this after ingestion succeeded and ``--save-intermediate`` wrote -``extraction.parquet``, but LanceDB write or index creation failed -(e.g. disk full on ``/tmp``). Set ``TMPDIR`` to a large filesystem -before calling:: - - export TMPDIR=/raid/$USER/tmp - mkdir -p "$TMPDIR" -""" - -from __future__ import annotations - -import logging -from pathlib import Path - -from nemo_retriever.common.io.dataframe import read_extraction_parquet -from nemo_retriever.common.vdb.lancedb_bulk import handle_lancedb - -logger = logging.getLogger(__name__) - -LANCEDB_URI = "lancedb" -LANCEDB_TABLE = "nv-ingest" - - -def reload_parquet_to_lancedb( - parquet_path: str | Path, - lancedb_uri: str = LANCEDB_URI, - table_name: str = LANCEDB_TABLE, - hybrid: bool = False, -) -> int: - """Read an extraction Parquet and write its records into LanceDB. - - Parameters - ---------- - parquet_path : str or Path - Path to ``extraction.parquet`` (from ``graph_pipeline --save-intermediate``). - lancedb_uri : str - LanceDB directory path. - table_name : str - Target table name inside LanceDB. - hybrid : bool - Enable hybrid (dense + sparse) indexing. - - Returns - ------- - int - Number of rows written. - """ - uri = str(Path(lancedb_uri).expanduser().resolve()) - parquet_path = Path(parquet_path).expanduser().resolve() - - logger.info("Reading %s ...", parquet_path) - df = read_extraction_parquet(parquet_path) - records = df.to_dict("records") - logger.info( - "Loaded %s rows; writing LanceDB at uri=%s table=%s ...", - len(records), - uri, - table_name, - ) - - handle_lancedb(records, uri, table_name, hybrid=hybrid, mode="overwrite") - logger.info("Done.") - return len(records) diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index e1052528bf..8ae7e624ad 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/nemo_retriever/src/nemo_retriever/common/vdb/README.md @@ -7,7 +7,7 @@ This package wraps **vector database backends** behind a small `VDB` interface ( The only built-in backend key today is **`lancedb`**, resolved by `get_vdb_op_cls()` in `factory.py` to the concrete **`LanceDB`** class in `lancedb.py`. -The root CLI is intentionally LanceDB-first: `retriever ingest ...` writes LanceDB tables, and `retriever query ...` queries LanceDB tables. Other VDB backends should plug in through the SDK/operator layer by implementing `VDB` and registering a backend key in `factory.py`; the root CLI does not currently expose a generic `--vdb-op` / `--vdb-kwargs-json` query surface. +The root CLI is intentionally LanceDB-first: `retriever ingest ...` writes LanceDB tables, and `retriever query ...` queries LanceDB tables. Other VDB backends should plug in through the SDK/operator layer by implementing `VDB` and registering a backend key in `factory.py`; the root CLI does not expose a backend-agnostic VDB configuration surface. --- @@ -39,7 +39,6 @@ Together, repartition + full batch mean **`process()`** receives **every row at ### Wiring ingestion today - **Root CLI** (`retriever ingest ...`): writes to LanceDB through the shared graph ingest path. -- **Compatibility CLI** (`retriever pipeline run ...`): builds `VdbUploadParams` and `GraphIngestor.vdb_upload(...)`, which appends `IngestVdbOperator` to the graph after embed/store and before webhook. - **Direct API**: ```python @@ -56,11 +55,12 @@ op = IngestVdbOperator( op(pandas_dataframe_of_embedded_rows) # or list of row dicts ``` -CLI-equivalent kwargs are often passed as JSON: +The root CLI exposes the built-in LanceDB target directly: ```bash -retriever pipeline run /data/pdfs --vdb-op lancedb \ - --vdb-kwargs-json '{"uri":"./kb","table_name":"nemo-retriever"}' +retriever ingest /data/pdfs \ + --lancedb-uri ./kb \ + --table-name nemo-retriever ``` --- diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py index c1ee30082f..3ab48161d1 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py @@ -32,8 +32,7 @@ def _normalize_on_bad_vectors(value: str) -> str: LanceDB's ``Table.create`` accepts a fixed set of policies for handling rows whose vector column does not match the declared fixed-size schema. We - surface the same vocabulary on this wrapper so callers can configure the - behavior through ``--vdb-kwargs-json``. + surface the same vocabulary on this wrapper for direct SDK configuration. Args: value: User-supplied policy name. Whitespace and case are ignored. @@ -156,6 +155,12 @@ def _table_schema(table: Any) -> pa.Schema: return schema() if callable(schema) else schema +def lancedb_row_count(uri: str, table_name: str) -> int: + """Return the number of rows in a LanceDB table.""" + table = lancedb.connect(uri).open_table(table_name) + return int(table.count_rows()) + + def _validate_append_schema(table: Any, expected_schema: pa.Schema, *, table_name: str, uri: str) -> None: """Fail before append when an existing table cannot accept this writer's rows.""" existing_schema = _table_schema(table) diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_read.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_read.py deleted file mode 100644 index 194f7db3c7..0000000000 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_read.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Read-only LanceDB helpers (table scans, counts) kept under ``vdb``.""" - -from __future__ import annotations - -import json -from typing import Any, Dict, Iterable, Optional, Tuple - - -def _safe_int(value: object, default: int = 0) -> int: - try: - if value is None: - return default - return int(value) - except Exception: - return default - - -def iter_lancedb_rows(uri: str, table_name: str) -> Iterable[Tuple[Any, Dict[str, Any], Dict[str, Any]]]: - """Yield ``(page_key, meta, row_dict)`` tuples from a LanceDB table.""" - import lancedb # type: ignore - - db = lancedb.connect(uri) - table = db.open_table(table_name) - df = table.to_pandas()[["source_id", "page_number", "metadata"]] - - for row in df.itertuples(index=False): - source_id = str(getattr(row, "source_id", "") or "") - page_number = _safe_int(getattr(row, "page_number", -1), default=-1) - raw_metadata = getattr(row, "metadata", None) - meta: dict = {} - if isinstance(raw_metadata, str) and raw_metadata.strip(): - try: - parsed = json.loads(raw_metadata) - if isinstance(parsed, dict): - meta = parsed - except Exception: - pass - yield (source_id, page_number), meta, {} - - -def lancedb_row_count(uri: str, table_name: str) -> int: - """Return ``table.count_rows()`` or 0 on failure.""" - import lancedb # type: ignore - - db = lancedb.connect(uri) - table = db.open_table(table_name) - return int(table.count_rows()) - - -def estimate_processed_pages(uri: str, table_name: str) -> Optional[int]: - """Estimate pages processed by counting unique (source_id, page_number) pairs. - - Falls back to table row count if page-level fields are unavailable. - """ - try: - import lancedb # type: ignore - - db = lancedb.connect(uri) - table = db.open_table(table_name) - except Exception: - return None - - try: - df = table.to_pandas()[["source_id", "page_number"]] - return int(df.dropna(subset=["source_id", "page_number"]).drop_duplicates().shape[0]) - except Exception: - try: - return int(table.count_rows()) - except Exception: - return None diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/records.py b/nemo_retriever/src/nemo_retriever/common/vdb/records.py index 2b10158f7c..c72721479a 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/records.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/records.py @@ -2,7 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Record adapters for the graph-pipeline VDB upload/retrieval path.""" +"""Record adapters for graph VDB upload and retrieval.""" from __future__ import annotations @@ -164,7 +164,7 @@ def _client_record_from_graph_row(row: dict[str, Any], *, require_embedding: boo def to_client_vdb_records(rows: list[dict[str, Any]]) -> list[list[dict[str, Any]]]: - """Convert graph-pipeline rows into the nested record shape expected by client VDBs. + """Convert graph-ingest rows into the nested record shape expected by client VDBs. Dense rows require an embedding and either nonblank text or concrete image backing. When no row survives conversion, returns ``[]`` — a falsy value so @@ -187,7 +187,7 @@ def to_client_vdb_records(rows: list[dict[str, Any]]) -> list[list[dict[str, Any def to_sparse_client_vdb_records(rows: Any) -> list[list[dict[str, Any]]]: - """Convert graph-pipeline rows into text/provenance records for sparse LanceDB ingest.""" + """Convert graph-ingest rows into text/provenance records for sparse LanceDB ingest.""" if hasattr(rows, "to_pandas"): rows = rows.to_pandas() if hasattr(rows, "to_dict"): diff --git a/nemo_retriever/src/nemo_retriever/examples/common.py b/nemo_retriever/src/nemo_retriever/examples/common.py deleted file mode 100644 index ee8d29c54e..0000000000 --- a/nemo_retriever/src/nemo_retriever/examples/common.py +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Shared helpers used by multiple example pipeline scripts.""" - -from __future__ import annotations - -from typing import Optional - -from nemo_retriever.common.vdb.lancedb_read import estimate_processed_pages - -__all__ = ["estimate_processed_pages", "print_pages_per_second"] - - -def print_pages_per_second( - processed_pages: Optional[int], - ingest_elapsed_s: float, - *, - label: str = "ingest only", -) -> None: - """Print a throughput summary line.""" - if ingest_elapsed_s <= 0: - print("Pages/sec: unavailable (ingest elapsed time was non-positive).") - return - if processed_pages is None: - print(f"Pages/sec: unavailable (could not estimate processed pages). Ingest time: {ingest_elapsed_s:.2f}s") - return - - pps = processed_pages / ingest_elapsed_s - print(f"Pages processed: {processed_pages}") - print(f"Pages/sec ({label}): {pps:.2f}") diff --git a/nemo_retriever/src/nemo_retriever/examples/graph_pipeline.py b/nemo_retriever/src/nemo_retriever/examples/graph_pipeline.py deleted file mode 100644 index 48e86f58a3..0000000000 --- a/nemo_retriever/src/nemo_retriever/examples/graph_pipeline.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Backward-compat shim for the graph ingestion pipeline. - -The implementation was moved to :mod:`nemo_retriever.pipeline` and is exposed -as the ``retriever pipeline run`` CLI subcommand. - -This module re-exports the same Typer :data:`app` and keeps the -``python -m nemo_retriever.examples.graph_pipeline `` entry point -working for existing pipeline callers. - -New code should invoke the pipeline via one of the following: - -* ``retriever pipeline run [OPTIONS]`` -* ``python -m nemo_retriever.pipeline [OPTIONS]`` -* ``from nemo_retriever.pipeline import app`` (Typer app) or - ``from nemo_retriever.pipeline import run`` (command callable) - -Module-level names that used to live in the monolithic script (``GraphIngestor``, -``_resolve_file_patterns``, etc.) are re-exported for tests and harness code that -patches or imports them from this module. -""" - -from __future__ import annotations - -from nemo_retriever.operators.extract.audio.asr_actor import asr_params_from_env -from nemo_retriever.ingestor.graph_ingestor import GraphIngestor -from nemo_retriever.cli.pipeline import __main__ as _pipeline_main - -# Typer app and supporting hooks (same as :mod:`nemo_retriever.pipeline.__main__`). -app = _pipeline_main.app -_resolve_file_patterns = _pipeline_main._resolve_file_patterns - -__all__ = [ - "GraphIngestor", - "_resolve_file_patterns", - "app", - "asr_params_from_env", -] - -if __name__ == "__main__": - app() diff --git a/nemo_retriever/src/nemo_retriever/harness/portal/app.py b/nemo_retriever/src/nemo_retriever/harness/portal/app.py index 799cef88dc..92eff5d708 100644 --- a/nemo_retriever/src/nemo_retriever/harness/portal/app.py +++ b/nemo_retriever/src/nemo_retriever/harness/portal/app.py @@ -834,7 +834,7 @@ async def get_run_lancedb_info(run_id: int): if not uri: return {"available": False, "uri": None, "row_count": 0} try: - from nemo_retriever.common.vdb.lancedb_read import lancedb_row_count + from nemo_retriever.common.vdb.lancedb import lancedb_row_count count = int(lancedb_row_count(uri, LANCEDB_TABLE)) return {"available": True, "uri": uri, "row_count": count, "table": LANCEDB_TABLE} diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index 69d2ab9369..66a628f64c 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -126,7 +126,7 @@ def _generate_batches(prompts: Iterable[str], batch_size: int = 100) -> List[Lis def _text_from_row(row: pd.Series, *, text_column: str) -> Optional[str]: """ - Extract text from a row with small fallbacks for graph-pipeline inputs. + Extract text from a row with small fallbacks for graph-ingest inputs. """ v = row.get(text_column) if isinstance(v, str) and v.strip(): diff --git a/nemo_retriever/src/nemo_retriever/operators/vdb.py b/nemo_retriever/src/nemo_retriever/operators/vdb.py index 16874e8bc2..b0c8efd2d3 100644 --- a/nemo_retriever/src/nemo_retriever/operators/vdb.py +++ b/nemo_retriever/src/nemo_retriever/operators/vdb.py @@ -130,7 +130,7 @@ def preprocess(self, data: Any, **kwargs: Any) -> Any: return data def process(self, data: Any, **kwargs: Any) -> Any: - # Compatibility shim: graph_pipeline emits flat embedded rows, while + # Graph ingest emits flat embedded rows, while # nv-ingest-client VDB.run still expects nested Nemo Retriever Library (NRL) records. records = to_client_vdb_records(data) if self._sidecar_spec is not None and self._sidecar_lookup is not None: diff --git a/nemo_retriever/src/nemo_retriever/tools/evaluation/README.md b/nemo_retriever/src/nemo_retriever/tools/evaluation/README.md index 066cc8f5c4..cad9696fb4 100644 --- a/nemo_retriever/src/nemo_retriever/tools/evaluation/README.md +++ b/nemo_retriever/src/nemo_retriever/tools/evaluation/README.md @@ -1,6 +1,6 @@ # QA Evaluation Pipeline -The evaluation framework lives in **`nemo_retriever.evaluation`** (install `nemo-retriever[llm]` from PyPI, or **`uv pip install -e "./nemo_retriever[llm]"` from this repo root** so `graph_pipeline` and local changes resolve). +The evaluation framework lives in **`nemo_retriever.evaluation`** (install `nemo-retriever[llm]` from PyPI, or **`uv pip install -e "./nemo_retriever[llm]"` from this repo root** to exercise local changes). Measures LLM answer quality over a RAG pipeline: retrieve context from a VDB, generate answers with one or more LLMs, and score each answer against ground-truth references using multi-tier scoring and an LLM-as-judge. @@ -39,28 +39,27 @@ End-to-end bo767 + LanceDB + full-page markdown touches these **artifacts** and | Stage | Artifacts produced | Code / APIs involved | |-------|-------------------|----------------------| -| **1. Ingest + embed** | `lancedb//
/` (embedded sub-page chunks); optionally `data/bo767_extracted/*.parquet` | `python -m nemo_retriever.examples.graph_pipeline` (extract, embed, VDB upload to LanceDB via the operator graph). Add `--save-intermediate` to also save extraction Parquet for full-page markdown (recommended for best results). **Table name must match** `retriever eval export` (`--lancedb-table`, default `nemo-retriever`). | -| **2. Full-page markdown index** | `data/bo767_page_markdown.json` (`source_id` -> page -> markdown) | `retriever eval build-page-index` -> `nemo_retriever.io.markdown.build_page_index()` (which calls `to_markdown_by_page` per document); numpy list columns are coerced so structured content is not dropped. | +| **1. Ingest + embed** | `lancedb//
/` (embedded sub-page chunks) | `retriever ingest` extracts, embeds, and writes LanceDB through the supported ingest workflow. **Table name must match** `retriever eval export` (`--lancedb-table`, default `nemo-retriever`). | +| **2. Optional full-page markdown index** | `data/bo767_page_markdown.json` (`source_id` -> page -> markdown) | If you already have extraction Parquet, `retriever eval build-page-index` calls `nemo_retriever.io.markdown.build_page_index()` to create a page index. | | **3. Retrieval export** | `data/eval/bo767_retrieval_fullpage.json` (or sub-page JSON) | `retriever eval export` -> `nemo_retriever.export.export_retrieval_json()` queries LanceDB; if `--page-index` is provided, hits are expanded/deduped by `(source_id, page)` and replaced with full-page markdown strings. | | **4. Ground truth** | `data/bo767_annotations.csv` (repo root) | Questions/answers for export and eval; must align with **query string normalization** in `FileRetriever` (see retrieval JSON rules). | | **5. Evaluation** | `qa_results_*.json` | `retriever eval run` or operator graph chain -> `nemo_retriever.evaluation`: `RetrievalLoaderOperator >> QAGenerationOperator >> JudgingOperator >> ScoringOperator`, or `QAEvalPipeline` for multi-model sweeps. | -**Data flow (conceptual):** PDFs -> (A) **chunked embeddings in LanceDB** for similarity search; (B) **Parquet** for full-page reconstruction. **Export** runs search on (A), then **replaces** hit chunks with pages from (B) via the index. In **file mode**, eval reads the retrieval JSON + ground-truth CSV. In **[lancedb mode](#in-memory-lancedb-retrieval)**, eval queries LanceDB directly in-memory (optionally saving the JSON for later re-runs). +**Data flow (conceptual):** `retriever ingest` writes chunk embeddings to +LanceDB. **Export** queries that table and writes retrieval JSON. If a separate +workflow supplies extraction Parquet, an optional page index can replace hit +chunks with full-page markdown during export. In **file mode**, evaluation reads +the retrieval JSON and ground-truth CSV. In **[lancedb mode](#in-memory-lancedb-retrieval)**, +evaluation queries LanceDB directly (optionally saving JSON for later re-runs). ``` - NeMo Retriever (steps 1-3) Universal (steps 4-5) - ────────────────────────── ───────────────────── - Step 1 Step 2 Step 3 - Ingest + Embed Index Export QA Eval -+-----------------------------+ +--------+ +----------+ +------------------+ -| graph_pipeline | | Parquet| | LanceDB | | RetrievalLoader | -| --vdb-kwargs-json ... | | -> page|->| queries |->| >> Generation | -| [--save-intermediate ]| | md idx | | + pages | | >> Judging | -| (always: LanceDB output) | +--------+ | -> JSON | | >> Scoring | -| (optional: Parquet output) | +----------+ +------------------+ -+-----------------------------+ | | | - | | page_md.json retrieval.json qa_results.json -lancedb/ *.parquet (opt.) ++--------------------+ +---------+ +----------------+ +------------------+ +| retriever ingest | -> | LanceDB | -> | eval export | -> | RetrievalLoader | +| --lancedb-uri ... | | table | | retrieval JSON | | >> Generation | +| --table-name ... | +---------+ +----------------+ | >> Judging | ++--------------------+ ^ | >> Scoring | + | optional page index +------------------+ + extraction Parquet -> build-page-index Bring Your Own Retrieval +---------+--------+ ───────────────────────── | Any pipeline that | @@ -77,11 +76,13 @@ Steps 4-5 can be re-run with different LLM configs without repeating retrieval. ## Reproducing the bo767 Run -Exact commands to reproduce the full-page markdown QA evaluation from scratch. +The default supported path evaluates sub-page chunks from a LanceDB table. A +full-page markdown run additionally requires a separately produced extraction +Parquet dataset for `retriever eval build-page-index`. **Working directory:** All commands below run from the **repo root** unless otherwise noted. -**Debug:** Lance index build can hit `No space left on device` when `/tmp` is a tiny tmpfs; set `export TMPDIR=/path/to/large/filesystem/tmp` and `mkdir -p "$TMPDIR"` before step 1. If `extraction.parquet` was written but LanceDB failed, retry with `python -c "from nemo_retriever.utils.parquet_to_lancedb import reload_parquet_to_lancedb; reload_parquet_to_lancedb('', '')""`; otherwise re-run `graph_pipeline`. +**Debug:** Lance index build can hit `No space left on device` when `/tmp` is a tiny tmpfs; set `export TMPDIR=/path/to/large/filesystem/tmp` and `mkdir -p "$TMPDIR"` before step 1, then re-run `retriever ingest`. **Before running any quick reference below**, complete the one-time setup: @@ -100,35 +101,6 @@ export NVIDIA_API_KEY="nvapi-..." See [Python environment](#python-environment) and [Prerequisites](#prerequisites-data-and-keys) for details. -
-Quick reference -- full-page markdown (all commands) - -```bash -cd /path/to/nemo-retriever - -# 1. Ingest + embed + save Parquet in one pass (~45-90 min) -python -m nemo_retriever.examples.graph_pipeline /path/to/bo767 \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' \ - --save-intermediate data/bo767_extracted - -# 2. Build page markdown index (~5-10 min) -retriever eval build-page-index \ - --parquet-dir data/bo767_extracted \ - --output data/bo767_page_markdown.json - -# 3. Export retrieval results (~5-15 min) -retriever eval export \ - --lancedb-uri lancedb \ - --query-csv data/bo767_annotations.csv \ - --output data/eval/bo767_retrieval_fullpage.json \ - --page-index data/bo767_page_markdown.json - -# 4. Run QA evaluation (~1-2 hrs) -retriever eval run --config nemo_retriever/examples/eval_sweep.yaml -``` - -
-
Quick reference -- sub-page chunks (skip full-page markdown) @@ -141,8 +113,9 @@ retrieval context, which may produce lower scores for structured content cd /path/to/nemo-retriever # 1. Ingest + embed into LanceDB -python -m nemo_retriever.examples.graph_pipeline /path/to/bo767 \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' +retriever ingest /path/to/bo767 \ + --lancedb-uri lancedb \ + --table-name nemo-retriever # 2. Export retrieval (sub-page chunks, no page index) retriever eval export \ @@ -157,35 +130,21 @@ retriever eval run --config nemo_retriever/examples/eval_sweep.yaml
-Quick reference -- end-to-end in-memory (single command: ingest + eval) +Quick reference -- live LanceDB evaluation -Alternative to the separable export+eval path above. One command ingests -PDFs, builds the full-page markdown index in-memory, queries LanceDB, and -runs generation + judging + scoring. Optionally saves the retrieval JSON -so you can re-run eval later without re-querying. +Alternative to the separable export+eval path above. Ingest first, then let the +evaluation command query LanceDB directly. Optionally save retrieval JSON so +you can re-run evaluation later without re-querying. ```bash cd /path/to/nemo-retriever -# Single command: ingest -> page index -> LanceDB query -> QA eval -python -m nemo_retriever.examples.graph_pipeline /path/to/bo767 \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' \ - --evaluation-mode qa \ - --eval-config nemo_retriever/examples/eval_sweep.yaml \ - --query-csv data/bo767_annotations.csv \ - --retrieval-save-path data/eval/bo767_retrieval.json -``` - -The page index is built automatically from the ingestion results (no -separate step needed). Pass `--page-index ` to use a pre-built -index instead, or `--save-intermediate ` if you also want the -extraction Parquet saved for other uses. - -Or, if you already have a LanceDB from a previous ingestion and want to -run eval standalone (no re-ingestion): +retriever ingest /path/to/bo767 \ + --lancedb-uri lancedb \ + --table-name nemo-retriever -```bash export LANCEDB_URI="lancedb" +export LANCEDB_TABLE="nemo-retriever" export QA_DATASET="csv:data/bo767_annotations.csv" export RETRIEVAL_SAVE_PATH="data/eval/bo767_retrieval.json" # optional retriever eval run --from-env @@ -195,19 +154,18 @@ retriever eval run --from-env ### Which path should I use? -| | **Option A: Separable (recommended)** | **Option B: End-to-end in-memory** | +| | **Option A: Separable (recommended)** | **Option B: Live LanceDB** | |---|---|---| -| **Steps** | Ingest -> Export JSON -> Run eval | Single command: Ingest -> Page index -> LanceDB query -> Eval | +| **Steps** | Ingest -> Export JSON -> Run eval | Ingest -> LanceDB query -> Eval | | **Re-run eval with a different model?** | Instant -- just re-run step 4 with the same JSON | Must re-query LanceDB (or pass `save_path` to cache) | | **Share results with teammates?** | Send the retrieval JSON file | Must share LanceDB or use `save_path` | | **Best for** | Benchmarking, CI, reproducible comparisons | Quick end-to-end validation, development iteration | -| **Commands** | `retriever eval export` + `retriever eval run` | `--evaluation-mode qa` on `graph_pipeline` or `LANCEDB_URI` with `--from-env` | +| **Commands** | `retriever eval export` + `retriever eval run` | `LANCEDB_URI` with `retriever eval run --from-env` | Option A is the primary design -- the retrieval JSON is the interface contract that lets you re-run eval N times without re-querying, swap retrievers, and -share results. Option B is a convenience shortcut for when you want a single -end-to-end pass (e.g. validating a new ingestion pipeline). Both produce -identical evaluation results. +share results. Option B is a convenience path for direct evaluation against an +existing LanceDB table. ### Bring your own retrieval (skip steps 1-3) @@ -230,7 +188,7 @@ the [interface contract](#retrieval-json-format-interface-contract) below. Steps 1-3 (ingest, build index, export) require the **`nemo_retriever`** library with LanceDB, CUDA, and Ray support. Step 4 (QA eval) additionally requires **`litellm`**. -**Recommended setup:** create an isolated Python 3.12 virtual environment and install the **local** `nemo_retriever` checkout in editable mode (required when working in this repo so `python -m nemo_retriever.examples.graph_pipeline` resolves): +**Recommended setup:** create an isolated Python 3.12 virtual environment and install the **local** `nemo_retriever` checkout in editable mode: ```bash uv venv qa-retriever --python 3.12 @@ -255,46 +213,29 @@ ls /path/to/bo767/*.pdf | wc -l # should be 767 ### Step 1: Ingest and embed PDFs (NeMo Retriever) -`graph_pipeline.py` builds an operator graph (`AbstractOperator` nodes connected via `>>`) and -executes it through either an `InprocessExecutor` (inprocess mode, default) or `RayDataExecutor` -(batch mode). The pipeline extracts, embeds, and uploads chunks to LanceDB in a -single pass. +`retriever ingest` resolves the manifest and executes the same graph ingest +implementation in local mode by default or with Ray Data in explicit `batch` +mode. It extracts, embeds, and uploads chunks to LanceDB in one workflow. Run from the **repo root**. **Estimated time: ~45-90 min** (767 PDFs, GPU-accelerated extraction + embedding). -**Recommended (full-page markdown):** pass `--save-intermediate ` to also write the -extraction DataFrame as Parquet. That preserves table/chart/infographic columns for step 2 -to reconstruct full pages and generally yields better results on structured content. - ```bash -python -m nemo_retriever.examples.graph_pipeline /path/to/bo767 \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' \ - --save-intermediate data/bo767_extracted -``` - -Output: -- `lancedb/nemo-retriever/` (~84k chunks) -- used by step 3 for retrieval queries. -- `data/bo767_extracted/*.parquet` -- used by step 2 for full-page markdown. - -**Minimal (skip Parquet / full-page path):** omit `--save-intermediate` if you only need -LanceDB and will skip step 2. Step 3 will use raw sub-page chunks instead of full-page -markdown. - -```bash -python -m nemo_retriever.examples.graph_pipeline /path/to/bo767 \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' +retriever ingest /path/to/bo767 \ + --lancedb-uri lancedb \ + --table-name nemo-retriever ``` Output: - `lancedb/nemo-retriever/` (~84k chunks) -- used by step 3 for retrieval queries. -**Note:** `graph_pipeline.py` defaults to `--run-mode inprocess`. Pass `--run-mode batch` for Ray Data scale-out. Both modes return a `pandas.DataFrame` and produce compatible chunk records for downstream evaluation. +Use `retriever ingest batch ...` for Ray Data scale-out. ### Step 2: Build page markdown index (NeMo Retriever) -> **Requires** `--save-intermediate` from step 1. If you skipped it, you can -> skip this step too -- step 3 will fall back to raw sub-page chunks. +> **Optional:** this step requires an extraction Parquet dataset produced by a +> separate workflow. Skip it to evaluate the sub-page chunks stored by +> `retriever ingest` directly. Steps 2-4 are run from the repo root (or adjust paths accordingly). @@ -313,8 +254,8 @@ Output: `data/bo767_page_markdown.json` (~180 MB, ~6k pages across 767 docs). ### Step 3: Export retrieval results (NeMo Retriever) -> **Skip this step** if using the [end-to-end in-memory path](#which-path-should-i-use) -> (`--evaluation-mode qa`), which queries LanceDB live instead of reading a JSON file. +> **Skip this step** if using the [live LanceDB path](#which-path-should-i-use), +> which queries LanceDB during evaluation instead of reading a JSON file. Queries LanceDB for each ground-truth question via `nemo_retriever.export.export_retrieval_json()`, then looks up the full-page markdown for each hit's page. Multiple sub-page hits @@ -342,15 +283,13 @@ retriever eval export \ Output: `data/eval/bo767_retrieval_fullpage.json` (~50 MB, 1005 queries). **Sub-page chunk mode:** omit `--page-index` to skip full-page -expansion. The export will use raw sub-page chunks directly from LanceDB, which requires -only step 1 without `--save-intermediate` (no Parquet or page index needed). +expansion. The export will use raw sub-page chunks directly from LanceDB, so no +Parquet or page index is needed. ### Step 4: Run QA evaluation -> **Alternative:** skip steps 3 and 4 and use `--evaluation-mode qa` on -> `graph_pipeline` to query LanceDB in-memory and run eval in one pass. -> See [Which path should I use?](#which-path-should-i-use) and the -> end-to-end quick reference above. +> **Alternative:** use the [live LanceDB path](#which-path-should-i-use) to +> query the table during `retriever eval run --from-env`. **Estimated time: ~15 min - 45 min** (1005 queries, ~12s per query for generation + judge, 8 concurrent workers). @@ -453,9 +392,8 @@ plus `generation_miss` (model had the context but answered incorrectly). ### Sub-page chunk mode -To skip full-page markdown and use raw sub-page chunks instead, omit -`--save-intermediate` from step 1, skip step 2, and omit -`--page-index` in step 3. This produces smaller context windows and may +To use raw sub-page chunks, skip step 2 and omit `--page-index` in step 3. +This produces smaller context windows and may result in lower scores for queries that span structured content (tables, charts, infographics). See the sub-page quick reference above. @@ -679,7 +617,7 @@ independently of any ingestion or retrieval pipeline. | Entry Point | Use Case | |-------------|----------| -| `python -m nemo_retriever.examples.graph_pipeline` | Ingest PDFs into LanceDB (extract + embed + VDB upload via operator graph). | +| `retriever ingest` | Ingest supported documents and media into LanceDB. | | `retriever eval build-page-index` | Build full-page markdown index from Parquet extraction results | | `retriever eval export` | Export retrieval from NeMo Retriever LanceDB (supports full-page markdown) | | `retriever eval run` | QA eval runner -- env-var (`--from-env`) or config-driven (`--config`), uses `QAEvalPipeline` | @@ -846,22 +784,21 @@ export QA_DATASET="csv:data/bo767_annotations.csv" retriever eval run --from-env ``` -**Graph pipeline (single command: ingest + page index + QA eval):** +**Live LanceDB evaluation:** -The page index is built automatically from the ingestion results -- no -separate `build-page-index` step is needed. +Ingest through the supported CLI, then run evaluation against that table. ```bash -python -m nemo_retriever.examples.graph_pipeline /data/pdfs \ - --vdb-kwargs-json '{"uri":"lancedb","table_name":"nemo-retriever"}' \ - --evaluation-mode qa \ - --eval-config nemo_retriever/examples/eval_sweep.yaml \ - --query-csv data/bo767_annotations.csv \ - --retrieval-save-path data/eval/bo767_retrieval.json -``` +retriever ingest /data/pdfs \ + --lancedb-uri lancedb \ + --table-name nemo-retriever -Pass `--page-index ` to use a pre-built index instead of the -automatic in-memory build. +export LANCEDB_URI="lancedb" +export LANCEDB_TABLE="nemo-retriever" +export QA_DATASET="csv:data/bo767_annotations.csv" +export RETRIEVAL_SAVE_PATH="data/eval/bo767_retrieval.json" +retriever eval run --from-env +``` ### Environment Variables diff --git a/nemo_retriever/src/nemo_retriever/tools/evaluation/cli.py b/nemo_retriever/src/nemo_retriever/tools/evaluation/cli.py index cb4a53cc33..d39b21d6b1 100644 --- a/nemo_retriever/src/nemo_retriever/tools/evaluation/cli.py +++ b/nemo_retriever/src/nemo_retriever/tools/evaluation/cli.py @@ -228,8 +228,7 @@ def _build_env_config() -> tuple[dict, str, str, str, float]: def run_qa_sweep_from_config_dict(cfg: dict) -> int: """Run a QA evaluation sweep from a loaded config dict (0 = success, 1 = failure). - Used by ``retriever eval run`` and by ``retriever pipeline run`` when - ``--evaluation-mode=qa`` so the LLM sweep stays in one implementation. + Used by ``retriever eval run`` so the LLM sweep stays in one implementation. """ from nemo_retriever.tools.evaluation.ground_truth import get_qa_dataset_loader from nemo_retriever.tools.evaluation.retrievers import FileRetriever diff --git a/nemo_retriever/src/nemo_retriever/tools/recall/beir_eval.py b/nemo_retriever/src/nemo_retriever/tools/recall/beir_eval.py index 971a0bd826..720eae6db2 100644 --- a/nemo_retriever/src/nemo_retriever/tools/recall/beir_eval.py +++ b/nemo_retriever/src/nemo_retriever/tools/recall/beir_eval.py @@ -5,7 +5,7 @@ Reuses the existing evaluation logic from ``nemo_retriever.recall.beir`` and prints the standard run summary via -``nemo_retriever.utils.detection_summary.print_run_summary``. +``nemo_retriever.tools.recall.core.print_run_summary``. """ from __future__ import annotations @@ -32,8 +32,8 @@ class BEIREvaluatorActor: """Designer BEIR evaluation node against an existing LanceDB table. Assumes vectors were already written (for example via - :class:`~nemo_retriever.vdb.operators.IngestVdbOperator` or the ``retriever - pipeline`` upload path). After evaluation, calls ``print_run_summary`` like + :class:`~nemo_retriever.vdb.operators.IngestVdbOperator` or ``retriever + ingest``). After evaluation, calls ``print_run_summary`` like the batch pipeline. """ @@ -78,7 +78,7 @@ def evaluate(self) -> dict[str, Any]: """ from nemo_retriever.models import resolve_embed_model from nemo_retriever.tools.recall.beir import BeirConfig, evaluate_lancedb_beir - from nemo_retriever.common.detection_summary import print_run_summary + from nemo_retriever.tools.recall.core import print_run_summary resolved_model = resolve_embed_model(self.embedding_model) diff --git a/nemo_retriever/src/nemo_retriever/tools/recall/core.py b/nemo_retriever/src/nemo_retriever/tools/recall/core.py index f08d85ccd5..79173c825e 100644 --- a/nemo_retriever/src/nemo_retriever/tools/recall/core.py +++ b/nemo_retriever/src/nemo_retriever/tools/recall/core.py @@ -9,6 +9,7 @@ import logging import time from dataclasses import dataclass, field +from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL @@ -20,6 +21,117 @@ import pandas as pd +def _fmt_time(seconds: float) -> str: + """Format *seconds* as ``raw / H:MM:SS.mmm``.""" + ms = int(round(seconds * 1000)) + h, remainder = divmod(ms, 3_600_000) + m, remainder = divmod(remainder, 60_000) + s, millis = divmod(remainder, 1000) + return f"{seconds:.2f}s / {h}:{m:02d}:{s:02d}.{millis:03d}" + + +def _evaluation_metric_sort_key(item: tuple[str, float]) -> tuple[str, int, str]: + """Sort metrics like ndcg@1, ndcg@3, ..., recall@1, recall@3, ... .""" + key, _value = item + metric_name, sep, suffix = str(key).partition("@") + if sep: + try: + return metric_name, int(suffix), str(key) + except ValueError: + pass + return metric_name, 0, str(key) + + +def print_run_summary( + processed_pages: Optional[int], + input_path: Path, + vdb_op: str, + vdb_kwargs: Optional[Dict[str, Any]], + total_time: float, + ingest_only_total_time: float, + ray_dataset_download_total_time: float, + vdb_upload_total_time: float, + evaluation_total_time: float = 0.0, + evaluation_metrics: Optional[Dict[str, float]] = None, + recall_total_time: float = 0.0, + recall_metrics: Optional[Dict[str, float]] = None, + processed_files: Optional[int] = None, + evaluation_label: str = "Recall", + evaluation_count: Optional[int] = None, +) -> Dict[str, Any]: + """Print a human-readable run summary and return all metrics as a dict.""" + if recall_metrics is None: + recall_metrics = {} + if evaluation_metrics is None: + evaluation_metrics = {} + pages = processed_pages if processed_pages is not None else 0 + + ingest_only_pps = pages / ingest_only_total_time if ingest_only_total_time > 0 else 0 + ingest_write_denom = ingest_only_total_time + vdb_upload_total_time + ingest_and_vdb_upload_pps = pages / ingest_write_denom if ingest_write_denom > 0 else 0 + recall_qps = pages / recall_total_time if recall_total_time > 0 else 0 + total_pps = pages / total_time if total_time > 0 else 0 + vdb_kwargs = dict(vdb_kwargs or {}) + + print(f"===== Run Summary - {datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')} UTC =====") + + print("Run Configuration:") + print(f"\tInput path: {input_path}") + print(f"\tVDB op: {vdb_op}") + if vdb_kwargs: + print(f"\tVDB kwargs: {json.dumps(vdb_kwargs, default=str, sort_keys=True)}") + + print("Runtimes:") + if processed_files is not None: + print(f"\tTotal files processed: {processed_files}") + print(f"\tTotal pages processed: {pages} from {input_path}") + print(f"\tIngestion only time: {_fmt_time(ingest_only_total_time)}") + print(f"\tRay dataset download time: {_fmt_time(ray_dataset_download_total_time)}") + print(f"\tVDB upload time: {_fmt_time(vdb_upload_total_time)}") + if recall_total_time > 0: + print(f"\tRecall time: {_fmt_time(recall_total_time)}") + if evaluation_total_time > 0: + print(f"\t{evaluation_label} time: {_fmt_time(evaluation_total_time)}") + + print("PPS:") + print(f"\tIngestion only PPS: {ingest_only_pps:.2f}") + print(f"\tIngestion + VDB upload PPS: {ingest_and_vdb_upload_pps:.2f}") + if recall_total_time > 0: + print(f"\tRecall QPS: {recall_qps:.2f}") + print(f"\tTotal - Processed: {pages} pages in {_fmt_time(total_time)} @ {total_pps:.2f} PPS") + + if recall_metrics: + print("Recall metrics:") + for k, v in sorted(recall_metrics.items(), key=_evaluation_metric_sort_key): + print(f" {k}: {v:.4f}") + elif not evaluation_metrics: + print("Recall metrics: skipped (no query CSV configured)") + + if evaluation_metrics: + print(f"{evaluation_label} metrics:") + for k, v in sorted(evaluation_metrics.items(), key=_evaluation_metric_sort_key): + print(f" {k}: {v:.4f}") + + return { + "pages": pages, + "files": processed_files, + "ingest_secs": round(ingest_only_total_time, 4), + "pages_per_sec_ingest": round(ingest_only_pps, 4), + "total_time_secs": round(total_time, 4), + "total_pps": round(total_pps, 4), + "ray_dataset_download_secs": round(ray_dataset_download_total_time, 4), + "vdb_op": str(vdb_op), + "vdb_kwargs": vdb_kwargs, + "vdb_upload_secs": round(vdb_upload_total_time, 4), + "recall_time_secs": round(recall_total_time, 4), + "evaluation_time_secs": round(evaluation_total_time, 4), + "evaluation_label": evaluation_label, + "evaluation_count": evaluation_count, + "recall_metrics": recall_metrics, + "evaluation_metrics": evaluation_metrics, + } + + @dataclass(frozen=True) class RecallConfig: vdb_op: str = "lancedb" diff --git a/nemo_retriever/src/nemo_retriever/tools/recall/recall_eval.py b/nemo_retriever/src/nemo_retriever/tools/recall/recall_eval.py index 6511b46b0e..f64a55a8b6 100644 --- a/nemo_retriever/src/nemo_retriever/tools/recall/recall_eval.py +++ b/nemo_retriever/src/nemo_retriever/tools/recall/recall_eval.py @@ -5,7 +5,7 @@ Reuses the existing evaluation logic from ``nemo_retriever.recall.core`` and ``nemo_retriever.recall.beir``, and prints the standard run summary via -``nemo_retriever.utils.detection_summary.print_run_summary``. +``nemo_retriever.tools.recall.core.print_run_summary``. """ from __future__ import annotations @@ -36,8 +36,8 @@ class RecallEvaluatorActor: """Designer evaluation node against an existing LanceDB table. Assumes vectors were already written (for example via - :class:`~nemo_retriever.vdb.operators.IngestVdbOperator` or the ``retriever - pipeline`` upload path). Supports ``audio_recall`` (ground-truth query CSV) + :class:`~nemo_retriever.vdb.operators.IngestVdbOperator` or ``retriever + ingest``). Supports ``audio_recall`` (ground-truth query CSV) and ``beir`` (HuggingFace BEIR dataset) modes, then calls ``print_run_summary`` like the batch pipeline. """ @@ -94,7 +94,7 @@ def evaluate(self) -> dict[str, Any]: Returns the ``summary_dict`` produced by ``print_run_summary``. """ from nemo_retriever.models import resolve_embed_model - from nemo_retriever.common.detection_summary import print_run_summary + from nemo_retriever.tools.recall.core import print_run_summary resolved_model = resolve_embed_model(self.embedding_model) diff --git a/nemo_retriever/tests/test_beir_evaluation.py b/nemo_retriever/tests/test_beir_evaluation.py index e8ee88b9d0..4bfc151ea7 100644 --- a/nemo_retriever/tests/test_beir_evaluation.py +++ b/nemo_retriever/tests/test_beir_evaluation.py @@ -3,7 +3,6 @@ import pytest -from nemo_retriever.cli.pipeline import __main__ as pipeline_main from nemo_retriever.tools.recall.beir import ( BeirConfig, BeirDataset, @@ -155,114 +154,6 @@ def test_resolve_beir_dataset_options_does_not_guess_unknown_dataset() -> None: assert options.ks == DEFAULT_BEIR_KS -def test_pipeline_beir_evaluation_keeps_custom_dataset_doc_id_default(monkeypatch) -> None: - import nemo_retriever.models as model_module - import nemo_retriever.tools.recall.beir as beir_module - - captured: dict[str, BeirConfig] = {} - - def _fake_evaluate_lancedb_beir(cfg: BeirConfig): - captured["cfg"] = cfg - return ( - BeirDataset(dataset_name=cfg.dataset_name, query_ids=["q1"], queries=["query"], qrels={"q1": {"doc": 1}}), - [], - {}, - {"recall@5": 1.0}, - ) - - monkeypatch.setattr(model_module, "resolve_embed_model", lambda model_name: model_name) - monkeypatch.setattr(beir_module, "evaluate_lancedb_beir", _fake_evaluate_lancedb_beir) - - label, _elapsed, _metrics, _query_count, ran = pipeline_main._run_evaluation( - evaluation_mode="beir", - vdb_op="lancedb", - vdb_kwargs={"uri": "lancedb", "table_name": "nv-ingest"}, - embed_model_name="nvidia/llama-nemotron-embed-1b-v2", - embed_invoke_url=None, - embed_remote_api_key=None, - embed_modality="text", - query_csv=Path("unused.csv"), - recall_match_mode="audio_segment", - audio_match_tolerance_secs=2.0, - reranker=False, - reranker_model_name="reranker", - reranker_invoke_url=None, - reranker_api_key="", - local_reranker_backend="vllm", - local_hf_batch_size=32, - local_query_max_length=128, - beir_loader="custom_csv", - beir_dataset_name="custom_dataset", - beir_split="test", - beir_query_language=None, - beir_doc_id_field=None, - beir_k=[], - ) - - assert label == "BEIR" - assert ran is True - assert captured["cfg"].loader == "custom_csv" - assert captured["cfg"].dataset_name == "custom_dataset" - assert captured["cfg"].doc_id_field == "pdf_basename" - - -def test_pipeline_beir_evaluation_resolves_known_dataset_name(monkeypatch) -> None: - import nemo_retriever.models as model_module - import nemo_retriever.tools.recall.beir as beir_module - - captured: dict[str, BeirConfig] = {} - - def _fake_evaluate_lancedb_beir(cfg: BeirConfig): - captured["cfg"] = cfg - return ( - BeirDataset(dataset_name=cfg.dataset_name, query_ids=["q1"], queries=["query"], qrels={"q1": {"doc": 1}}), - [], - {}, - {"recall@5": 1.0}, - ) - - monkeypatch.setattr(model_module, "resolve_embed_model", lambda model_name: model_name) - monkeypatch.setattr(beir_module, "evaluate_lancedb_beir", _fake_evaluate_lancedb_beir) - - label, _elapsed, metrics, query_count, ran = pipeline_main._run_evaluation( - evaluation_mode="beir", - vdb_op="lancedb", - vdb_kwargs={"uri": "lancedb", "table_name": "nv-ingest"}, - embed_model_name="nvidia/llama-nemotron-embed-1b-v2", - embed_invoke_url=None, - embed_remote_api_key=None, - embed_modality="text", - query_csv=Path("unused.csv"), - recall_match_mode="audio_segment", - audio_match_tolerance_secs=2.0, - reranker=False, - reranker_model_name="reranker", - reranker_invoke_url=None, - reranker_api_key="", - local_reranker_backend="vllm", - local_hf_batch_size=32, - local_query_max_length=256, - beir_loader=None, - beir_dataset_name="bo767", - beir_split="validation", - beir_query_language="fr", - beir_doc_id_field=None, - beir_k=[3, 7], - ) - - assert label == "BEIR" - assert metrics == {"recall@5": 1.0} - assert query_count == 1 - assert ran is True - assert captured["cfg"].loader == "bo767_csv" - assert captured["cfg"].dataset_name == str(BO767_ANNOTATIONS_PATH) - assert captured["cfg"].split == "validation" - assert captured["cfg"].query_language == "fr" - assert captured["cfg"].doc_id_field == "pdf_page" - assert tuple(captured["cfg"].ks) == (3, 7) - assert captured["cfg"].local_query_max_length == 256 - - def test_load_beir_dataset_supports_bo767_csv_pdf_page_modality(tmp_path: Path) -> None: annotations = tmp_path / "bo767_annotations.csv" annotations.write_text( diff --git a/nemo_retriever/tests/test_graph_pipeline_cli.py b/nemo_retriever/tests/test_graph_pipeline_cli.py deleted file mode 100644 index 9bfda8861d..0000000000 --- a/nemo_retriever/tests/test_graph_pipeline_cli.py +++ /dev/null @@ -1,768 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import inspect -import json -import sys -from types import SimpleNamespace -from typing import Any - -import pandas as pd -import pytest -from typer.testing import CliRunner - -import nemo_retriever.cli.pipeline.__main__ as pipeline_main -import nemo_retriever.common.detection_summary as detection_summary_module -import nemo_retriever.examples.graph_pipeline as batch_pipeline -import nemo_retriever.ingest.execution as ingest_execution -import nemo_retriever.models as model_module -import nemo_retriever.tools.recall.beir as beir_module -from nemo_retriever.common.input_files import resolve_input_patterns - -RUNNER = CliRunner() - - -def test_staged_pipeline_caption_default_and_help_are_omni() -> None: - option = inspect.signature(pipeline_main.run).parameters["caption_model_name"].default - result = RUNNER.invoke(pipeline_main.app, ["run", "--help"]) - - assert option.default == "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" - assert result.exit_code == 0 - assert "Omni 30B BF16" in result.output - - -class _FakeDataset: - def materialize(self): - return self - - def take_all(self): - return [] - - def to_pandas(self): - return pd.DataFrame() - - def groupby(self, _key): - class _FakeGrouped: - @staticmethod - def count(): - class _FakeCounted: - @staticmethod - def count(): - return 1 - - return _FakeCounted() - - return _FakeGrouped() - - -class _FakeErrorRows: - def materialize(self): - return self - - def count(self) -> int: - return 0 - - -class _FakeIngestor: - def __init__(self) -> None: - self.extract_params = None - self.extract_kwargs = {} - self.audio_extract_params = None - self.audio_asr_params = None - self.embed_params = None - self.file_patterns = None - - def files(self, file_patterns): - self.file_patterns = file_patterns - return self - - def extract(self, params=None, **kwargs): - self.extract_params = params - self.extract_kwargs = kwargs - self.audio_extract_params = kwargs.get("audio_chunk_params") - self.audio_asr_params = kwargs.get("asr_params") - return self - - def extract_image_files(self, params): - self.extract_params = params - return self - - def extract_audio(self, params=None, asr_params=None): - self.audio_extract_params = params - self.audio_asr_params = asr_params - return self - - def extract_txt(self, params): - return self - - def extract_html(self, params): - return self - - def split(self, params): - return self - - def embed(self, params): - self.embed_params = params - return self - - def vdb_upload(self, params): - """In-graph VDB stage; real :class:`GraphIngestor` chains this before :meth:`ingest`.""" - self.vdb_upload_params = params - return self - - def ingest(self, params=None): - return _FakeDataset() - - def get_error_rows(self, dataset=None): - return _FakeErrorRows() - - -def _install_fake_lancedb(monkeypatch, *, row_count: int = 0) -> None: - class _FakeTable: - def count_rows(self) -> int: - return row_count - - class _FakeDb: - def open_table(self, _name): - return _FakeTable() - - monkeypatch.setitem(sys.modules, "lancedb", SimpleNamespace(connect=lambda _uri: _FakeDb())) - - -def test_resolve_input_file_patterns_recurses_for_directory_inputs(tmp_path) -> None: - dataset_dir = tmp_path / "earnings_consulting" - dataset_dir.mkdir() - - pdf_patterns = resolve_input_patterns(dataset_dir, "pdf") - txt_patterns = resolve_input_patterns(dataset_dir, "txt") - doc_patterns = resolve_input_patterns(dataset_dir, "doc") - - assert pdf_patterns == [str(dataset_dir / "**" / "*.pdf")] - assert txt_patterns == [ - str(dataset_dir / "**" / "*.txt"), - str(dataset_dir / "**" / "*.md"), - str(dataset_dir / "**" / "*.json"), - str(dataset_dir / "**" / "*.sh"), - ] - assert doc_patterns == [str(dataset_dir / "**" / "*.docx"), str(dataset_dir / "**" / "*.pptx")] - - -def test_graph_pipeline_resolves_nested_pdf_directories(tmp_path) -> None: - dataset_dir = tmp_path / "earnings_consulting" - nested_dir = dataset_dir / "amazon_earnings_call" - nested_dir.mkdir(parents=True) - (nested_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - patterns = batch_pipeline._resolve_file_patterns(dataset_dir, "pdf") - - assert patterns == [str(dataset_dir / "**" / "*.pdf")] - - -def test_graph_pipeline_cli_accepts_multimodal_embed_and_page_image_flags(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - missing_query_csv = tmp_path / "missing.csv" - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - - _install_fake_lancedb(monkeypatch) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--evaluation-mode", - "beir", - "--beir-loader", - "vidore_hf", - "--beir-dataset-name", - "vidore_v3_computer_science", - "--query-csv", - str(missing_query_csv), - "--embed-modality", - "text_image", - "--embed-granularity", - "page", - "--extract-infographics", - "--no-extract-page-as-image", - ], - ) - - assert result.exit_code == 0 - assert isinstance(fake_ingestor.file_patterns, list) - assert fake_ingestor.extract_params.extract_infographics is True - assert fake_ingestor.extract_params.extract_page_as_image is False - assert fake_ingestor.embed_params.embed_modality == "text_image" - assert fake_ingestor.embed_params.embed_granularity == "page" - - -def test_graph_pipeline_cli_defaults_vdb_overwrite(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - create_calls: list[dict[str, Any]] = [] - - def fake_create_ingestor(**kwargs: Any) -> _FakeIngestor: - create_calls.append(kwargs) - return fake_ingestor - - monkeypatch.setattr(ingest_execution, "create_ingestor", fake_create_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke(batch_pipeline.app, [str(dataset_dir), "--evaluation-mode", "none"]) - - assert result.exit_code == 0 - assert create_calls == [{"run_mode": "inprocess", "ray_log_to_driver": True}] - assert fake_ingestor.extract_params.dpi == 200 - assert fake_ingestor.extract_params.extract_infographics is False - assert fake_ingestor.extract_params.extract_page_as_image is True - assert fake_ingestor.vdb_upload_params.vdb_kwargs["overwrite"] is True - - -def test_graph_pipeline_cli_vdb_append_forwards_overwrite_false(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke(batch_pipeline.app, [str(dataset_dir), "--evaluation-mode", "none", "--vdb-append"]) - - assert result.exit_code == 0 - assert fake_ingestor.vdb_upload_params.vdb_kwargs["overwrite"] is False - - -def test_graph_pipeline_cli_vdb_flag_overrides_json(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--evaluation-mode", - "none", - "--vdb-kwargs-json", - json.dumps({"uri": "/tmp/custom-lancedb", "overwrite": True}), - "--vdb-append", - ], - ) - - assert result.exit_code == 0 - assert fake_ingestor.vdb_upload_params.vdb_kwargs == { - "uri": "/tmp/custom-lancedb", - "overwrite": False, - } - - -def test_graph_pipeline_cli_routes_audio_input_to_audio_ingestor(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.mp3").write_text("placeholder", encoding="utf-8") - missing_query_csv = tmp_path / "missing.csv" - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - - _install_fake_lancedb(monkeypatch) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--input-type", - "audio", - "--evaluation-mode", - "audio_recall", - "--query-csv", - str(missing_query_csv), - "--recall-match-mode", - "audio_segment", - "--audio-match-tolerance-secs", - "3.0", - "--segment-audio", - "--audio-split-type", - "time", - "--audio-split-interval", - "45", - ], - ) - - assert result.exit_code == 0 - assert isinstance(fake_ingestor.file_patterns, list) - assert fake_ingestor.audio_extract_params.split_type == "time" - assert fake_ingestor.audio_extract_params.split_interval == 45 - assert fake_ingestor.audio_asr_params.segment_audio is True - - -@pytest.mark.parametrize( - ("input_type", "filename", "param_key"), - [ - ("txt", "sample.txt", "text_params"), - ("html", "sample.html", "html_params"), - ], -) -def test_graph_pipeline_cli_text_chunk_for_text_inputs_uses_dedicated_params_only( - tmp_path, - monkeypatch, - input_type: str, - filename: str, - param_key: str, -) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / filename).write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - - _install_fake_lancedb(monkeypatch, row_count=1) - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--input-type", - input_type, - "--run-mode", - "batch", - "--text-chunk", - "--text-chunk-max-tokens", - "64", - "--text-chunk-overlap-tokens", - "8", - "--evaluation-mode", - "none", - ], - ) - - assert result.exit_code == 0, result.output - assert fake_ingestor.extract_kwargs[param_key].max_tokens == 64 - assert fake_ingestor.extract_kwargs[param_key].overlap_tokens == 8 - assert "split_config" not in fake_ingestor.extract_kwargs - - -def test_graph_pipeline_cli_text_chunk_auto_pdf_uses_plan_split_config(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - _install_fake_lancedb(monkeypatch, row_count=1) - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--run-mode", - "batch", - "--text-chunk", - "--text-chunk-max-tokens", - "64", - "--text-chunk-overlap-tokens", - "8", - "--evaluation-mode", - "none", - ], - ) - - assert result.exit_code == 0, result.output - assert fake_ingestor.extract_kwargs["split_config"]["pdf"]["max_tokens"] == 64 - assert fake_ingestor.extract_kwargs["split_config"]["pdf"]["overlap_tokens"] == 8 - - -def test_graph_pipeline_cli_allows_default_evaluation_for_pdf_inputs(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke(batch_pipeline.app, [str(dataset_dir), "--input-type", "pdf"]) - - assert result.exit_code == 0 - assert isinstance(fake_ingestor.file_patterns, list) - - -def test_graph_pipeline_cli_rejects_invalid_recall_mode(tmp_path) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - result = RUNNER.invoke(batch_pipeline.app, [str(dataset_dir), "--evaluation-mode", "recall"]) - - assert result.exit_code != 0 - assert result.exception is not None - assert "Unsupported --evaluation-mode: 'recall'" in str(result.exception) - - -def test_graph_pipeline_cli_rejects_audio_recall_for_pdf_inputs(tmp_path) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - result = RUNNER.invoke( - batch_pipeline.app, - [str(dataset_dir), "--input-type", "pdf", "--evaluation-mode", "audio_recall"], - ) - - assert result.exit_code != 0 - assert result.exception is not None - assert "--evaluation-mode=audio_recall is only supported with --input-type=audio" in str(result.exception) - - -def test_graph_pipeline_cli_routes_beir_mode_to_evaluator(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setattr(pipeline_main, "_count_uploadable_vdb_records", lambda _records: 1) - monkeypatch.setattr(detection_summary_module, "print_run_summary", lambda *args, **kwargs: None) - - _install_fake_lancedb(monkeypatch, row_count=1) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - captured = {} - - def _fake_evaluate(cfg): - captured["cfg"] = cfg - return type("Dataset", (), {"query_ids": ["1", "2"]})(), [], {}, {"ndcg@10": 0.75, "recall@5": 0.6} - - monkeypatch.setattr(beir_module, "evaluate_lancedb_beir", _fake_evaluate) - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--evaluation-mode", - "beir", - "--beir-loader", - "vidore_hf", - "--beir-dataset-name", - "vidore_v3_computer_science", - "--beir-k", - "5", - "--beir-k", - "10", - ], - ) - - assert result.exit_code == 0 - assert captured["cfg"].loader == "vidore_hf" - assert captured["cfg"].dataset_name == "vidore_v3_computer_science" - assert tuple(captured["cfg"].ks) == (5, 10) - - -def test_graph_pipeline_cli_accepts_harness_runtime_metric_flags(tmp_path, monkeypatch) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - missing_query_csv = tmp_path / "missing.csv" - runtime_dir = tmp_path / "runtime_metrics" - - fake_ingestor = _FakeIngestor() - monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_kwargs: fake_ingestor) - monkeypatch.setitem( - sys.modules, - "ray", - SimpleNamespace(shutdown=lambda: None, is_initialized=lambda: True), - ) - - _install_fake_lancedb(monkeypatch) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - monkeypatch.setattr(pipeline_main, "_count_uploadable_vdb_records", lambda _records: 1) - monkeypatch.setattr( - pipeline_main, - "_run_evaluation", - lambda **_kwargs: ("BEIR", 0.0, {"recall@5": 0.6}, 2, True), - ) - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--evaluation-mode", - "beir", - "--beir-loader", - "vidore_hf", - "--beir-dataset-name", - "vidore_v3_computer_science", - "--query-csv", - str(missing_query_csv), - "--runtime-metrics-dir", - str(runtime_dir), - "--runtime-metrics-prefix", - "sample-run", - "--no-recall-details", - ], - ) - - assert result.exit_code == 0 - summary_path = runtime_dir / "sample-run.runtime.summary.json" - assert summary_path.exists() - payload = json.loads(summary_path.read_text(encoding="utf-8")) - assert payload["recall_details"] is False - assert payload["evaluation_mode"] == "beir" - - -def test_graph_pipeline_cli_service_mode_rejects_ingest_flag(tmp_path) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--run-mode", - "service", - "--ocr-invoke-url", - "http://localhost:9000/v1/infer", - ], - ) - - assert result.exit_code != 0 - assert "--run-mode=service" in result.output - assert "--ocr-invoke-url" in result.output - - -def test_graph_pipeline_cli_service_mode_lists_all_incompatible_flags(tmp_path) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--run-mode", - "service", - "--ocr-invoke-url", - "http://localhost:9000/v1/infer", - "--ray-address", - "ray://localhost:10001", - "--caption-device", - "cuda:0", - ], - ) - - assert result.exit_code != 0 - assert "--ocr-invoke-url" in result.output - assert "--ray-address" in result.output - assert "--caption-device" in result.output - - -def test_graph_pipeline_cli_service_mode_allows_extract_and_embed_flags(tmp_path, monkeypatch) -> None: - """Flags whose values flow through to ``ServiceIngestor`` must not be rejected.""" - import nemo_retriever.service.service_ingestor as service_ingestor_module - - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - save_dir = tmp_path / "save" - - captured: dict[str, Any] = {} - - class _FakeServiceIngestor(list): - def __init__(self, *args, **kwargs) -> None: - super().__init__() - - def files(self, _files): - return self - - def extract(self, params=None, *, split_config=None, extraction_mode="auto", **_kwargs): - captured["extract_params"] = params - captured["split_config"] = split_config - captured["extraction_mode"] = extraction_mode - return self - - def dedup(self, params=None, **_kwargs): - captured["dedup_params"] = params - return self - - def embed(self, params=None, **_kwargs): - captured["embed_params"] = params - return self - - def ingest(self, *args, **kwargs): - return self - - monkeypatch.setattr(service_ingestor_module, "ServiceIngestor", _FakeServiceIngestor) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--run-mode", - "service", - "--service-url", - "http://localhost:7670", - "--embed-model-name", - "nvidia/llama-3.2-nv-embedqa-1b-v2", - "--method", - "ocr", - "--dpi", - "300", - "--no-extract-text", - "--embed-granularity", - "page", - "--dedup", - "--dedup-iou-threshold", - "0.6", - "--text-chunk", - "--text-chunk-max-tokens", - "64", - "--evaluation-mode", - "none", - "--save-intermediate", - str(save_dir), - ], - ) - - assert result.exit_code == 0, result.output - assert captured["extract_params"].method == "ocr" - assert captured["extract_params"].dpi == 300 - assert captured["extract_params"].extract_text is False - assert captured["embed_params"].embed_granularity == "page" - assert captured["dedup_params"].iou_threshold == 0.6 - assert captured["split_config"]["pdf"]["max_tokens"] == 64 - - -def test_graph_pipeline_cli_service_mode_rejects_vdb_flags(tmp_path) -> None: - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - - result_no_vdb = RUNNER.invoke( - batch_pipeline.app, - [str(dataset_dir), "--run-mode", "service", "--no-vdb"], - ) - assert result_no_vdb.exit_code != 0 - assert "--no-vdb" in result_no_vdb.output - - result_overwrite = RUNNER.invoke( - batch_pipeline.app, - [str(dataset_dir), "--run-mode", "service", "--vdb-overwrite"], - ) - assert result_overwrite.exit_code != 0 - assert "--vdb-overwrite" in result_overwrite.output - - result_append = RUNNER.invoke( - batch_pipeline.app, - [str(dataset_dir), "--run-mode", "service", "--vdb-append"], - ) - assert result_append.exit_code != 0 - assert "--vdb-overwrite" in result_append.output - - -def test_graph_pipeline_cli_service_mode_accepts_allowlisted_flags(tmp_path, monkeypatch) -> None: - import nemo_retriever.service.service_ingestor as service_ingestor_module - - dataset_dir = tmp_path / "dataset" - dataset_dir.mkdir() - (dataset_dir / "sample.pdf").write_text("placeholder", encoding="utf-8") - save_dir = tmp_path / "save" - - class _FakeServiceIngestor(list): - def __init__(self, *args, **kwargs) -> None: - super().__init__() - - def files(self, _files): - return self - - def extract(self, *args, **kwargs): - return self - - def embed(self, *args, **kwargs): - return self - - def ingest(self, *args, **kwargs): - return self - - monkeypatch.setattr(service_ingestor_module, "ServiceIngestor", _FakeServiceIngestor) - monkeypatch.setattr(model_module, "resolve_embed_model", lambda _name: "fake-embed-model") - - result = RUNNER.invoke( - batch_pipeline.app, - [ - str(dataset_dir), - "--run-mode", - "service", - "--service-url", - "http://localhost:7670", - "--service-concurrency", - "2", - "--embed-model-name", - "nvidia/llama-3.2-nv-embedqa-1b-v2", - "--evaluation-mode", - "none", - "--save-intermediate", - str(save_dir), - ], - ) - - assert result.exit_code == 0, result.output diff --git a/nemo_retriever/tests/test_ingest_service.py b/nemo_retriever/tests/test_ingest_service.py new file mode 100644 index 0000000000..54272f0632 --- /dev/null +++ b/nemo_retriever/tests/test_ingest_service.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from nemo_retriever.common.params import EmbedParams, ExtractParams, TextChunkParams +from nemo_retriever.common.policy import validate_pipeline_spec +from nemo_retriever.common.schemas.pipeline_spec import PipelineSpec +from nemo_retriever.ingest.service import ServiceIngestRequest, build_service_ingestor, execute_service_ingest_request +from nemo_retriever.service.config import PipelineOverridesConfig +from nemo_retriever.service.service_ingestor import ServiceIngestor + + +def test_build_service_ingestor_wires_extract_embed_and_chunking(tmp_path: Path) -> None: + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4") + + ingestor = build_service_ingestor( + ServiceIngestRequest( + documents=[str(pdf)], + input_type="pdf", + extract_params=ExtractParams(method="ocr", extract_text=False, dpi=300), + embed_params=EmbedParams(embed_granularity="page"), + text_chunk_params=TextChunkParams(max_tokens=64, overlap_tokens=8), + enable_text_chunk=True, + ) + ) + + assert isinstance(ingestor, ServiceIngestor) + payload = ingestor._pipeline_payload() + assert payload is not None + assert payload["extraction_mode"] == "pdf" + assert payload["extract_params"]["method"] == "ocr" + assert payload["extract_params"]["extract_text"] is False + assert payload["extract_params"]["dpi"] == 300 + assert "batch_tuning" not in payload["extract_params"] + assert payload["split_config"]["pdf"]["max_tokens"] == 64 + assert payload["split_config"]["pdf"]["overlap_tokens"] == 8 + assert payload["embed_params"]["embed_granularity"] == "page" + assert "model_name" not in payload["embed_params"] + + validate_pipeline_spec( + PipelineSpec.model_validate(ingestor._pipeline_spec), + PipelineOverridesConfig().to_policy(), + ) + + +def test_build_service_ingestor_does_not_forward_environment_api_key(monkeypatch, tmp_path: Path) -> None: + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(b"%PDF-1.4") + monkeypatch.setenv("NGC_API_KEY", "environment-secret") + stage_params: dict[str, object] = {} + original_extract = ServiceIngestor.extract + original_embed = ServiceIngestor.embed + + def capture_extract(self, params=None, **kwargs): + stage_params["extract"] = params + return original_extract(self, params, **kwargs) + + def capture_embed(self, params=None, **kwargs): + stage_params["embed"] = params + return original_embed(self, params, **kwargs) + + monkeypatch.setattr(ServiceIngestor, "extract", capture_extract) + monkeypatch.setattr(ServiceIngestor, "embed", capture_embed) + + ingestor = build_service_ingestor( + ServiceIngestRequest( + documents=[str(pdf)], + input_type="pdf", + extract_params=ExtractParams(method="pdfium", use_table_structure=True), + embed_params=EmbedParams(embed_granularity="page"), + ) + ) + + payload = ingestor._pipeline_payload() + assert payload is not None + assert isinstance(stage_params["extract"], ExtractParams) + assert isinstance(stage_params["embed"], EmbedParams) + assert payload["extract_params"]["method"] == "pdfium" + assert payload["extract_params"]["use_table_structure"] is True + assert payload["embed_params"]["embed_granularity"] == "page" + assert not ({"api_key", "page_elements_api_key", "ocr_api_key"} & payload["extract_params"].keys()) + assert "api_key" not in payload["embed_params"] + + +def test_execute_service_ingest_request_raises_for_document_failures(monkeypatch, tmp_path: Path) -> None: + request = ServiceIngestRequest(documents=[str(tmp_path / "doc.pdf")], input_type="pdf") + failed_result = SimpleNamespace(failures=[("doc.pdf", "HTTP 400: invalid request")]) + monkeypatch.setattr( + "nemo_retriever.ingest.service.build_service_ingestor", + lambda _request: SimpleNamespace(ingest=lambda: failed_result), + ) + + with pytest.raises(RuntimeError, match=r"failed for 1 document\(s\).+doc.pdf.+HTTP 400"): + execute_service_ingest_request(request) diff --git a/nemo_retriever/tests/test_pipeline_helpers.py b/nemo_retriever/tests/test_pipeline_helpers.py deleted file mode 100644 index aa1b24b184..0000000000 --- a/nemo_retriever/tests/test_pipeline_helpers.py +++ /dev/null @@ -1,304 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from pathlib import Path -from types import SimpleNamespace - -import pandas as pd -import pytest -import typer - -import nemo_retriever.cli.pipeline as pipeline_pkg -from nemo_retriever.ingest.service import ( - execute_service_ingest_request, - ServiceIngestRequest, - build_service_ingestor, -) -from nemo_retriever.common.params import EmbedParams, ExtractParams, TextChunkParams -from nemo_retriever.cli.pipeline.__main__ import ( - _build_embed_params, - _collect_results, - _count_input_units, - _count_uploadable_vdb_records, - _parse_vdb_kwargs_json, - _resolve_file_patterns, -) -from nemo_retriever.service.config import PipelineOverridesConfig -from nemo_retriever.common.schemas.pipeline_spec import PipelineSpec -from nemo_retriever.common.policy import validate_pipeline_spec -from nemo_retriever.service.service_ingestor import ServiceIngestor - - -def test_pipeline_package_exports_cli_app_and_run() -> None: - from nemo_retriever.cli.pipeline.__main__ import app, run - - assert pipeline_pkg.app is app - assert pipeline_pkg.run is run - assert set(pipeline_pkg.__all__) == {"app", "run"} - - -@pytest.mark.parametrize( - ("input_type", "files", "expected_globs"), - [ - ("pdf", ["nested/doc.pdf"], ["*.pdf"]), - ("doc", ["deck.pptx", "report.docx"], ["*.docx", "*.pptx"]), - ("image", ["plot.png"], ["*.png"]), - ("video", ["clip.mp4"], ["*.mp4"]), - ], -) -def test_resolve_file_patterns_recurses_directory_inputs( - tmp_path: Path, - input_type: str, - files: list[str], - expected_globs: list[str], -) -> None: - for name in files: - path = tmp_path / name - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(b"x") - - assert _resolve_file_patterns(tmp_path, input_type) == [str(tmp_path / "**" / glob) for glob in expected_globs] - - -def test_build_service_ingestor_wires_extract_embed_and_chunking(tmp_path: Path) -> None: - pdf = tmp_path / "doc.pdf" - pdf.write_bytes(b"%PDF-1.4") - - ingestor = build_service_ingestor( - ServiceIngestRequest( - documents=[str(pdf)], - input_type="pdf", - extract_params=ExtractParams(method="ocr", extract_text=False, dpi=300), - embed_params=EmbedParams(embed_granularity="page"), - text_chunk_params=TextChunkParams(max_tokens=64, overlap_tokens=8), - enable_text_chunk=True, - ) - ) - - assert isinstance(ingestor, ServiceIngestor) - payload = ingestor._pipeline_payload() - assert payload is not None - assert payload["extraction_mode"] == "pdf" - assert payload["extract_params"]["method"] == "ocr" - assert payload["extract_params"]["extract_text"] is False - assert payload["extract_params"]["dpi"] == 300 - assert "batch_tuning" not in payload["extract_params"] - assert payload["split_config"]["pdf"]["max_tokens"] == 64 - assert payload["split_config"]["pdf"]["overlap_tokens"] == 8 - assert payload["embed_params"]["embed_granularity"] == "page" - assert "model_name" not in payload["embed_params"] - - validate_pipeline_spec( - PipelineSpec.model_validate(ingestor._pipeline_spec), - PipelineOverridesConfig().to_policy(), - ) - - -def test_build_service_ingestor_does_not_forward_environment_api_key(monkeypatch, tmp_path: Path) -> None: - pdf = tmp_path / "doc.pdf" - pdf.write_bytes(b"%PDF-1.4") - monkeypatch.setenv("NGC_API_KEY", "environment-secret") - stage_params: dict[str, object] = {} - original_extract = ServiceIngestor.extract - original_embed = ServiceIngestor.embed - - def capture_extract(self, params=None, **kwargs): - stage_params["extract"] = params - return original_extract(self, params, **kwargs) - - def capture_embed(self, params=None, **kwargs): - stage_params["embed"] = params - return original_embed(self, params, **kwargs) - - monkeypatch.setattr(ServiceIngestor, "extract", capture_extract) - monkeypatch.setattr(ServiceIngestor, "embed", capture_embed) - - ingestor = build_service_ingestor( - ServiceIngestRequest( - documents=[str(pdf)], - input_type="pdf", - extract_params=ExtractParams(method="pdfium", use_table_structure=True), - embed_params=EmbedParams(embed_granularity="page"), - ) - ) - - payload = ingestor._pipeline_payload() - assert payload is not None - assert isinstance(stage_params["extract"], ExtractParams) - assert isinstance(stage_params["embed"], EmbedParams) - assert payload["extract_params"]["method"] == "pdfium" - assert payload["extract_params"]["use_table_structure"] is True - assert payload["embed_params"]["embed_granularity"] == "page" - assert not ({"api_key", "page_elements_api_key", "ocr_api_key"} & payload["extract_params"].keys()) - assert "api_key" not in payload["embed_params"] - - -def test_execute_service_ingest_request_raises_for_document_failures(monkeypatch, tmp_path: Path) -> None: - request = ServiceIngestRequest(documents=[str(tmp_path / "doc.pdf")], input_type="pdf") - failed_result = SimpleNamespace(failures=[("doc.pdf", "HTTP 400: invalid request")]) - monkeypatch.setattr( - "nemo_retriever.ingest.service.build_service_ingestor", - lambda _request: SimpleNamespace(ingest=lambda: failed_result), - ) - - with pytest.raises(RuntimeError, match=r"failed for 1 document\(s\).+doc.pdf.+HTTP 400"): - execute_service_ingest_request(request) - - -def test_resolve_file_patterns_returns_existing_file_verbatim(tmp_path: Path) -> None: - path = tmp_path / "doc.pdf" - path.write_bytes(b"x") - - assert _resolve_file_patterns(path, "audio") == [str(path)] - - -def test_resolve_file_patterns_rejects_missing_or_empty_inputs(tmp_path: Path) -> None: - with pytest.raises(typer.BadParameter, match="Path does not exist"): - _resolve_file_patterns(tmp_path / "missing", "pdf") - - (tmp_path / "sidecar.json").write_text("{}", encoding="utf-8") - with pytest.raises(typer.BadParameter, match="No files found"): - _resolve_file_patterns(tmp_path, "pdf") - - -def test_parse_vdb_kwargs_json_keeps_backend_kwargs_opaque() -> None: - assert _parse_vdb_kwargs_json(None) == {} - assert _parse_vdb_kwargs_json('{"collection_name": "docs", "uri": "http://localhost:19530"}') == { - "collection_name": "docs", - "uri": "http://localhost:19530", - } - - -def test_parse_vdb_kwargs_json_rejects_non_object_json() -> None: - with pytest.raises(typer.BadParameter, match="JSON object"): - _parse_vdb_kwargs_json('["not", "an", "object"]') - - -def test_build_embed_params_forwards_remote_and_modality_flags() -> None: - params = _build_embed_params( - embed_model_name="nvidia/test-embed", - embed_invoke_url="http://embed.example/v1", - embed_remote_api_key="nvapi-secret", - embed_modality="text_image", - text_elements_modality="text", - structured_elements_modality="image", - embed_granularity="element", - embed_actors=2, - embed_batch_size=16, - embed_cpus_per_actor=1.5, - embed_gpus_per_actor=0.5, - ) - - assert isinstance(params, EmbedParams) - assert params.model_name == "nvidia/test-embed" - assert params.embed_invoke_url == "http://embed.example/v1" - assert params.api_key == "nvapi-secret" - assert params.embed_modality == "text_image" - assert params.text_elements_modality == "text" - assert params.structured_elements_modality == "image" - assert params.embed_granularity == "element" - assert params.inference_batch_size == 16 - assert params.batch_tuning.embed_workers == 2 - assert params.batch_tuning.embed_batch_size == 16 - assert params.batch_tuning.embed_cpus_per_actor == 1.5 - assert params.batch_tuning.gpu_embed == 0.0 - - -class TestCollectResults: - """Ingest returns a DataFrame (``ingestor.ingest()`` → ``ds.to_pandas()``); _collect_results consumes it.""" - - def test_batch_mode_accepts_ingest_dataframe(self): - rows = [ - {"source_id": "a", "text": "hello"}, - {"source_id": "a", "text": "world"}, - {"source_id": "b", "text": "!"}, - ] - # Same shape as the graph executor return after ``Dataset.to_pandas()``. - result_df = pd.DataFrame(rows) - - records, df, download_time, num_units = _collect_results("batch", result_df) - - assert records == rows - assert df is result_df - assert isinstance(df, pd.DataFrame) - assert list(df.columns) == ["source_id", "text"] - assert len(df) == 3 - # ``source_id`` has two distinct values → that is the unit count. - assert num_units == 2 - assert download_time >= 0.0 - - def test_batch_mode_handles_empty_result(self): - result_df = pd.DataFrame() - records, df, download_time, num_units = _collect_results("batch", result_df) - assert records == [] - assert df.empty - # Empty DataFrame has no columns → falls through to len(df.index) == 0. - assert num_units == 0 - assert download_time >= 0.0 - - def test_inprocess_mode_accepts_dataframe_directly(self): - rows = [ - {"source_id": "a", "text": "x"}, - {"source_id": "b", "text": "y"}, - ] - df_in = pd.DataFrame(rows) - - records, df_out, download_time, num_units = _collect_results("inprocess", df_in) - - # The DataFrame is passed through unchanged (same object). - assert df_out is df_in - assert records == rows - # inprocess mode never incurs Ray download time. - assert download_time == 0.0 - assert num_units == 2 - - -def test_collect_results_accepts_inprocess_dataframe() -> None: - df_in = pd.DataFrame([{"source_path": "/a.pdf"}, {"source_path": "/b.pdf"}]) - - records, df_out, download_time, num_units = _collect_results("inprocess", df_in) - - assert df_out is df_in - assert records == [{"source_path": "/a.pdf"}, {"source_path": "/b.pdf"}] - assert download_time == 0.0 - assert num_units == 2 - - -def test_to_client_vdb_records_returns_empty_list_when_nothing_uploadable() -> None: - from nemo_retriever.common.vdb.records import to_client_vdb_records - - assert to_client_vdb_records([]) == [] - assert to_client_vdb_records([{"text": "no embedding"}]) == [] - - -def test_count_uploadable_vdb_records_filters_rows_without_embedding_or_text() -> None: - rows = [ - { - "text": "keep", - "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, - "source_id": "/tmp/doc-a.pdf", - "page_number": 1, - }, - { - "text": "drop missing embedding", - "source_id": "/tmp/doc-a.pdf", - "page_number": 2, - }, - { - "text_embeddings_1b_v2": {"embedding": [0.3, 0.4]}, - "source_id": "/tmp/doc-a.pdf", - "page_number": 3, - }, - ] - - assert _count_uploadable_vdb_records(rows) == 1 - assert _count_uploadable_vdb_records([]) == 0 - - -def test_count_input_units_prefers_source_id_then_source_path() -> None: - assert _count_input_units(pd.DataFrame({"source_id": ["a", "a", "b"]})) == 2 - assert _count_input_units(pd.DataFrame({"source_path": ["/a", "/b", "/b"]})) == 2 - assert _count_input_units(pd.DataFrame({"text": ["x", "y", "z"]})) == 3 diff --git a/nemo_retriever/tests/test_detection_summary.py b/nemo_retriever/tests/test_recall_summary.py similarity index 94% rename from nemo_retriever/tests/test_detection_summary.py rename to nemo_retriever/tests/test_recall_summary.py index 9575ccee79..beadb58afd 100644 --- a/nemo_retriever/tests/test_detection_summary.py +++ b/nemo_retriever/tests/test_recall_summary.py @@ -1,6 +1,6 @@ from pathlib import Path -from nemo_retriever.common.detection_summary import print_run_summary +from nemo_retriever.tools.recall.core import print_run_summary def test_print_run_summary_sorts_evaluation_metrics_by_numeric_k(capsys) -> None: diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index 378fd4433b..42a13a6cd2 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -93,15 +93,9 @@ def test_root_help_lists_only_product_workflows() -> None: assert f"│ {developer_command} " not in result.output -def test_pipeline_compatibility_command_is_hidden_but_callable() -> None: - result = RUNNER.invoke(cli_main.app, ["pipeline", "--help"]) - - assert result.exit_code == 0 - - @pytest.mark.parametrize( "removed_command", - ("txt", "html", "local", "audio", "image", "pdf", "chart", "compare"), + ("txt", "html", "local", "audio", "image", "pdf", "chart", "compare", "pipeline"), ) def test_removed_root_commands_are_not_callable(removed_command: str) -> None: result = RUNNER.invoke(cli_main.app, [removed_command, "--help"]) diff --git a/nemo_retriever/tests/test_slim_imports_no_triton.py b/nemo_retriever/tests/test_slim_imports_no_triton.py index a83c4ed7dd..9ca92c8ee5 100644 --- a/nemo_retriever/tests/test_slim_imports_no_triton.py +++ b/nemo_retriever/tests/test_slim_imports_no_triton.py @@ -2,7 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Ensure core pipeline / API imports never load ``tritonclient`` at import time. +"""Ensure core ingest imports never load ``tritonclient`` at import time. Remote-NIM and slim installs omit ``tritonclient``; gRPC code paths import it lazily. """ @@ -33,14 +33,13 @@ def _guard(name, globals=None, locals=None, fromlist=(), level=0): from nemo_retriever.models.nim.util import create_inference_client # noqa: F401 from nemo_retriever.ingestor.graph_ingestor import GraphIngestor # noqa: F401 -from nemo_retriever.cli.pipeline import __main__ as _pipeline_main # noqa: F401 print("slim_imports_ok") """ def test_core_imports_do_not_require_tritonclient_at_import_time() -> None: - """Fresh interpreter: block tritonclient, then import hot paths used by remote pipeline.""" + """Fresh interpreter: block tritonclient, then import remote-ingest paths.""" root = Path(__file__).resolve().parents[1] src = root / "src" env = os.environ.copy() diff --git a/nemo_retriever/tests/test_src_documentation_snippets.py b/nemo_retriever/tests/test_src_documentation_snippets.py index 2a04233d4a..0bf52ea2a1 100644 --- a/nemo_retriever/tests/test_src_documentation_snippets.py +++ b/nemo_retriever/tests/test_src_documentation_snippets.py @@ -60,11 +60,6 @@ def _iter_markdown_python_blocks() -> list[tuple[str, str]]: "nemo_retriever/src/nemo_retriever/tools/evaluation/README.md", "nemo_retriever/src/nemo_retriever/common/vdb/README.md", ) -_PUBLIC_GRAPH_PIPELINE_DOCS = ( - "docs/docs/extraction/workflow-document-ingestion.md", - "nemo_retriever/README.md", - "nemo_retriever/src/nemo_retriever/tools/evaluation/README.md", -) _UNSUPPORTED_DIRECT_RETRIEVER_KWARGS = frozenset( { "vdb", @@ -76,7 +71,6 @@ def _iter_markdown_python_blocks() -> list[tuple[str, str]]: "reranker", } ) -_UNSUPPORTED_GRAPH_PIPELINE_OPTIONS = frozenset({"--lancedb-uri"}) def _public_doc_path(root: Path, rel_path: str) -> Path | None: @@ -120,30 +114,6 @@ def _iter_public_retriever_doc_code() -> list[tuple[str, str]]: return blocks -def _iter_public_graph_pipeline_commands() -> list[tuple[str, str]]: - root = _repo_root() - commands: list[tuple[str, str]] = [] - for rel_path in _PUBLIC_GRAPH_PIPELINE_DOCS: - path = _public_doc_path(root, rel_path) - if path is None: - continue - text = path.read_text(encoding="utf-8", errors="replace") - for i, code in enumerate(re.findall(r"```bash\n(.*?)```", text, re.DOTALL)): - lines = code.splitlines() - command_idx = 0 - for line_idx, line in enumerate(lines): - if "python -m nemo_retriever.examples.graph_pipeline" not in line: - continue - command_lines = [line] - next_idx = line_idx + 1 - while command_lines[-1].rstrip().endswith("\\") and next_idx < len(lines): - command_lines.append(lines[next_idx]) - next_idx += 1 - commands.append((f"{rel_path}#bash-{i}-cmd-{command_idx}", "\n".join(command_lines))) - command_idx += 1 - return commands - - def _retriever_call_unsupported_kwargs(code: str) -> list[str]: tree = ast.parse(code) found: list[str] = [] @@ -182,17 +152,6 @@ def test_public_retriever_examples_do_not_use_unsupported_constructor_kwargs() - assert not violations, "Unsupported kwargs in public direct Retriever(...) examples:\n" + "\n".join(violations) -def test_public_graph_pipeline_examples_do_not_use_unsupported_options() -> None: - """Public ``graph_pipeline`` examples should not use options that command rejects.""" - violations = [] - for block_id, command in _iter_public_graph_pipeline_commands(): - unsupported_options = [option for option in _UNSUPPORTED_GRAPH_PIPELINE_OPTIONS if option in command] - if unsupported_options: - violations.append(f"{block_id}: {', '.join(sorted(unsupported_options))}") - - assert not violations, "Unsupported options in public graph_pipeline examples:\n" + "\n".join(violations) - - def test_graph_readme_smallest_example() -> None: """``graph/README.md`` — single :class:`UDFOperator` on a :class:`Graph`.""" from nemo_retriever.graph import Graph, UDFOperator diff --git a/skills/nemo-retriever/references/cli/ingest.md b/skills/nemo-retriever/references/cli/ingest.md index b663151b17..ea526535ad 100644 --- a/skills/nemo-retriever/references/cli/ingest.md +++ b/skills/nemo-retriever/references/cli/ingest.md @@ -109,11 +109,7 @@ when available, and an embedding vector. The root ingest entrypoint expands inputs, builds a manifest, resolves the selected profile into typed ingest options, and calls the canonical ingest execution path. The manifest planner routes PDF/document, image, text, HTML, -audio, and video branches without relying on `retriever pipeline run`. - -`retriever pipeline run` remains callable, but hidden from root help, for legacy -or development behavior such as intermediate Parquet artifacts, pipeline -reports, eval, recall, or harness work. +audio, and video branches through the same supported interface. ## Common failure modes