nemo_automodel.components.datasets.llm

Subpackages

nemo_automodel.components.datasets.llm.megatron

Submodules

Package Contents

Classes

Name	Description
`BiEncoderCollator`	Collator for encoder retrieval training.
`BiEncoderDistillCollator`	Collator for bi-encoder distillation with explicit positive/negative splits.
`ChatDataset`	Dataset for OpenAI-format tool-calling chat transcripts.
`ColumnMappedTextInstructionDataset`	Generic instruction-tuning dataset that maps arbitrary column names.
`ColumnMappedTextInstructionIterableDataset`	Streaming iterable variant that reuses the column-mapping/tokenization logic.
`CrossEncoderCollator`	Collate query-document pairs for cross-encoder reranking.
`DeltaLakeDataset`	HuggingFace datasets-compatible wrapper for Delta Lake tables.
`NanogptDataset`	Dataset class for NanoGPT Dataset.
`ProcessorMethodCollator`	Expose one method of a multimodal processor as a dataloader collator.

Functions

Name	Description
`is_delta_lake_path`	Check if a path refers to a Delta Lake table.
`make_agent_chat_dataset`	Load a multi-turn function-calling SFT dataset.
`make_normalized_retrieval_dataset`	Build a normalized portable retrieval dataset from local Arrow bundles.
`make_retrieval_dataset`	Load and return dataset in retrieval format for encoder training.
`make_squad_dataset`	Load and preprocess a SQuAD-style QA dataset for model fine-tuning.
`make_vision_collator_from_processor_method`	Turns a method of a processor into a collator function.
`make_xlam_dataset`	Load and preprocess the xLAM function-calling dataset to OpenAI messages

API

class nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator(
    tokenizer: transformers.PreTrainedTokenizerBase,
    q_max_len: int = 512,
    p_max_len: int = 512,
    query_prefix: str = '',
    passage_prefix: str = '',
    padding: typing.Union[bool, str, transformers.file_utils.PaddingStrategy] = True,
    pad_to_multiple_of: int = None,
    use_dataset_instruction: bool = False
)

Collator for encoder retrieval training.

This collator handles tokenization of queries and documents at batch time, which is more memory-efficient than pre-tokenization and allows for dynamic padding based on batch max length.

Based on EncoderCollator from nemo-retriever-research but adapted for Automodel.

nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator.__call__(
    batch: typing.List[typing.Dict[str, typing.Any]]
) -> typing.Dict[str, torch.Tensor]

Collate a batch of examples.

Parameters:

batch

List[Dict[str, Any]]

List of examples, each with ‘question’, ‘doc_text’, ‘doc_image’ keys

Returns: Dict[str, torch.Tensor]

Dictionary with:

nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator._convert_dict_to_list(
    input_dict: typing.Dict[str, typing.List]
) -> typing.List[typing.Dict[str, typing.Any]]

Convert dictionary of lists to list of dictionaries.

Example: Input: {‘a’: [1, 2], ‘b’: [3, 4]} Output: [{‘a’: 1, ‘b’: 3}, {‘a’: 2, ‘b’: 4}]

nemo_automodel.components.datasets.llm.retrieval_collator.BiEncoderCollator._merge_batch_dict(
    query_batch_dict: typing.Dict[str, typing.List],
    doc_batch_dict: typing.Dict[str, typing.List],
    train_n_passages: int
) -> typing.Dict[str, typing.List]

Merge query and document batches into a single dictionary.

Adapted from nemo-retriever-research/src/loaders/loader_utils.py

class nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator(
    tokenizer: transformers.PreTrainedTokenizerBase,
    q_max_len: int = 512,
    p_max_len: int = 512,
    query_prefix: str = '',
    passage_prefix: str = '',
    padding: transformers.file_utils.PaddingStrategy | str | bool = True,
    pad_to_multiple_of: int | None = None,
    use_dataset_instruction: bool = False,
    teacher_embeddings_cache: str | None = None,
    teacher_cache_lookup_with_prefixes: bool = False
)

Collator for bi-encoder distillation with explicit positive/negative splits.

_teacher_cache

CachedTeacherEmbeddings | None = None

teacher_cache

CachedTeacherEmbeddings | None

nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__call__(
    batch: typing.List[dict[str, typing.Any]]
) -> dict[str, torch.Tensor]

nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__getstate__() -> dict[str, typing.Any]

nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._add_cached_teacher_embeddings(
    out: dict[str, torch.Tensor],
    raw_queries: list[str],
    raw_pos_docs: list[str],
    raw_neg_lists: list[list[str]],
    query_texts: list[str],
    pos_docs: list[str],
    neg_lists: list[list[str]],
    max_negs: int
) -> None

nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_passage_prefix(
    text: str,
    instruction: str | None = None
) -> str

nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_query_prefix(
    text: str,
    instruction: str | None = None
) -> str

class nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset(
    path_or_dataset_id: typing.Union[str, typing.Sequence[str]],
    tokenizer,
    split: str | None = None,
    name: str | None = None,
    seq_length: int | None = None,
    padding: typing.Union[str, bool] = 'do_not_pad',
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    start_of_turn_token: str | None = None,
    chat_template: str | None = None,
    shuffle_seed: int | None = None,
    mask_reasoning_content: bool = False,
    mask_history: bool = False,
    unshifted: bool = False,
    skip_invalid_samples: bool = False,
    mask_generation_prompt: bool = False
)

Bases: Dataset

Dataset for OpenAI-format tool-calling chat transcripts.

Each row should contain a messages list in OpenAI chat format (role / content), potentially including tool calls and tool responses. Rows that instead carry a ShareGPT conversations list (from / value, as used by PerfectBlend and similar) are auto-converted, so no manual column rename is needed. The conversation is formatted via the tokenizer’s chat template to produce input_ids, labels, and attention_mask suitable for SFT.

dataset

pad_token_id

= _add_pad_token(self.tokenizer) or eos_token_id

nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset.__getitem__(
    idx: int
) -> typing.Dict[str, typing.List[int]]

nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset.__len__() -> int

nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset._keep_last_supervised_run(
    seq: typing.List[int],
    unsupervised_value: int
) -> None

staticmethod

In place, keep only the final contiguous supervised run; mask the rest.

Supervised positions are those != unsupervised_value (0 for loss_mask, -100 for labels). Used for mask_history: a multi-turn conversation has one supervised run per assistant turn separated by unsupervised user turns; this collapses it to the last turn so the supervised tokens form a single suffix.

class nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset(
    path_or_dataset_id: typing.Union[str, typing.List[str]],
    column_mapping: typing.Dict[str, str],
    tokenizer,
    split: str | None = 'train',
    name: str | None = None,
    answer_only_loss_mask: bool = True,
    seq_length: int | None = None,
    padding: typing.Union[str, bool] = 'do_not_pad',
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    limit_dataset_samples: int | None = None,
    use_hf_chat_template: bool = False
)

Bases: Dataset

Generic instruction-tuning dataset that maps arbitrary column names.

The class is intentionally lightweight: it simply loads the raw samples (either from HF or from local JSON/JSONL files) and remaps the columns so that downstream components can rely on a consistent field interface.

Optionally, if answer_only_loss_mask is requested, the dataset will also compute a loss_mask indicating which tokens should contribute to the loss (typically only those belonging to the assistant answer).

dataset

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset.__getitem__(
    idx
)

Returns the item at the given index.

Parameters:

idx

The index of the item to return.

Returns:

A dictionary with the mapped columns.

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.List[int]]]

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset.__len__() -> int

Returns the length of the dataset.

Returns: int

The length of the dataset.

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset._apply_tokenizer(
    sample: typing.Dict[str, str]
) -> typing.Dict[str, typing.List[int]]

Tokenize a mapped sample and compute auxiliary fields.

If the tokenizer is provided:

If the tokenizer supports a chat template, the dataset will be tokenized in a conversation style.
Otherwise, the dataset will be tokenized in a simple prompt-completion style.

Parameters:

sample

Dict[str, str]

A dictionary with the mapped columns.

Returns: Dict[str, List[int]]

A dictionary with the tokenized columns.

class nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset(
    path_or_dataset_id: typing.Union[str, typing.List[str]],
    column_mapping: typing.Dict[str, str],
    tokenizer,
    split: str | None = None,
    name: str | None = None,
    answer_only_loss_mask: bool = True,
    seq_length: int | None = None,
    padding: typing.Union[str, bool] = 'do_not_pad',
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    start_of_turn_token: str | None = None,
    limit_dataset_samples: int | None = None,
    repeat_on_exhaustion: bool = True,
    use_hf_chat_template: bool = False,
    delta_storage_options: typing.Dict[str, str] | None = None,
    delta_version: int | None = None,
    delta_sql_query: str | None = None
)

Bases: IterableDataset, ColumnMappedTextInstructionDataset

Streaming iterable variant that reuses the column-mapping/tokenization logic.

This wraps a Hugging Face streaming dataset (IterableDataset from datasets) or Delta Lake table and yields tokenized samples compatible with the non-streaming variant, while supporting sharding and epoch-setting for deterministic shuffles upstream.

Supports the following data sources:

HuggingFace Hub datasets
Local JSON/JSONL files
Delta Lake tables (via delta://, dbfs:/, or local directories with _delta_log)

_current_epoch_for_repeat

= 0

num_shards

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__getitem__(
    idx: int
) -> typing.Dict[str, typing.List[int]]

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.List[int]]]

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__len__() -> int

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.set_epoch(
    epoch: int
) -> None

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.shard(
    num_shards: int,
    index: int
)

nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.shuffle(
    buffer_size: int = 1000,
    seed: int | None = None
)

class nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator(
    rerank_max_length: int,
    args = (),
    prompt_template: str = 'question:{query} \n \n pas...,
    kwargs = {}
)

Bases: DataCollatorWithPadding

Collate query-document pairs for cross-encoder reranking.

nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator.__call__(
    features: typing.List[typing.Dict[str, typing.Any]]
) -> transformers.BatchEncoding

nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator.configure_tokenizer_for_export() -> None

Save the training format as a standard named tokenizer chat template.

The reranking template accepts flat query/document messages. It includes the tokenizer’s single-sequence special-token prefix and suffix because tokenizer chat templates bypass automatic special-token insertion. Sentence Transformers restores the special-token suffix after truncation.

Raises:

ValueError: If truncation is not right-sided, the prompt ends in fixed text, or the tokenizer does not use a portable special-token prefix/suffix with zero token type IDs.

class nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset(
    table_path: str,
    columns: list | None = None,
    storage_options: typing.Dict[str, str] | None = None,
    version: int | None = None,
    sql_query: str | None = None
)

HuggingFace datasets-compatible wrapper for Delta Lake tables.

This class provides better integration with the HuggingFace datasets library, supporting features like sharding, shuffling, and epoch setting for distributed training scenarios.

Parameters:

table_path

str

Path to the Delta Lake table.

columns

list | NoneDefaults to None

Optional list of column names to read.

storage_options

Dict[str, str] | NoneDefaults to None

Optional dict of storage options for cloud authentication.

version

int | NoneDefaults to None

Optional specific version of the Delta table to read.

_data_iterator

_epoch

int = 0

_shard_info

tuple | None = None

_shuffle_info

tuple | None = None

storage_options

= storage_options or {}

table_path

= _normalize_delta_path(table_path)

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__getitem__(
    idx: int
) -> typing.Dict[str, typing.Any]

Get a specific row by index.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.Any]]

Iterate over rows in the dataset.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__len__() -> int

Return the number of rows in the table.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.set_epoch(
    epoch: int
) -> None

Set the current epoch for deterministic shuffling.

Parameters:

epoch

int

The epoch number.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.shard(
    num_shards: int,
    index: int
) -> 'DeltaLakeDataset'

Shard the dataset for distributed processing.

Parameters:

num_shards

int

Total number of shards.

index

int

Index of this shard (0-based).

Returns: 'DeltaLakeDataset'

Self for method chaining.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.shuffle(
    buffer_size: int = 1000,
    seed: int | None = None
) -> 'DeltaLakeDataset'

Configure shuffling for the dataset.

Note: For streaming Delta Lake datasets, shuffling is performed on-the-fly using a shuffle buffer. The actual shuffling happens during iteration.

Parameters:

buffer_size

intDefaults to 1000

Size of the shuffle buffer.

seed

int | NoneDefaults to None

Random seed for reproducibility.

Returns: 'DeltaLakeDataset'

