> This page is for version 0.5.0 · 26.06. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > 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 ## Subpackages * **[`nemo_automodel.components.datasets.llm.megatron`](/nemo-automodel/nemo_automodel/components/datasets/llm/megatron)** ## Submodules * **[`nemo_automodel.components.datasets.llm.agent_chat`](/nemo-automodel/nemo_automodel/components/datasets/llm/agent_chat)** * **[`nemo_automodel.components.datasets.llm.chat_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/chat_dataset)** * **[`nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/column_mapped_text_instruction_dataset)** * **[`nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/column_mapped_text_instruction_iterable_dataset)** * **[`nemo_automodel.components.datasets.llm.delta_lake_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/delta_lake_dataset)** * **[`nemo_automodel.components.datasets.llm.dspark_cache`](/nemo-automodel/nemo_automodel/components/datasets/llm/dspark_cache)** * **[`nemo_automodel.components.datasets.llm.eagle3`](/nemo-automodel/nemo_automodel/components/datasets/llm/eagle3)** * **[`nemo_automodel.components.datasets.llm.eagle3_cache`](/nemo-automodel/nemo_automodel/components/datasets/llm/eagle3_cache)** * **[`nemo_automodel.components.datasets.llm.formatting_utils`](/nemo-automodel/nemo_automodel/components/datasets/llm/formatting_utils)** * **[`nemo_automodel.components.datasets.llm.hellaswag`](/nemo-automodel/nemo_automodel/components/datasets/llm/hellaswag)** * **[`nemo_automodel.components.datasets.llm.length_grouped_sampler`](/nemo-automodel/nemo_automodel/components/datasets/llm/length_grouped_sampler)** * **[`nemo_automodel.components.datasets.llm.megatron_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/megatron_dataset)** * **[`nemo_automodel.components.datasets.llm.mock`](/nemo-automodel/nemo_automodel/components/datasets/llm/mock)** * **[`nemo_automodel.components.datasets.llm.mock_iterable_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/mock_iterable_dataset)** * **[`nemo_automodel.components.datasets.llm.mock_packed`](/nemo-automodel/nemo_automodel/components/datasets/llm/mock_packed)** * **[`nemo_automodel.components.datasets.llm.mock_prefix_tree`](/nemo-automodel/nemo_automodel/components/datasets/llm/mock_prefix_tree)** * **[`nemo_automodel.components.datasets.llm.mock_seq_cls`](/nemo-automodel/nemo_automodel/components/datasets/llm/mock_seq_cls)** * **[`nemo_automodel.components.datasets.llm.nanogpt_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/nanogpt_dataset)** * **[`nemo_automodel.components.datasets.llm.neat_packing`](/nemo-automodel/nemo_automodel/components/datasets/llm/neat_packing)** * **[`nemo_automodel.components.datasets.llm.offline_cache`](/nemo-automodel/nemo_automodel/components/datasets/llm/offline_cache)** * **[`nemo_automodel.components.datasets.llm.packed_sequence`](/nemo-automodel/nemo_automodel/components/datasets/llm/packed_sequence)** * **[`nemo_automodel.components.datasets.llm.prefix_tree`](/nemo-automodel/nemo_automodel/components/datasets/llm/prefix_tree)** * **[`nemo_automodel.components.datasets.llm.retrieval_collator`](/nemo-automodel/nemo_automodel/components/datasets/llm/retrieval_collator)** * **[`nemo_automodel.components.datasets.llm.retrieval_dataset`](/nemo-automodel/nemo_automodel/components/datasets/llm/retrieval_dataset)** * **[`nemo_automodel.components.datasets.llm.retrieval_dataset_inline`](/nemo-automodel/nemo_automodel/components/datasets/llm/retrieval_dataset_inline)** * **[`nemo_automodel.components.datasets.llm.retrieval_dataset_normalized`](/nemo-automodel/nemo_automodel/components/datasets/llm/retrieval_dataset_normalized)** * **[`nemo_automodel.components.datasets.llm.retrieval_distill_collator`](/nemo-automodel/nemo_automodel/components/datasets/llm/retrieval_distill_collator)** * **[`nemo_automodel.components.datasets.llm.seq2seq`](/nemo-automodel/nemo_automodel/components/datasets/llm/seq2seq)** * **[`nemo_automodel.components.datasets.llm.seq_cls`](/nemo-automodel/nemo_automodel/components/datasets/llm/seq_cls)** * **[`nemo_automodel.components.datasets.llm.squad`](/nemo-automodel/nemo_automodel/components/datasets/llm/squad)** * **[`nemo_automodel.components.datasets.llm.xlam`](/nemo-automodel/nemo_automodel/components/datasets/llm/xlam)** ## Package Contents ### Classes | Name | Description | | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- | | [`BiEncoderCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-BiEncoderCollator) | Collator for encoder retrieval training. | | [`BiEncoderDistillCollator`](#nemo_automodel-components-datasets-llm-retrieval_distill_collator-BiEncoderDistillCollator) | Collator for bi-encoder distillation with explicit positive/negative splits. | | [`ChatDataset`](#nemo_automodel-components-datasets-llm-chat_dataset-ChatDataset) | Dataset for OpenAI-format tool-calling chat transcripts. | | [`ColumnMappedTextInstructionDataset`](#nemo_automodel-components-datasets-llm-column_mapped_text_instruction_dataset-ColumnMappedTextInstructionDataset) | Generic instruction-tuning dataset that maps arbitrary column names. | | [`ColumnMappedTextInstructionIterableDataset`](#nemo_automodel-components-datasets-llm-column_mapped_text_instruction_iterable_dataset-ColumnMappedTextInstructionIterableDataset) | Streaming iterable variant that reuses the column-mapping/tokenization logic. | | [`CrossEncoderCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-CrossEncoderCollator) | Collate query-document pairs for cross-encoder reranking. | | [`DeltaLakeDataset`](#nemo_automodel-components-datasets-llm-delta_lake_dataset-DeltaLakeDataset) | HuggingFace datasets-compatible wrapper for Delta Lake tables. | | [`NanogptDataset`](#nemo_automodel-components-datasets-llm-nanogpt_dataset-NanogptDataset) | Dataset class for NanoGPT Dataset. | | [`ProcessorMethodCollator`](#nemo_automodel-components-datasets-llm-retrieval_collator-ProcessorMethodCollator) | Expose one method of a multimodal processor as a dataloader collator. | ### Functions | Name | Description | | ----------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ | | [`is_delta_lake_path`](#nemo_automodel-components-datasets-llm-delta_lake_dataset-is_delta_lake_path) | Check if a path refers to a Delta Lake table. | | [`make_agent_chat_dataset`](#nemo_automodel-components-datasets-llm-agent_chat-make_agent_chat_dataset) | Load a multi-turn function-calling SFT dataset. | | [`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. | | [`make_retrieval_dataset`](#nemo_automodel-components-datasets-llm-retrieval_dataset-make_retrieval_dataset) | Load and return dataset in retrieval format for encoder training. | | [`make_squad_dataset`](#nemo_automodel-components-datasets-llm-squad-make_squad_dataset) | Load and preprocess a SQuAD-style QA dataset for model fine-tuning. | | [`make_vision_collator_from_processor_method`](#nemo_automodel-components-datasets-llm-retrieval_collator-make_vision_collator_from_processor_method) | Turns a method of a processor into a collator function. | | [`make_xlam_dataset`](#nemo_automodel-components-datasets-llm-xlam-make_xlam_dataset) | Load and preprocess the xLAM function-calling dataset to OpenAI messages | ### API ```python 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. ```python 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: ```python 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}] ```python 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 ```python 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` --- ```python nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__call__( batch: typing.List[dict[str, typing.Any]] ) -> dict[str, torch.Tensor] ``` ```python nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator.__getstate__() -> dict[str, typing.Any] ``` ```python 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 ``` ```python nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_passage_prefix( text: str, instruction: str | None = None ) -> str ``` ```python nemo_automodel.components.datasets.llm.retrieval_distill_collator.BiEncoderDistillCollator._apply_query_prefix( text: str, instruction: str | None = None ) -> str ``` ```python 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` --- ```python nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset.__getitem__( idx: int ) -> typing.Dict[str, typing.List[int]] ``` ```python nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset.__len__() -> int ``` ```python 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. ```python 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`** --- ```python 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. ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_dataset.ColumnMappedTextInstructionDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.List[int]]] ``` ```python 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. ```python 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. ```python 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](/nemo-automodel/nemo_automodel/components/datasets/llm/column_mapped_text_instruction_dataset#nemo_automodel-components-datasets-llm-column_mapped_text_instruction_dataset-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`** --- ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__getitem__( idx: int ) -> typing.Dict[str, typing.List[int]] ``` ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.List[int]]] ``` ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.__len__() -> int ``` ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.set_epoch( epoch: int ) -> None ``` ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.shard( num_shards: int, index: int ) ``` ```python nemo_automodel.components.datasets.llm.column_mapped_text_instruction_iterable_dataset.ColumnMappedTextInstructionIterableDataset.shuffle( buffer_size: int = 1000, seed: int | None = None ) ``` ```python 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. ```python nemo_automodel.components.datasets.llm.retrieval_collator.CrossEncoderCollator.__call__( features: typing.List[typing.Dict[str, typing.Any]] ) -> transformers.BatchEncoding ``` ```python 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. ```python 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 | None` — default: None Optional list of column names to read. --- **`storage_options`** `Dict[str, str] | None` — default: None Optional dict of storage options for cloud authentication. --- **`version`** `int | None` — default: 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)` --- ```python nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__getitem__( idx: int ) -> typing.Dict[str, typing.Any] ``` Get a specific row by index. ```python nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__iter__() -> typing.Iterator[typing.Dict[str, typing.Any]] ``` Iterate over rows in the dataset. ```python nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.__len__() -> int ``` Return the number of rows in the table. ```python 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. --- ```python 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. ```python 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`** `int` — default: 1000 Size of the shuffle buffer. --- **`seed`** `int | None` — default: None Random seed for reproducibility. --- **Returns:** `'DeltaLakeDataset'` Self for method chaining. ```python nemo_automodel.components.datasets.llm.delta_lake_dataset.DeltaLakeDataset.take( n: int ) -> 'DeltaLakeDataset' ``` Limit the dataset to the first n samples. **Parameters:** **`n`** `int` Number of samples to take. --- **Returns:** `'DeltaLakeDataset'` A new DeltaLakeDataset limited to n samples. ```python 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)` --- ```python nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__getitem__( index: int ) ``` ```python nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__iter__() -> typing.Iterator[dict] ``` Iterate over training samples from the dataset. ```python nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset.__len__() -> int ``` ```python 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) --- ```python 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) --- ```python 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) ```python 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. --- ```python 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. ```python 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: 1. It starts with "delta://" protocol prefix 2. It's a local directory containing a "\_delta\_log" subdirectory 3. 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. ```python 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 | None` — default: None HF Hub dataset id, e.g. `llamafactory/glaive_toolcall_en`. --- **`path`** `Union[str, List[str]] | None` — default: None Local JSON/JSONL file path or list of paths. --- **`split`** `str` — default: 'train' Dataset split (only used with `dataset_name`). --- **`revision`** `str | None` — default: None Optional Hugging Face dataset revision. Agent trace repositories may expose training-ready rows under `refs/convert/parquet`. --- **`filter_no_assistant`** `bool` — default: False If True, discard incomplete traces that contain no assistant message and therefore have no supervised target. --- **`filter_overlong`** `bool` — default: 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`** `bool` — default: 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 | None` — default: None Optional Jinja template string or file path overriding `tokenizer.chat_template`. --- **`seq_length`** `int | None` — default: None Optional max sequence length for the tokenizer. --- **`limit_dataset_samples`** `int | None` — default: None If set, keep only the first N examples. --- **`padding`** `Union[str, bool]` — default: False Padding strategy forwarded to the tokenizer. --- **`truncation`** `Union[str, bool]` — default: False Truncation strategy forwarded to the tokenizer. --- **`mask_reasoning_content`** `bool` — default: 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`** `bool` — default: 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`** `bool` — default: 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`** `bool` — default: 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`** `bool` — default: 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` ```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.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: > **Note** > > 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. ```python 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`** `int` — default: None If set, pad/truncate each example to this length. --- **`limit_dataset_samples`** `int` — default: None If set, limit the number of examples loaded from the split. --- **`fp8`** `bool` — default: False Flag for future use (e.g., mixed precision). Currently unused. --- **`split`** `str` — default: 'train' Which split of the dataset to load (e.g. 'train', 'validation'). --- **`dataset_name`** `str` — default: 'squad' Identifier for the Hugging Face dataset (default "rajpurkar/squad"). --- **`padding`** `Optional[str | bool]` — default: False Optional padding strategy. --- **`truncation`** `Optional[str | bool]` — default: False Optional truncation strategy. --- **`chat_template`** `str | None` — default: None Optional Jinja template string or path overriding `tokenizer.chat_template`. --- **Returns:** A Hugginggth Face Dataset where each example is a dict with keys: ```python 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. ```python 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