Retrieval SDG | NVIDIA NeMo Helix

NeMo Helix 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 Helix has no default Stage 0 chat model: you pass the names Inference Gateway serves. The Nemotron embed recipe uses nvidia/nemotron-3-ultra-550b-a55b for artifact extraction, Q&A generation, and quality judging. That choice is first-order for retrieval accuracy, because those queries and answers become the embedding training set. The published NVDocs dump was generated with that pipeline.

The examples below use nvidia/nemotron-3.5-lightning-30b-a3b so they can run on a typical already-deployed IGW model. Outside of examples, always use a larger model — for example nvidia/nemotron-3-ultra-550b-a55b when the provider serves it. Copying the examples as-is produces shallower queries and lower judge scores than the recipe, and a weaker fine-tune than NVDocs.

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.

Corpus and runtime

Corpus files must be UTF-8 text. Stage 0 walks directories recursively and, by default, reads .txt, .md, .text, and files with no extension. Use file_extensions to replace that list. Aim for documents of 200–2,000 tokens that represent the target domain. Start with at least 50–100 documents; 500 or more usually gives better domain coverage.

Stage 0 makes approximately four model calls per document for artifact extraction, Q&A generation, deduplication, and quality judging. As planning estimates, the recipe takes about 30 minutes for its roughly 70-file sample and tens of hours to 1–2 days for 10,000 or more files. The served model, provider rate limits, chunk and pair counts, and max_parallel_requests_for_gen can move those numbers substantially.

Stage 0 quality controls

Field	Default	Guidance
`file_extensions`	`[".txt", ".md", ".text", ""]`	Corpus files to process recursively.
`sentences_per_chunk`	`5`	About 100–150 tokens for typical prose. Match production chunk size as closely as practical; larger chunks cost more context.
`max_artifacts_per_type`	`2`	Entities, terms, concepts, and relationships retained per chunk. Raise it when rare identifiers matter.
`num_pairs`	`7`	Q&A pairs per chunk. Raise it for more query diversity on multifaceted content. `query_counts` and `reasoning_counts` must each sum to this value.
`multi_doc`	`false`	Bundle chunks for questions whose evidence spans sections or documents. Tune `bundle_size`, `bundle_strategy`, and `max_docs_per_bundle` when enabled.
`max_parallel_requests_for_gen`	Provider default (`null`)	Cap concurrent generation calls. Raise cautiously for throughput; lower it when the provider rate-limits or responses become unreliable.

Preview 1–5 records after changing these fields. Inspect identifier coverage, query depth, and judge scores before paying for a full corpus run.

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. quality_threshold belongs to this conversion step, not Stage 0: it defaults to 7.0 and removes Q&A pairs whose judge score is lower. Lower it only after checking why too few pairs survive.

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

nemo data-designer retrieval-prepare --spec '{
  "sdg_input": "default/stage0-out",
  "enable_mining": true,
  "quality_threshold": 7.0,
  "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,
    "quality_threshold": 7.0,
    "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.

Automodel dataset.training and retrieve-eval dataset take a fileset name with no path fragment. Copy training.jsonl and eval_beir/ to the root of a dataset fileset, and drop duplicate (query-id, corpus-id) rows from eval_beir/qrels/test.tsv before upload. Stage 1 unique-keys qrels when it writes them; unique-key again if you are uploading an older dump, because the BEIR loader still rejects duplicates. The copy-and-dedup cell in Embedding Model Customization does that procedure.

Automodel dataset.training — dataset discovery selects training.jsonl and ignores the non-JSONL train.json beside it. Stage 1 does not produce a training-validation file, so Automodel creates a deterministic holdout from training.jsonl using schedule.validation_split, which defaults to 0.1. If you add val*.jsonl, validation*.jsonl, or dev*.jsonl at the fileset root, Automodel uses that optional validation data instead.
nemo evals retrieve-eval dataset — the BEIR loader accepts either the BEIR directory or a fileset root containing eval_beir. It still rejects duplicate (query-id, corpus-id) rows in qrels/test.tsv.

The Automodel holdout is used for validation loss and checkpoint selection. The frozen eval_beir/ split is separate and remains the dataset for comparing base and tuned retrieval quality.

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_helix_plugin.client.client import NemoClient
from nemo_data_designer_plugin.sdk.resources import DataDesignerResource
from nemo_data_designer_plugin.jobs.retrieval_spec import RetrievalGenerateJobConfig
client = NemoClient(base_url="http://localhost:8080", workspace="default")
data_designer = DataDesignerResource(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.