Retrieval Dataset

NeMo AutoModel supports retrieval model fine-tuning using a retrieval-style dataset: each training example is a query paired with one positive document and one or more negative documents.

This dataset is used by the retrieval recipes in examples/retrieval/bi_encoder/ and examples/retrieval/cross_encoder/, together with the retrieval collators. For an end-to-end training workflow, refer to Retrieval Fine-Tuning.

AutoModel supports three raw retrieval input formats: local corpus ID-based JSON, AutoModel-format Hugging Face datasets, and inline JSONL. It also supports prepared normalized Arrow bundles for large corpus-backed datasets. Regardless of how the source data is stored, AutoModel converts it into the following common format for training. Later sections explain how to use the raw formats with RetrievalDatasetConfig and normalized Arrow with NormalizedRetrievalDatasetConfig.

Training-Time Record Format

With model_type: bi_encoder, both dataset configs build a Hugging Face datasets.Dataset and transform each stored record into the same training-time schema:

• question: query string
• doc_text: list of document texts in the order [positive, negative_1, negative_2, ...]
• doc_image: list of images (or empty strings), aligned with doc_text
• doc_id: list of document identifiers aligned with doc_text. Corpus-backed and hf:// sources preserve document IDs. Pure inline JSONL produces empty IDs, so use a corpus-backed source or custom preprocessing when reliable same-document masking is required.
• query_instruction / passage_instruction: optional, used when use_dataset_instruction: true and the corpus provides instructions through metadata

With model_type: cross_encoder, the configs consume the same retrieval records but flatten each query with its positive and negative passages. CrossEncoderCollator then serializes each query-passage pair for reranking.

For each query, both configs select exactly one positive passage. They use the first item in pos_doc by default. Set cycle_positive_docs: true to rotate through multiple positives deterministically by epoch (epoch % len(pos_doc)). For guidance about expanding multi-positive data and preventing sibling positives from becoming negatives, refer to Convert Qrels and Multi-Positive Data.

Raw Input Formats

RetrievalDatasetConfig selects the loader for the three raw formats from the source URI or file extension. Use .jsonl for inline records, .json for local corpus-backed data, and hf:// for AutoModel-format Hugging Face sources.

The following table compares the three supported source types:

Corpus ID-Based JSON

Use this format when documents live in a separate corpus and training examples reference documents by ID.

The following example shows one training JSON file:

1
{
2
  "corpus": [
3
    { "path": "/abs/path/to/wiki_corpus" }
4
  ],
5
  "data": [
6
    {
7
      "question_id": "q_001",
8
      "question": "Explain transformers",
9
      "corpus_id": "wiki_corpus",
10
      "pos_doc": [{ "id": "d_123" }],
11
      "neg_doc": [{ "id": "d_456" }, "d_789"]
12
    }
13
  ]
14
}

For the text retrieval workflow in this guide, configure the corpus as a TextQADataset. The corpus directory must contain a merlin_metadata.json file and a Hugging Face-loadable train split with id and text columns. AutoModel calls datasets.load_dataset(<corpus path>)["train"], then resolves pos_doc and neg_doc IDs against that split.

Other registered corpus classes support image-oriented data and have different column requirements. Use the matching retrieval example when you configure one of those classes.

The following example is a minimal merlin_metadata.json file:

1
{ "class": "TextQADataset", "corpus_id": "wiki_corpus" }

The corresponding minimal local layout is:

retrieval-data/
  train.json
  wiki_corpus/
    merlin_metadata.json
    train.parquet   # or another load_dataset-compatible train split with id,text columns

The corpus_id in merlin_metadata.json must match the corpus_id in each training record. Relative corpus paths in train.json are resolved relative to the JSON file.

Hugging Face Sources

Direct hf:// loading expects the AutoModel retrieval schema, not arbitrary Hugging Face retrieval datasets. The URI format is:

hf://<org>/<repo>/<subset>

Each subset must provide:

• <subset>/dataset_metadata.json with corpus_id metadata and ids_only set to false or omitted
• a <subset>_corpus train split with id and text columns
• a <subset> train split with question and pos_doc, plus at least one neg_doc when n_passages > 1

Datasets in BEIR, DPR, MS MARCO, MIRACL, or other layouts need preprocessing unless they have already been converted to the AutoModel schema above.

Inline-Text JSONL (No Corpus Required)

This is convenient for custom fine-tuning pipelines where the documents are included inline.

The following JSONL example contains one record per line:

1
{"query":"Explain transformers","pos_doc":"Transformers are a type of neural network...","neg_doc":["RNNs are...","CNNs are..."]}
2
{"query":"What is Python?","pos_doc":["A programming language."],"neg_doc":"A snake."}

Raw Record Requirements

The following requirements apply to all three raw input types:

• pos_doc must be non-empty.
• Local JSON and JSONL records must include neg_doc; it may be empty only when n_passages: 1.
• An hf:// source may omit neg_doc only when n_passages: 1.
• With n_passages > 1, every record must provide at least one negative. If it provides fewer than n_passages - 1, AutoModel repeats the available negatives to fill the requested passage count.

