> 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.

# nemo_automodel.components.datasets.llm.retrieval_dataset

## Module Contents

### Classes

| Name                                                                                                         | Description                                                                            |
| ------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------- |
| [`AbstractDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)               | Interface for corpus datasets addressable by document id.                              |
| [`ColPaliDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-ColPaliDataset)                 | Load ColPali corpus documents from a dataset path.                                     |
| [`CorpusInfo`](#nemo_automodel-components-datasets-llm-retrieval_dataset-CorpusInfo)                         | Data structure to hold corpus metadata and dataset object together.                    |
| [`DocMatixDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-DocMatixDataset)               | Load DocMatix corpus documents from a dataset path.                                    |
| [`HFCorpusDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-HFCorpusDataset)               | Wraps an already-loaded HuggingFace Dataset as a corpus (in-memory, no local Parquet). |
| [`RetrievalDatasetConfig`](#nemo_automodel-components-datasets-llm-retrieval_dataset-RetrievalDatasetConfig) | Construction-time configuration for the retrieval dataset.                             |
| [`RetrievalTransform`](#nemo_automodel-components-datasets-llm-retrieval_dataset-RetrievalTransform)         | Stateful transform for retrieval datasets with epoch-based positive cycling.           |
| [`TextQADataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-TextQADataset)                   | Load TextQA corpus documents from a HuggingFace dataset path.                          |
| [`WikiSSNQDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-WikiSSNQDataset)               | Load Wiki-SS corpus documents from a dataset path.                                     |

### Functions

| Name                                                                                                                       | Description                                                                       |
| -------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`_cross_encoder_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_cross_encoder_transform_func) | Transform function to convert from raw format to cross-encoder training format.   |
| [`_list_hf_subsets`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_list_hf_subsets)                           | Discover all subset names in *repo\_id* by finding `dataset_metadata.json` files. |
| [`_load_hf_sources`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_load_hf_sources)                           | Load one or more `hf://` URIs and return `(Dataset, corpus_dict)`.                |
| [`_load_hf_subset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_load_hf_subset)                             | Load a single HF subset and return `(normalized_data_list, CorpusInfo)`.          |
| [`_normalize_data_entries`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_normalize_data_entries)             | Normalize a single source or list of sources into parsed entries.                 |
| [`_parse_data_entry`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_parse_data_entry)                         | Parse a data entry.                                                               |
| [`_parse_hf_uri`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_parse_hf_uri)                                 | Parse an `hf://` URI into `(repo_id, subset_or_none)`.                            |
| [`_sample_data_items`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_sample_data_items)                       | -                                                                                 |
| [`_transform_func`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_transform_func)                             | Transform function to convert from raw format to training format.                 |
| [`add_corpus`](#nemo_automodel-components-datasets-llm-retrieval_dataset-add_corpus)                                       | Add one or more corpus paths to a corpus dictionary.                              |
| [`load_corpus`](#nemo_automodel-components-datasets-llm-retrieval_dataset-load_corpus)                                     | Instantiate a corpus dataset from a path and optional metadata.                   |
| [`load_corpus_metadata`](#nemo_automodel-components-datasets-llm-retrieval_dataset-load_corpus_metadata)                   | Load Merlin corpus metadata from a corpus directory.                              |
| [`load_datasets`](#nemo_automodel-components-datasets-llm-retrieval_dataset-load_datasets)                                 | Load datasets from JSON files.                                                    |
| [`make_retrieval_dataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-make_retrieval_dataset)               | Load and return dataset in retrieval format for encoder training.                 |

### Data

[`DATASETS`](#nemo_automodel-components-datasets-llm-retrieval_dataset-DATASETS)

[`DataEntry`](#nemo_automodel-components-datasets-llm-retrieval_dataset-DataEntry)

[`EXAMPLE_TEMPLATE`](#nemo_automodel-components-datasets-llm-retrieval_dataset-EXAMPLE_TEMPLATE)

[`_HF_PREFIX`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_HF_PREFIX)

[`_OVERSAMPLING_WARNED_CORPORA`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_OVERSAMPLING_WARNED_CORPORA)

[`_VALID_MODEL_TYPES`](#nemo_automodel-components-datasets-llm-retrieval_dataset-_VALID_MODEL_TYPES)

[`args`](#nemo_automodel-components-datasets-llm-retrieval_dataset-args)

[`dataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-dataset)

[`example`](#nemo_automodel-components-datasets-llm-retrieval_dataset-example)

[`parser`](#nemo_automodel-components-datasets-llm-retrieval_dataset-parser)

### API

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset()
```

Abstract

Interface for corpus datasets addressable by document id.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset.get_all_ids()
```

abstract

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset.get_document_by_id(
    id
)
```

abstract

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.ColPaliDataset(
    path
)
```

**Bases:** [AbstractDataset](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)

Load ColPali corpus documents from a dataset path.

**`data`** `= load_dataset(path)['train']`

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.ColPaliDataset.get_all_ids()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.ColPaliDataset.get_document_by_id(
    id
)
```

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo(
    metadata: dict,
    corpus: nemo_automodel.components.datasets.llm.retrieval_dataset.AbstractDataset
)
```

Dataclass

Data structure to hold corpus metadata and dataset object together.
Provides easy access to both components with descriptive attribute names.

**`corpus`** `AbstractDataset`

---

**`corpus_id`** `str`

Get corpus ID from metadata

---

**`metadata`** `dict`

---

**`passage_instruction`** `str`

Get passage instruction from metadata

---

**`path`** `str`

Get corpus path from the corpus object

---

**`query_instruction`** `str`

Get query instruction from metadata

---

**`task_type`** `str`

Get task type from metadata

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo.get_all_ids()
```

Delegate to corpus for convenience

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo.get_document_by_id(
    doc_id: str
)
```

Delegate to corpus for convenience

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.DocMatixDataset(
    path
)
```

**Bases:** [AbstractDataset](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)

Load DocMatix corpus documents from a dataset path.

**`data`** `= load_dataset(path, 'images')['train']`

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.DocMatixDataset.get_all_ids()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.DocMatixDataset.get_document_by_id(
    id
)
```

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.HFCorpusDataset(
    hf_dataset: datasets.Dataset,
    path: str = ''
)
```

**Bases:** [AbstractDataset](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)

Wraps an already-loaded HuggingFace Dataset as a corpus (in-memory, no local Parquet).

**`_docid2idx`**

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.HFCorpusDataset.get_all_ids()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.HFCorpusDataset.get_document_by_id(
    id
)
```

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig(
    data_dir_list: list[nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry] | nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry | None = None,
    model_type: str = 'bi_encoder',
    data_type: str = 'train',
    n_passages: int = 5,
    eval_negative_size: int | None = None,
    seed: int = 42,
    do_shuffle: bool = False,
    max_train_samples: int | None = None,
    train_data_select_offset: int = 0,
    use_dataset_instruction: bool = False,
    cycle_positive_docs: bool = False,
    use_text_in_document: bool = False
)
```

Dataclass

Construction-time configuration for the retrieval dataset.

**`cycle_positive_docs`** `bool = False`

Whether to rotate through multiple positive documents by epoch during training.

---

**`data_dir_list`** `list[DataEntry] | DataEntry | None = None`

Path(s) to corpus JSON or inline JSONL files, `hf://` URIs, or dict entries with path and num\_samples.

---

**`data_type`** `str = 'train'`

Type of data (`train` or `eval`).

---

**`do_shuffle`** `bool = False`

Shuffle dataset rows before subset selection (only when `max_train_samples` is set).

---

**`eval_negative_size`** `int | None = None`

Number of negative documents for evaluation.

---

**`max_train_samples`** `int | None = None`

Maximum number of training samples to use.

---

**`model_type`** `str = 'bi_encoder'`

`bi_encoder` or `cross_encoder`.

---

**`n_passages`** `int = 5`

Number of passages (1 positive + n-1 negatives).

---

**`seed`** `int = 42`

Random seed for shuffling / sampling.

---

**`train_data_select_offset`** `int = 0`

Offset for selecting training samples.

---

**`use_dataset_instruction`** `bool = False`

Whether to use the instruction from the dataset's metadata.

---

**`use_text_in_document`** `bool = False`

Whether image documents should also include their text.

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig.build() -> datasets.Dataset
```

Build the retrieval :class:`~datasets.Dataset` from this :class:`RetrievalDatasetConfig`.

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalTransform(
    num_neg_docs: int,
    corpus_dict: dict,
    use_dataset_instruction: bool = False,
    model_type: str = 'bi_encoder',
    cycle_positive_docs: bool = False,
    use_text_in_document: bool = False
)
```

Stateful transform for retrieval datasets with epoch-based positive cycling.

This class encapsulates the transform state (epoch, corpus\_dict, etc.) and
provides a clean interface for updating the epoch without recreating the transform.

**`epoch`** `= 0`

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalTransform.__call__(
    examples
)
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalTransform.set_epoch(
    epoch: int
)
```

Update the epoch for positive document cycling.

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.TextQADataset(
    path
)
```

**Bases:** [AbstractDataset](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)

Load TextQA corpus documents from a HuggingFace dataset path.

**`data`** `= load_dataset(path)['train']`

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.TextQADataset.get_all_ids()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.TextQADataset.get_document_by_id(
    id
)
```

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset.WikiSSNQDataset(
    path
)
```

**Bases:** [AbstractDataset](#nemo_automodel-components-datasets-llm-retrieval_dataset-AbstractDataset)

Load Wiki-SS corpus documents from a dataset path.

**`data`** `= load_dataset(path)['train']`

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.WikiSSNQDataset.get_all_ids()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.WikiSSNQDataset.get_document_by_id(
    id
)
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._cross_encoder_transform_func(
    examples,
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False,
    epoch: int = 0,
    use_text_in_document: bool = False
)
```

Transform function to convert from raw format to cross-encoder training format.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._list_hf_subsets(
    repo_id: str
) -> typing.List[str]
```

Discover all subset names in *repo\_id* by finding `dataset_metadata.json` files.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._load_hf_sources(
    hf_entries: typing.List[typing.Tuple[int | None, str]],
    seed: int = 42
)
```

Load one or more `hf://` URIs and return `(Dataset, corpus_dict)`.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._load_hf_subset(
    repo_id: str,
    subset: str
)
```

Load a single HF subset and return `(normalized_data_list, CorpusInfo)`.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._normalize_data_entries(
    data_dir_list: typing.Union[typing.List[nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry], nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry]
) -> typing.List[typing.Tuple[int | None, str]]
```

Normalize a single source or list of sources into parsed entries.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._parse_data_entry(
    entry: nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry
) -> typing.Tuple[int | None, str]
```

Parse a data entry.

Supported forms:

* "path\_or\_hf\_uri": use all samples
* \{"path": "path\_or\_hf\_uri", "num\_samples": N}: sample N examples once from that source

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._parse_hf_uri(
    uri: str
)
```

Parse an `hf://` URI into `(repo_id, subset_or_none)`.

Examples::

"hf://nvidia/embed-nemotron-dataset-v1/FEVER"  -> ("nvidia/embed-nemotron-dataset-v1", "FEVER")
"hf://nvidia/embed-nemotron-dataset-v1"         -> ("nvidia/embed-nemotron-dataset-v1", None)

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._sample_data_items(
    data_items: typing.List[dict],
    num_samples: int | None,
    source: str,
    seed: int
) -> typing.List[dict]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._transform_func(
    examples,
    num_neg_docs,
    corpus_dict,
    use_dataset_instruction: bool = False,
    epoch: int = 0,
    use_text_in_document: bool = False
)
```

Transform function to convert from raw format to training format.

**Parameters:**

**`examples`**

Batch of examples with question, corpus\_id, pos\_doc, neg\_doc

---

**`num_neg_docs`**

Number of negative documents to use

---

**`corpus_dict`**

Dictionary mapping corpus\_id to corpus objects

---

**`use_dataset_instruction`** `bool` — default: False

Whether to use instruction from dataset's metadata

---

**`epoch`** `int` — default: 0

Current epoch for cycling through positive documents

---

**`use_text_in_document`** `bool` — default: False

Whether image documents should also include their text

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.add_corpus(
    qa_corpus_paths: typing.Union[dict, list],
    corpus_dict: dict
)
```

Add one or more corpus paths to a corpus dictionary.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.load_corpus(
    path,
    metadata: dict | None = None
)
```

Instantiate a corpus dataset from a path and optional metadata.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.load_corpus_metadata(
    path: str
)
```

Load Merlin corpus metadata from a corpus directory.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.load_datasets(
    data_dir_list: typing.Union[typing.List[nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry], nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry],
    concatenate: bool = True,
    seed: int = 42
)
```

Load datasets from JSON files.

Entries can be strings (use all samples) or dictionaries with path and optional
num\_samples fields (sample a fixed subset once while loading).

**Returns:**

Tuple of (dataset, corpus\_dict)

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.make_retrieval_dataset(
    data_dir_list: typing.Union[typing.List[nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry], nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry] = None,
    model_type: str = 'bi_encoder',
    data_type: str = 'train',
    n_passages: int = 5,
    eval_negative_size: int = None,
    seed: int = 42,
    do_shuffle: bool = False,
    max_train_samples: int = None,
    train_data_select_offset: int = 0,
    use_dataset_instruction: bool = False,
    cycle_positive_docs: bool = False,
    use_text_in_document: bool = False
)
```

Load and return dataset in retrieval format for encoder training.

Entries in *data\_dir\_list* can be local corpus JSON or inline JSONL file paths or `hf://` URIs
pointing to a Hugging Face dataset repository (for example,
`hf://nvidia/embed-nemotron-dataset-v1/SciFact`). A source can also be
provided as `{"path": path_or_uri, "num_samples": N}` to sample a fixed
subset once while loading. Uses `set_transform()` for lazy evaluation —
tokenization is handled by the collator.

**Parameters:**

**`data_dir_list`** `Union[List[DataEntry], DataEntry]` — default: None

Path(s) to corpus JSON or inline JSONL files, `hf://` URIs, or dictionary entries with path and
num\_samples.

---

**`model_type`** `str` — default: 'bi\_encoder'

"bi\_encoder" (default) or "cross\_encoder"

---

**`data_type`** `str` — default: 'train'

Type of data ("train" or "eval")

---

**`n_passages`** `int` — default: 5

Number of passages (1 positive + n-1 negatives)

---

**`eval_negative_size`** `int` — default: None

Number of negative documents for evaluation

---

**`seed`** `int` — default: 42

Random seed for reproducibility (for shuffling if needed)

---

**`do_shuffle`** `bool` — default: False

Shuffle dataset rows before subset selection. Only applied when
`max_train_samples` is set; otherwise iteration order is controlled by
the dataloader's sampler (e.g. `StatefulDistributedSampler`).

---

**`max_train_samples`** `int` — default: None

Maximum number of training samples to use

---

**`train_data_select_offset`** `int` — default: 0

Offset for selecting training samples

---

**`use_dataset_instruction`** `bool` — default: False

Whether to use instruction from dataset's metadata

---

**`cycle_positive_docs`** `bool` — default: False

Whether training should cycle through positive documents across epochs.
Defaults to `False` (always use the first positive document). Set to `True` only
when a query has multiple positive documents and you want to rotate through them by epoch.

---

**`use_text_in_document`** `bool` — default: False

Whether image documents should also include their text.

---

**Returns:**

A HuggingFace Dataset where each example is a dict with keys:

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.DATASETS = {'TextQADataset': TextQADataset, 'ColPaliDataset': ColPaliDataset, 'WikiSSNQData...
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.DataEntry = Union[str, dict[str, Any]]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.EXAMPLE_TEMPLATE = {'text': '', 'image': '', 'nr_ocr': ''}
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._HF_PREFIX = 'hf://'
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._OVERSAMPLING_WARNED_CORPORA: set[str] = set()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset._VALID_MODEL_TYPES = ('bi_encoder', 'cross_encoder')
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.args = parser.parse_args()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.dataset = make_retrieval_dataset(data_dir_list=(args.data_dir_list), data_type=(args.data_...
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.example = dataset[0]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset.parser = argparse.ArgumentParser(description='Load and transform dataset to retrieval for...
```