Self for method chaining.

nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.take(
    n: int
) -> 'DeltaLakeDataset'

Limit the dataset to the first n samples.

Parameters:

int

Number of samples to take.

Returns: 'DeltaLakeDataset'

A new DeltaLakeDataset limited to n samples.

class nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset(
    file_pattern: str | typing.Sequence[str],
    seq_len: int,
    bos_token: int | None = None,
    shuffle_files: bool = False,
    align_to_bos: bool = False
)

Bases: IterableDataset

Dataset class for NanoGPT Dataset.

A NanoGPT Dataset is a dataset that stores tokens in a binary file. The header contains:

256x4-byte header (magic number, version, num_tokens, dtype.itemsize)
And the tokens themselves.

Optionally, a corresponding .bos.idx file can be present alongside each .bin file containing precomputed BOS token positions for efficient alignment when align_to_bos=True. If the index file is not present, the dataset falls back to linear search for BOS tokens.

Parameters:

file_pattern

str | Sequence[str] Glob pattern (e.g. "data/fineweb_*_train_*.bin") or an explicit list of file paths.

seq_len

int Length of the training sample returned (not counting the next-token target). labels are simply inputs[1:].

shuffle_files

bool, default False Shuffle the order of shards each epoch/iteration.

align_to_bos

bool, default False Ensure that every slice starts with bos_token. When enabled, the dataset searches forward from the current position until it finds the next BOS token and starts there. Uses .bos.idx files when available for efficient search, falls back to linear search otherwise. Requires bos_token to be provided.

bos_token

int, optional, default None. Token ID marking beginning-of-document.

files

List[str] = sorted(glob.glob(str(file_pattern)))

seq_len

= int(seq_len)

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__getitem__(
    index: int
)

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__iter__() -> typing.Iterator[dict]

Iterate over training samples from the dataset.

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__len__() -> int

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset._get_file_iterator(
    worker_files: typing.List[str],
    rng: random.Random,
    split_single_file: bool,
    file_start_pos: int,
    file_end_pos: int
) -> typing.Iterator[dict]

Generate training samples from all assigned files, handling infinite iteration.

Parameters:

worker_files

List[str]

List of files assigned to this worker

rng

random.Random

Random number generator for shuffling

split_single_file

bool

Whether we’re splitting a single file among workers

file_start_pos

int

Starting position in file (for single file splitting)

file_end_pos

int

Ending position in file (for single file splitting)

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset._process_file_tokens(
    file: str,
    split_single_file: bool,
    file_start_pos: int,
    file_end_pos: int
) -> typing.Iterator[dict]

Process tokens from a single file and yield training samples.

Parameters:

file

str

Path to the .bin file to process

split_single_file

bool

Whether we’re splitting a single file among workers

file_start_pos

int

Starting position in the file (for single file splitting)

file_end_pos

int

Ending position in the file (for single file splitting)

nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset._setup_worker_context(
    files,
    shuffle
) -> tuple[typing.List[str], random.Random, bool, int, int]

Set up worker-specific context including file assignment and splitting parameters.

Returns: tuple[List[str], random.Random, bool, int, int]

Tuple of (worker_files, rng, split_single_file, file_start_pos, file_end_pos)

class nemo_automodel.components.datasets.llm.retrieval_collator.ProcessorMethodCollator(
    tokenizer: transformers.ProcessorMixin,
    collator_fn_name: str
)

Expose one method of a multimodal processor as a dataloader collator.

collate_fn

passage_prefix

str

Return the passage prefix owned by the underlying processor.

query_prefix

str

Return the query prefix owned by the underlying processor.

nemo_automodel.components.datasets.llm.retrieval_collator.ProcessorMethodCollator.__call__(
    batch: list[dict[str, object]]
) -> dict[str, object]

Collate retrieval examples with the resolved processor method.

Parameters:

batch

list[dict[str, object]]

Retrieval examples for one local batch.

Returns: dict[str, object]

Processor-produced tensor batch.

nemo_automodel.components.datasets.llm.delta_lake_dataset.is_delta_lake_path(
    path: str
) -> bool

Check if a path refers to a Delta Lake table.

A path is considered a Delta Lake path if:

