Retrieval Dataset (Embedding 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.

The retrieval recipes use this dataset with the BiEncoderCollator. Example implementations are in examples/retrieval/bi_encoder/ and examples/retrieval/cross_encoder/.

What the Bi-Encoder Consumes

The dataset factory nemo_automodel.components.datasets.llm.make_retrieval_dataset returns a Hugging Face datasets.Dataset. At runtime, it transforms each raw record into the 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
query_instruction / passage_instruction: optional, used when use_dataset_instruction: true and the corpus provides instructions through metadata

Supported Input Formats

NeMo AutoModel supports two input schemas:

Corpus ID-Based JSON (Merlin/NeMo-Retriever Style)

This is the format used by NeMo retriever pipelines where documents live in a separate corpus and training examples reference documents by ID.

Training File Example (Single 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"]
    }
  ]
}

Corpus Requirements

Each corpus directory must contain a merlin_metadata.json file.

Minimal example:

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

pos_doc and neg_doc can be lists of {"id": ...} dicts or raw IDs (they are normalized internally).
cycle_positive_docs defaults to false, which always uses the first positive document. When a training record has multiple positive documents, set cycle_positive_docs: true to rotate through them deterministically by epoch order (epoch % len(pos_doc)).
If you set use_dataset_instruction: true, optional fields like query_instruction and passage_instruction in merlin_metadata.json are surfaced to the collator.

Inline-Text JSONL (No Corpus Required)

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

The format uses one JSON object per line:

{"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."}

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 (optionally image, nr_ocr) for multimodal use cases.
If corpus_id is not provided, it defaults to __inline__.
use_dataset_instruction: true has no effect for pure inline records (instructions come from corpus metadata).

Configure the Dataset and Collator in YAML

Use the dataset factory plus the bi-encoder collator:

dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  dataset:
    _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
    data_dir_list:
      - /abs/path/to/train.jsonl   # or train.json (corpus-id format)
    data_type: train
    n_passages: 5                 # 1 positive + 4 negatives
    do_shuffle: true
    use_dataset_instruction: false
  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

Requirements

pos_doc must be non-empty.
If training requests negatives (e.g., n_passages > 1), neg_doc must contain at least one document.