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

Normalized portable retrieval dataset backed by local Arrow shards.

This format keeps the original retrieval data model:

* train rows contain query text plus positive/negative document ids;
* corpus shards store each referenced document/image once;
* training still resolves `doc_id -> document` through the retrieval transform.

It is intended as a portable alternative to the original corpus-id JSON plus
external Hugging Face corpus paths.

## Module Contents

### Classes

| Name                                                                                                                                        | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`NormalizedArrowCorpusDataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-NormalizedArrowCorpusDataset)         | Local Arrow corpus addressable by document id.                            |
| [`NormalizedDataEntryConfig`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-NormalizedDataEntryConfig)               | One normalized retrieval source and its optional sample cap.              |
| [`NormalizedRetrievalDatasetConfig`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-NormalizedRetrievalDatasetConfig) | Construction-time configuration for a normalized retrieval Arrow dataset. |

### Functions

| Name                                                                                                                                            | Description                                                             |
| ----------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_build_transform`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_build_transform)                                     | -                                                                       |
| [`_coerce_bundle_root`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_coerce_bundle_root)                               | -                                                                       |
| [`_decode_image_bytes`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_decode_image_bytes)                               | -                                                                       |
| [`_import_datasets`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_import_datasets)                                     | -                                                                       |
| [`_load_arrow_shards`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_load_arrow_shards)                                 | -                                                                       |
| [`_load_corpus_dict`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_load_corpus_dict)                                   | -                                                                       |
| [`_load_metadata`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_load_metadata)                                         | -                                                                       |
| [`_load_normalized_bundle_components`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_load_normalized_bundle_components) | -                                                                       |
| [`_merge_corpus_dicts`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_merge_corpus_dicts)                               | -                                                                       |
| [`_parse_data_entries`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_parse_data_entries)                               | -                                                                       |
| [`_select_source_samples`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_select_source_samples)                         | -                                                                       |
| [`make_normalized_retrieval_dataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-make_normalized_retrieval_dataset)   | Build a normalized portable retrieval dataset from local Arrow bundles. |

### Data

[`NormalizedDataEntry`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-NormalizedDataEntry)

[`NormalizedDataPath`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-NormalizedDataPath)

[`_NORMALIZED_RETRIEVAL_FORMAT`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_NORMALIZED_RETRIEVAL_FORMAT)

[`_NORMALIZED_RETRIEVAL_FORMAT_VERSION`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-_NORMALIZED_RETRIEVAL_FORMAT_VERSION)

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

[`logger`](#nemo_automodel-components-datasets-llm-retrieval_dataset_normalized-logger)

### API

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedArrowCorpusDataset(
    dataset: typing.Any,
    path: str
)
```

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

Local Arrow corpus addressable by document id.

**`docid2idx`**

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedArrowCorpusDataset.get_all_ids() -> list[str]
```

Return all document IDs in sorted order.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedArrowCorpusDataset.get_document_by_id(
    doc_id: str
) -> dict[str, typing.Any]
```

Return one normalized corpus document by ID.

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataEntryConfig
```

**Bases:** `typing.TypedDict`

One normalized retrieval source and its optional sample cap.

**`num_samples`** `int | None`

---

**`path`** `str | Path | PathLike[str]`

---

```python
class nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedRetrievalDatasetConfig(
    data_path: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath | 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,
    data_dir_list: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath | None = None
)
```

Dataclass

Construction-time configuration for a normalized retrieval Arrow dataset.

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

---

**`data_dir_list`** `NormalizedDataPath | None = None`

---

**`data_path`** `NormalizedDataPath | None = None`

---

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

---

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

---

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

---

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

---

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

---

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

---

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

---

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

---

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

---

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

---

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedRetrievalDatasetConfig.build() -> typing.Any
```

Build the normalized retrieval dataset.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._build_transform(
    corpus_dict: dict[str, nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo],
    model_type: str,
    data_type: str,
    n_passages: int,
    eval_negative_size: int | None,
    use_dataset_instruction: bool,
    cycle_positive_docs: bool,
    use_text_in_document: bool
)
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._coerce_bundle_root(
    data_path: str | pathlib.Path | os.PathLike[str]
) -> pathlib.Path
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._decode_image_bytes(
    image_bytes: bytes | None
) -> typing.Any
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._import_datasets()
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._load_arrow_shards(
    paths: typing.Sequence[pathlib.Path]
)
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._load_corpus_dict(
    bundle_root: pathlib.Path,
    metadata: dict[str, typing.Any]
) -> dict[str, nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._load_metadata(
    bundle_root: pathlib.Path
) -> dict[str, typing.Any]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._load_normalized_bundle_components(
    bundle_root: pathlib.Path,
    num_samples: int | None,
    seed: int
) -> tuple[typing.Any, dict[str, nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo]]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._merge_corpus_dicts(
    corpus_dicts: typing.Sequence[dict[str, nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo]]
) -> dict[str, nemo_automodel.components.datasets.llm.retrieval_dataset.CorpusInfo]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._parse_data_entries(
    data_path: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath
) -> list[tuple[pathlib.Path, int | None]]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._select_source_samples(
    dataset: typing.Any,
    num_samples: int | None,
    seed: int
) -> typing.Any
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.make_normalized_retrieval_dataset(
    data_path: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath | 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,
    data_dir_list: nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath | None = None
) -> typing.Any
```

Build a normalized portable retrieval dataset from local Arrow bundles.

**Parameters:**

**`data_path`** `NormalizedDataPath | None` — default: None

Bundle path or source entries. Use either this field or `data_dir_list`.

---

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

Retrieval model type, either `bi_encoder` or `cross_encoder`.

---

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

Dataset mode, either `train` or `eval`.

---

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

Number of passages per training query, including the positive passage.

---

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

Number of negative passages per evaluation query.

---

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

Random seed used for source sampling and optional shuffling.

---

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

Whether to shuffle before applying `max_train_samples`.

---

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

Optional cap on the combined training rows.

---

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

Starting row for the training sample cap.

---

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

Whether to attach corpus query and passage instructions.

---

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

Whether to cycle through positive documents by epoch.

---

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

Whether image documents should also include their text.

---

**`data_dir_list`** `NormalizedDataPath | None` — default: None

Legacy alias for `data_path` used by existing retrieval configs.

---

**Returns:** `Any`

A Hugging Face dataset with a lazy retrieval transform attached.

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataEntry = str | Path | os.PathLike[str] | NormalizedDataEntryConfig
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.NormalizedDataPath = NormalizedDataEntry | Sequence[NormalizedDataEntry]
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._NORMALIZED_RETRIEVAL_FORMAT = 'nemo_automodel_normalized_vl_retrieval_arrow'
```

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized._NORMALIZED_RETRIEVAL_FORMAT_VERSION = 1
```

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

```python
nemo_automodel.components.datasets.llm.retrieval_dataset_normalized.logger = logging.getLogger(__name__)
```