It starts with “delta://” protocol prefix
It’s a local directory containing a “_delta_log” subdirectory
It starts with “dbfs:/” (Databricks file system)

Parameters:

path

str

The path to check.

Returns: bool

True if the path is a Delta Lake table, False otherwise.

nemo_automodel.components.datasets.llm.agent_chat.make_agent_chat_dataset(
    tokenizer,
    dataset_name: str | None = None,
    path: typing.Union[str, typing.List[str]] | None = None,
    split: str = 'train',
    revision: str | None = None,
    filter_no_assistant: bool = False,
    filter_overlong: bool = False,
    trim_incomplete_last_turn: bool = False,
    chat_template: str | None = None,
    seq_length: int | None = None,
    limit_dataset_samples: int | None = None,
    padding: typing.Union[str, bool] = False,
    truncation: typing.Union[str, bool] = False,
    mask_reasoning_content: bool = False,
    train_on_last_turn_only: bool = False,
    drop_history_reasoning_content: bool = False,
    truncate_history: bool = False,
    mask_generation_prompt: bool = False
) -> nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset

Load a multi-turn function-calling SFT dataset.

Exactly one of dataset_name (HuggingFace Hub id) or path (local JSON/JSONL file or list of files) must be provided. The loaded examples are lazily rendered through the tokenizer’s chat template; tool/user tokens are excluded from the loss via answer_only_loss_mask=True.

Parameters:

tokenizer

HuggingFace tokenizer with a chat template.

dataset_name

str | NoneDefaults to None

HF Hub dataset id, e.g. llamafactory/glaive_toolcall_en.

path

Union[str, List[str]] | NoneDefaults to None

Local JSON/JSONL file path or list of paths.

split

strDefaults to 'train'

Dataset split (only used with dataset_name).

revision

str | NoneDefaults to None

Optional Hugging Face dataset revision. Agent trace repositories may expose training-ready rows under refs/convert/parquet.

filter_no_assistant

boolDefaults to False

If True, discard incomplete traces that contain no assistant message and therefore have no supervised target.

filter_overlong

boolDefaults to False

If True, discard traces that still exceed seq_length after whole-exchange history truncation. This avoids token-level truncation changing chat-template prefixes while answer-only loss masks are constructed.

trim_incomplete_last_turn

boolDefaults to False

If True, discard messages after the final assistant message so an unfinished user turn cannot replace the supervised target during history truncation.

chat_template

str | NoneDefaults to None

Optional Jinja template string or file path overriding tokenizer.chat_template.

seq_length

int | NoneDefaults to None

Optional max sequence length for the tokenizer.

limit_dataset_samples

int | NoneDefaults to None

If set, keep only the first N examples.

padding

Union[str, bool]Defaults to False

Padding strategy forwarded to the tokenizer.

truncation

Union[str, bool]Defaults to False

Truncation strategy forwarded to the tokenizer.

mask_reasoning_content

boolDefaults to False

If True, exclude assistant reasoning_content (thinking) tokens from the loss while still rendering them into the prompt. Requires a chat template that emits reasoning_content. Defaults to False, which trains on reasoning tokens like any other assistant content.

mask_generation_prompt

boolDefaults to False

If True, exclude from the loss the tokens of each assistant turn that the chat template’s generation prompt supplies at inference: the role header and any template-inserted empty reasoning block (for example the <think></think> Nemotron templates emit for non-thinking turns). The model never generates those tokens, so supervising them only reinforces template boilerplate. Detected per template by rendering the generation prompt, so no tag strings are hardcoded. Defaults to False.

train_on_last_turn_only

boolDefaults to False

If True, supervise only the final assistant turn of each dialogue (mask_history); all earlier assistant turns are excluded from the loss. Defaults to False, which supervises every assistant turn.

drop_history_reasoning_content

boolDefaults to False

If True, strip reasoning_content from all but the final assistant turn so historical thinking is not rendered into the prompt (matching inference, where prior-turn thinking is not visible). Orthogonal to mask_reasoning_content: this controls whether history thinking appears in the prompt at all, the latter controls whether rendered thinking contributes to loss. Defaults to False, which keeps every turn’s reasoning_content.

truncate_history

boolDefaults to False

