Retrieval SDG

View as Markdown

NeMo Platform exposes Nemotron embed/rerank Stage 0 (sdg) and Stage 1 (prep) as dedicated Data Designer jobs. They wrap data-designer-retrieval-sdg and do not go through nemo data-designer create.

Prerequisites

  • A text corpus published as a fileset in the job workspace, or an hf:// dataset URI.
  • An Inference Gateway provider for the Stage 0 chat model and, when different, an embedding provider.
  • For GPU mining, a platform model entity with an attached fileset containing the encoder and tokenizer.

Generate (Stage 0, CPU)

nemo data-designer retrieval-generate --workspace default --spec '{
"corpus": "default/my-docs",
"provider": "default/nvidia-build",
"artifact_extraction_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"qa_generation_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"quality_judge_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"embed_model": "nvidia/nemotron-3-embed-1b"
}'

Pass the chat and embedding model names served by Inference Gateway. All four model roles call IGW through provider (workspace/name), with optional chat_provider / embed_provider overrides. Do not pass raw NVIDIA_API_KEY.

Equivalent helper:

nemo data-designer retrieval generate \
--corpus default/my-docs \
--provider default/nvidia-build \
--chat-model nvidia/nemotron-3.5-lightning-30b-a3b \
--embed-model nvidia/nemotron-3-embed-1b

Output fileset includes Q&A JSONL and generation_result.json.

Preview without a full job:

nemo data-designer retrieval-preview --spec '{"generate":{"corpus":"default/my-docs","provider":"default/nvidia-build","artifact_extraction_model":"nvidia/nemotron-3.5-lightning-30b-a3b","qa_generation_model":"nvidia/nemotron-3.5-lightning-30b-a3b","quality_judge_model":"nvidia/nemotron-3.5-lightning-30b-a3b","embed_model":"nvidia/nemotron-3-embed-1b"},"num_records":1}'

Prepare (Stage 1)

Conversion produces eval_beir/ and training JSON. Automodel needs hard negatives, so run with enable_mining: true and a fileset-backed encoder model. Without mining, rows keep neg_doc: [], which training rejects.

Skip live Stage 0 by pointing sdg_input at a fileset holding an existing generation_result.json, or at the published NVDocs dump. For a differently named dump, put the filename in the reference or set generation_file. Directories default to generation_result.json; the job does not guess among JSON files.

nemo data-designer retrieval-prepare --spec '{
"sdg_input": "default/stage0-out",
"enable_mining": true,
"model": "default/nemotron-3-embed-1b"
}'
nemo data-designer retrieval-prepare --spec '{
"sdg_input": "hf://nvidia/Retrieval-Synthetic-NVDocs-v1@1c0d1856f3fb595b2dda98d4b61061fa6d782d51/nv_pp_dd_sdg.json",
"enable_mining": true,
"model": "default/nemotron-3-embed-1b"
}'
nemo data-designer retrieval-prepare --spec '{
"sdg_input": "default/retrieval-synthetic-nvdocs-v1",
"generation_file": "nv_pp_dd_sdg.json",
"enable_mining": true,
"model": "default/nemotron-3-embed-1b"
}'

Mining runs on GPU and needs about 40 GB. model is a platform entity with an attached encoder fileset; an Inference Gateway endpoint entity has no fileset and cannot serve as the mining encoder.

Mining embeds every training query and the corpus, so it dominates Stage 1 wall time. The mining.query_embedding_batch_size and mining.document_embedding_batch_size defaults of 16 leave a 48 GB GPU mostly idle; raise them to 64–128 for a large split.

All retrieval steps use the same container-backed execution profile so they share job storage across generation, conversion, model staging, and mining. Configure it with data_designer.job_executor_profile, or pass --profile for an individual job. The local platform uses gpu; Kubernetes uses default.

Chain generate then prepare

nemo data-designer retrieval-run --spec '{
"generate": {
"corpus": "default/my-docs",
"provider": "default/nvidia-build",
"artifact_extraction_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"qa_generation_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"quality_judge_model": "nvidia/nemotron-3.5-lightning-30b-a3b",
"embed_model": "nvidia/nemotron-3-embed-1b"
},
"prepare": {
"enable_mining": true,
"model": "default/nemotron-3-embed-1b"
}
}'

This is jobs-service multi-step execution (CPU then optional GPU), not Data Designer in-config workflow chaining.

Use the Stage 1 output

Stage 1 saves one artifacts fileset containing training.jsonl, eval_beir/ (corpus.jsonl, queries.jsonl, qrels/test.tsv), and the wrapped train.json that hard-negative mining consumes.

Pass that fileset to both downstream consumers without repackaging it:

  • Automodel dataset.training — dataset discovery selects training.jsonl and ignores the non-JSONL train.json beside it.
  • nemo evaluator retrieve-eval dataset — the BEIR loader accepts either the BEIR directory or a fileset root containing eval_beir.

Prepare without mining (enable_mining: false) writes neg_doc: [] on every training row. Encoder fine-tuning still samples train_n_passages - 1 negatives (default 4) and fails with neg_doc must contain at least 1 document to sample N negatives. Before nemo customization automodel submit, download training.jsonl and require a non-empty neg_doc list on every checked row. A single empty list can be selected by the collator and fail the run. If any row is empty, re-run retrieval-prepare with "enable_mining": true and a fileset-backed encoder model. To mine an unmined fileset without regenerating frozen eval_beir, set train_input_file to that fileset instead of sdg_input.

nemo files list <artifacts-fileset> --workspace default

Freeze the fileset once base-versus-tuned comparison starts, so both runs score against the same eval_beir split.

SDK

from nemo_platform import NeMoPlatform
from nemo_data_designer_plugin.jobs.retrieval_spec import RetrievalGenerateJobConfig
client = NeMoPlatform(base_url="http://localhost:8080", workspace="default")
client.data_designer.retrieval_generate(
RetrievalGenerateJobConfig(
corpus="default/my-docs",
provider="default/nvidia-build",
artifact_extraction_model="nvidia/nemotron-3.5-lightning-30b-a3b",
qa_generation_model="nvidia/nemotron-3.5-lightning-30b-a3b",
quality_judge_model="nvidia/nemotron-3.5-lightning-30b-a3b",
embed_model="nvidia/nemotron-3-embed-1b",
)
)

Next Steps

  • Worked Stage 0 notebook: Generate retrieval training data.
  • Continue with mining, Automodel bi_encoder, Retriever NIM 2.2.0, and retrieve-eval in Embedding Model Customization.
  • See the Data Designer CLI for configuration and command details.
  • Stage 1 emits training.jsonl plus eval_beir/ (corpus.jsonl, queries.jsonl, qrels/test.tsv). Tiny corpora can leave the train split empty — generate with enough files (50+) before mining.