> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # Retrieval Dataset > Prepare raw or normalized datasets for bi-encoder and cross-encoder fine-tuning. 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](/recipes-e2e-examples/retrieval-finetune). 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](#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: | Source | Query field | Required structure | Typical use | | ------------------------ | --------------------- | -------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- | | Corpus ID JSON | `question` | `question_id`, `corpus_id`, `pos_doc`, and `neg_doc` IDs resolved through a local corpus | Local corpus-backed training and hard-negative mining | | `hf://` AutoModel schema | `question` | `pos_doc`, a companion Hugging Face corpus split, and at least one `neg_doc` when `n_passages > 1` | Datasets hosted on Hugging Face that already use the AutoModel retrieval schema | | Inline JSONL | `query` or `question` | Inline text in `pos_doc` and `neg_doc` | Small custom text datasets | ### 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: ```json { "corpus": [ { "path": "/abs/path/to/wiki_corpus" } ], "data": [ { "question_id": "q_001", "question": "Explain transformers", "corpus_id": "wiki_corpus", "pos_doc": [{ "id": "d_123" }], "neg_doc": [{ "id": "d_456" }, "d_789"] } ] } ``` 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()["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: ```json { "class": "TextQADataset", "corpus_id": "wiki_corpus" } ``` The corresponding minimal local layout is: ```text 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. The following behavior applies to corpus-backed records: * `pos_doc` and `neg_doc` can be lists of `{"id": ...}` dicts or raw IDs (they are normalized internally). * To train with corpus instructions, set `use_dataset_instruction: true` on both the dataset and the bi-encoder collator. The dataset surfaces `query_instruction` and `passage_instruction` from `merlin_metadata.json`. The collator prepends them before tokenization. ### Hugging Face Sources Direct `hf://` loading expects the AutoModel retrieval schema, not arbitrary Hugging Face retrieval datasets. The URI format is: ```text hf://// ``` Each subset must provide: * `/dataset_metadata.json` with `corpus_id` metadata and `ids_only` set to `false` or omitted * a `_corpus` train split with `id` and `text` columns * a `` 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: ```json {"query":"Explain transformers","pos_doc":"Transformers are a type of neural network...","neg_doc":["RNNs are...","CNNs are..."]} {"query":"What is Python?","pos_doc":["A programming language."],"neg_doc":"A snake."} ``` The following behavior applies to inline records: * `query` is accepted (`question` is also accepted as an alias). * `pos_doc` and `neg_doc` can be either: * strings (interpreted as document text), or * lists of strings, or * dicts with at least `text`. * The current LLM retrieval collators tokenize text only. Do not rely on inline `image` or OCR fields unless you add a custom preprocessing and collator path. * If `corpus_id` is not provided, it defaults to `__inline__`. * Keep `use_dataset_instruction: false` for pure inline records. Dataset instructions come from corpus metadata, which inline JSONL does not provide. ## 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. `n_passages: 1` is supported, but it is not the usual training setup. The standard bi-encoder and cross-encoder recipes need at least one negative candidate for meaningful contrastive or reranking supervision, unless you add a custom negative strategy such as qrels-aware in-batch negatives. ## 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. | Data | Dataset config | Use when | | ---------------------------------------- | ---------------------------------- | ------------------------------------------------------------------------------------------------ | | Corpus ID JSON, `hf://`, or inline JSONL | `RetrievalDatasetConfig` | Direct loading for typical text retrieval, Hugging Face-hosted data, and small runs | | Prepared normalized Arrow bundle | `NormalizedRetrievalDatasetConfig` | Full-scale or image-heavy VL retrieval, expensive corpus startup, or a portable prepared dataset | ## 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: ```bash 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: ```yaml dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedRetrievalDatasetConfig data_path: /path/to/normalized_retrieval model_type: bi_encoder data_type: train 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](https://github.com/NVIDIA-NeMo/Automodel/blob/main/tools/retrieval/README.md) 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: ```yaml dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: bi_encoder data_dir_list: - /abs/path/to/train.json # or hf://nvidia/embed-nemotron-dataset-v1/FEVER data_type: train n_passages: 5 # 1 positive + 4 negatives use_dataset_instruction: false dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator q_max_len: 512 p_max_len: 512 query_prefix: "query:" passage_prefix: "passage:" use_dataset_instruction: false 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: ```yaml dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: bi_encoder data_dir_list: - /abs/path/to/train.jsonl data_type: train n_passages: 5 # 1 positive + 4 negatives dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator q_max_len: 512 p_max_len: 512 query_prefix: "query:" passage_prefix: "passage:" pad_to_multiple_of: 8 ``` For cross-encoder training, keep the same dataset config, set `model_type: cross_encoder`, and use `CrossEncoderCollator` arguments: ```yaml dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: cross_encoder data_dir_list: - /abs/path/to/train.jsonl data_type: train n_passages: 5 dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.CrossEncoderCollator rerank_max_length: 512 prompt_template: "question:{query} \n \n passage:{passage}" 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.