Choose How to Load the Data

Load any of the three raw input types directly with RetrievalDatasetConfig. For large corpus-ID datasets, you can instead prepare a normalized Arrow bundle on CPU and load it with NormalizedRetrievalDatasetConfig.

Normalized Arrow is not a fourth raw input format. It is a prepared version of corpus ID-based JSON data that keeps the same query-to-document relationships in portable local Arrow shards.

Normalized Arrow for Large VL Datasets

A normalized Arrow bundle is a prepared version of corpus-backed retrieval data. Query rows keep their document ID references, while local Arrow shards store the referenced text and image documents. Training reads those prepared shards instead of loading the source corpora and building their Hugging Face dataset caches after GPUs have been allocated.

Use normalized Arrow for full-scale or image-heavy vision-language (VL) retrieval. For typical text-only retrieval, loading the source data directly with RetrievalDatasetConfig is usually simpler. Normalization can still help when a large text corpus has expensive startup or when you need a portable prepared dataset.

The CPU preparation tool currently accepts corpus ID-based JSON sources. For hf:// or inline JSONL sources, continue using RetrievalDatasetConfig unless you first convert them to corpus ID-based JSON.

Prepare the bundle before launching GPU training:

$
uv run python tools/retrieval/prepare_normalized_vl_retrieval_data.py \
>
  --config /path/to/retrieval_config_with_local_json_sources.yaml \
>
  --output-dir /path/to/normalized_retrieval

Then train with the normalized dataset config:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedRetrievalDatasetConfig
3
  data_path: /path/to/normalized_retrieval
4
  model_type: bi_encoder
5
  data_type: train
6
  n_passages: 5

For large datasets on a Slurm cluster, use the CPU normalized dataset preparation scripts to process sources in parallel. See the retrieval data preparation tools for local and Slurm commands, source selection, append and resume behavior, and the cache-warming alternative.

Configure the Dataset and Collator in YAML

Use RetrievalDatasetConfig with the appropriate collator for every supported raw source. The following corpus-backed example also works with an hf:// source:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
3
  model_type: bi_encoder
4
  data_dir_list:
5
    - /abs/path/to/train.json    # or hf://nvidia/embed-nemotron-dataset-v1/FEVER
6
  data_type: train
7
  n_passages: 5                 # 1 positive + 4 negatives
8
  use_dataset_instruction: false
9
10
dataloader:
11
  collate_fn:
12
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
13
    q_max_len: 512
14
    p_max_len: 512
15
    query_prefix: "query:"
16
    passage_prefix: "passage:"
17
    use_dataset_instruction: false
18
    pad_to_multiple_of: 8

For all supported source types, do_shuffle: true shuffles rows only when max_train_samples is set, before subsampling. Otherwise, the dataloader or distributed sampler controls training order.

For inline JSONL, keep the same public config and change the source path:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
3
  model_type: bi_encoder
4
  data_dir_list:
5
    - /abs/path/to/train.jsonl
6
  data_type: train
7
  n_passages: 5                 # 1 positive + 4 negatives
8
9
dataloader:
10
  collate_fn:
11
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
12
    q_max_len: 512
13
    p_max_len: 512
14
    query_prefix: "query:"
15
    passage_prefix: "passage:"
16
    pad_to_multiple_of: 8

For cross-encoder training, keep the same dataset config, set model_type: cross_encoder, and use CrossEncoderCollator arguments:

1
dataset:
2
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
3
  model_type: cross_encoder
4
  data_dir_list:
5
    - /abs/path/to/train.jsonl
6
  data_type: train
7
  n_passages: 5
8
9
dataloader:
10
  collate_fn:
11
    _target_: nemo_automodel.components.datasets.llm.CrossEncoderCollator
12
    rerank_max_length: 512
13
    prompt_template: "question:{query} \n \n passage:{passage}"
14
    pad_to_multiple_of: 8

Convert Qrels and Multi-Positive Data

The training dataset does not consume query relevance judgments (qrels) as a separate input. Convert qrels-style data into retrieval records before training:

1. Put every passage in a corpus split with stable id and text values.
2. For each query, write one or more records with question_id, question, corpus_id, pos_doc, and neg_doc. Use unique question_id values within each mining file because hard-negative mining writes results back by ID.
3. Put the canonical positive first in pos_doc. Set cycle_positive_docs: true to rotate through all positives by epoch, or expand the query into multiple records if every positive must be supervised in the same epoch.
4. For hard-negative mining, include every known positive document ID for the query in pos_doc. The miner excludes only IDs in the input record, not an external qrels file.
5. If you expand a query into multiple records, keep sibling-positive records out of the same in-batch-negative training batch or add a custom known-positive mask.

Preserve the complete qrels separately for offline full-corpus retrieval evaluation and use them to audit mined negatives before reusing the output for training.