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`	Local Arrow corpus addressable by document id.
`NormalizedDataEntryConfig`	One normalized retrieval source and its optional sample cap.
`NormalizedRetrievalDatasetConfig`	Construction-time configuration for a normalized retrieval Arrow dataset.

Functions

Name	Description
`_build_transform`	-
`_coerce_bundle_root`	-
`_decode_image_bytes`	-
`_import_datasets`	-
`_load_arrow_shards`	-
`_load_corpus_dict`	-
`_load_metadata`	-
`_load_normalized_bundle_components`	-
`_merge_corpus_dicts`	-
`_parse_data_entries`	-
`_select_source_samples`	-
`make_normalized_retrieval_dataset`	Build a normalized portable retrieval dataset from local Arrow bundles.

Data

NormalizedDataEntry

NormalizedDataPath

_NORMALIZED_RETRIEVAL_FORMAT

_NORMALIZED_RETRIEVAL_FORMAT_VERSION

_VALID_MODEL_TYPES

logger

API

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

Bases: AbstractDataset

Local Arrow corpus addressable by document id.

docid2idx

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

Return all document IDs in sorted order.

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.

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]

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

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

Build the normalized retrieval dataset.

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
)

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

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

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

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

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]

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

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

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]

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

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

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 | NoneDefaults to None

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

model_type

strDefaults to 'bi_encoder'

Retrieval model type, either bi_encoder or cross_encoder.

data_type

strDefaults to 'train'

Dataset mode, either train or eval.

n_passages

intDefaults to 5

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

eval_negative_size

int | NoneDefaults to None

Number of negative passages per evaluation query.

seed

intDefaults to 42

Random seed used for source sampling and optional shuffling.

do_shuffle

boolDefaults to False

Whether to shuffle before applying max_train_samples.

max_train_samples

int | NoneDefaults to None

Optional cap on the combined training rows.

train_data_select_offset

intDefaults to 0

Starting row for the training sample cap.

use_dataset_instruction

boolDefaults to False

Whether to attach corpus query and passage instructions.

cycle_positive_docs

boolDefaults to False

Whether to cycle through positive documents by epoch.

use_text_in_document

boolDefaults to False

Whether image documents should also include their text.

data_dir_list

NormalizedDataPath | NoneDefaults to None

Legacy alias for data_path used by existing retrieval configs.

Returns: Any

A Hugging Face dataset with a lazy retrieval transform attached.

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

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

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

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

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

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