Retrieval SDG
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)
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:
Output fileset includes Q&A JSONL and generation_result.json.
Preview without a full job:
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.
CLI
Python SDK
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
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 selectstraining.jsonland ignores the non-JSONLtrain.jsonbeside it. nemo evaluator retrieve-evaldataset— the BEIR loader accepts either the BEIR directory or a fileset root containingeval_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.
Freeze the fileset once base-versus-tuned comparison starts, so both runs score
against the same eval_beir split.
SDK
Next Steps
- Worked Stage 0 notebook: Generate retrieval training data.
- Continue with mining, Automodel
bi_encoder, Retriever NIM 2.2.0, andretrieve-evalin Embedding Model Customization. - See the Data Designer CLI for configuration and command details.
- Stage 1 emits
training.jsonlpluseval_beir/(corpus.jsonl,queries.jsonl,qrels/test.tsv). Tiny corpora can leave the train split empty — generate with enough files (50+) before mining.