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

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

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

```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:**

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

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

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

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

Number of negative passages per evaluation query.

Random seed used for source sampling and optional shuffling.

Whether to shuffle before applying `max_train_samples`.

Optional cap on the combined training rows.

Starting row for the training sample cap.

Whether to attach corpus query and passage instructions.

Whether to cycle through positive documents by epoch.

Whether image documents should also include their text.

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__)
```