If True and seq_length is set, drop the oldest conversation exchanges (keeping any leading system message, the tool definitions, and the final exchange) until the dialogue fits seq_length. Unlike token-level truncation, which clips from the tokenizer side and usually drops the final assistant answer, this preserves the supervised target. Defaults to False.

Returns: LazyMappedDataset

A LazyMappedDataset yielding dicts with input_ids, labels

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.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]Defaults to None

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

model_type

strDefaults to 'bi_encoder'

“bi_encoder” (default) or “cross_encoder”

data_type

strDefaults to 'train'

Type of data (“train” or “eval”)

n_passages

intDefaults to 5

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

eval_negative_size

intDefaults to None

Number of negative documents for evaluation

seed

intDefaults to 42

Random seed for reproducibility (for shuffling if needed)

do_shuffle

boolDefaults to 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

intDefaults to None

Maximum number of training samples to use

train_data_select_offset

intDefaults to 0

Offset for selecting training samples

use_dataset_instruction

boolDefaults to False

Whether to use instruction from dataset’s metadata

cycle_positive_docs

boolDefaults to 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

boolDefaults to False

Whether image documents should also include their text.

Returns:

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

Direct hf:// loading currently supports HF datasets that already follow the Automodel retrieval schema (corpus-id based layout used by nvidia/embed-nemotron-dataset-v1 subsets such as FEVER and SyntheticClassificationData). For other HF dataset formats, implement a custom adapter/preprocessor before calling this loader.

Tokenization should be handled by a collator (e.g., BiEncoderCollator) which is more efficient for batch padding and supports dynamic processing.

nemo_automodel.components.datasets.llm.squad.make_squad_dataset(
    tokenizer,
    seq_length = None,
    limit_dataset_samples = None,
    fp8 = False,
    split = 'train',
    dataset_name = 'squad',
    padding = False,
    truncation = False,
    chat_template: str | None = None
)

Load and preprocess a SQuAD-style QA dataset for model fine-tuning.

This function retrieves the specified split of the SQuAD dataset, applies either a simple prompt–completion format or a chat‐template format (if tokenizer.chat_template is set), tokenizes each example, constructs input_ids and labels, and optionally pads all sequences to a fixed length.

Parameters:

tokenizer

A Hugging Face tokenizer with attributes eos_token_id, optional bos_id, optional eos_id, and optionally chat_template/apply_chat_template.

seq_length

intDefaults to None

If set, pad/truncate each example to this length.

limit_dataset_samples

intDefaults to None

If set, limit the number of examples loaded from the split.

fp8

boolDefaults to False

Flag for future use (e.g., mixed precision). Currently unused.

split

strDefaults to 'train'

Which split of the dataset to load (e.g. ‘train’, ‘validation’).

dataset_name

strDefaults to 'squad'

Identifier for the Hugging Face dataset (default “rajpurkar/squad”).

padding

Optional[str | bool]Defaults to False

Optional padding strategy.

truncation

Optional[str | bool]Defaults to False

Optional truncation strategy.

chat_template

str | NoneDefaults to None

Optional Jinja template string or path overriding tokenizer.chat_template.

Returns:

A Hugginggth Face Dataset where each example is a dict with keys:

nemo_automodel.components.datasets.llm.retrieval_collator.make_vision_collator_from_processor_method(
    tokenizer: transformers.ProcessorMixin,
    collator_fn_name: str
) -> typing.Callable[[list[dict[str, object]]], dict[str, object]]

Turns a method of a processor into a collator function.

Parameters:

tokenizer

ProcessorMixin

The processor instance.

collator_fn_name

str

The name of the processor method to turn into a collator function.

Returns: Callable[[list[dict[str, object]]], dict[str, object]]

A collator for vision/multimodal retrieval datasets.

nemo_automodel.components.datasets.llm.xlam.make_xlam_dataset(
    tokenizer,
    seq_length = None,
    limit_dataset_samples = None,
    fp8 = False,
    split = 'train',
    dataset_name = 'Salesforce/xlam-function-c...,
    padding = False,
    truncation = False
)

Load and preprocess the xLAM function-calling dataset to OpenAI messages compatible with the bulbasaur chat template (tool-calling aware).

Each example is formatted as:

user: the natural language query
assistant: emits tool_calls with serialized arguments
tools: OpenAI function schema derived from the dataset tool specs