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

# Integrate Your Own Text Dataset

This guide shows you how to integrate your own dataset into NeMo AutoModel for training. You will learn about two main dataset types: **completion datasets** for language modeling (such as [HellaSwag](https://huggingface.co/datasets/rowan/hellaswag)) and **instruction datasets** for question-answering tasks (such as [SQuAD](https://huggingface.co/datasets/rajpurkar/squad)). This guide covers how to create custom datasets by implementing the required methods and preprocessing functions, and finally shows you how to specify your own data logic using YAML configuration with file paths, allowing you to define custom dataset processing without modifying the main codebase.

## Quick Start Summary

| Type        | Use Case           | Example   | Preprocessor                | Section                       |
| :---------- | :----------------- | :-------- | :-------------------------- | :---------------------------- |
| Completion  | Language Modeling  | HellaSwag | `SFTSingleTurnPreprocessor` | [Jump](#completion-datasets)  |
| Instruction | Question Answering | SQuAD     | `make_*` Function           | [Jump](#instruction-datasets) |

## Types of Supported Datasets

NeMo AutoModel supports a variety of datasets, depending on the task.

### Completion Datasets

**Completion datasets** are single text sequences designed for language modeling where the model learns to predict the next token given a context. These datasets typically contain a context (prompt) and a target (completion) that the model should learn to generate.

#### Example: HellaSwag

The [HellaSwag](https://huggingface.co/datasets/rowan/hellaswag) dataset is a popular completion dataset used for commonsense reasoning. It contains situations with multiple-choice endings where the model must choose the most plausible continuation.

**HellaSwag dataset structure:**

* **Context (`ctx`)**: A situation or scenario description
* **Endings**: Multiple possible completions (four options)
* **Label**: Index of the correct ending

**Example:**

```
Context: "A man is sitting at a piano in a large room."
Endings: [
  "He starts playing a beautiful melody.",
  "He eats a sandwich while sitting there.",
  "He suddenly becomes invisible.",
  "He transforms into a robot."
]
Label: 0  # First ending is correct
```

#### Preprocess with SFTSingleTurnPreprocessor

NeMo AutoModel provides the `SFTSingleTurnPreprocessor` class to handle completion datasets. This processor:

1. **Extracts context and target** using `get_context()` and `get_target()`.
2. **Tokenizes and cleans** the context and target separately.
3. **Concatenates** the context and target into one sequence.
4. **Creates a loss mask** with `-100` for the context and target IDs for the target.
5. **Pads** the sequences to an equal length.

#### Create Your Own Completion Dataset

To adapt your dataset into this format, define a class such as the following:

```python
from datasets import load_dataset
from nemo_automodel.components.datasets.utils import SFTSingleTurnPreprocessor

class MyCompletionDataset:
    def __init__(self, path_or_dataset, tokenizer, split="train"):
        raw_datasets = load_dataset(path_or_dataset, split=split)
        processor = SFTSingleTurnPreprocessor(tokenizer)
        self.dataset = processor.process(raw_datasets, self)

    def get_context(self, examples):
        """Extract context/prompt from your dataset"""
        return examples["context_field"]  # Replace with your context field

    def get_target(self, examples):
        """Extract target/completion from your dataset"""
        return examples["target_field"]   # Replace with your target field

    def __getitem__(self, index):
        return self.dataset[index]

    def __len__(self):
        return len(self.dataset)
```

### Instruction Datasets

**Instruction datasets** are question-answer pairs where the model learns to respond to specific instructions or questions. These datasets are structured as context-question pairs with corresponding answers, making them ideal for teaching models to follow instructions and provide accurate responses.

#### Example: SQuAD

The [SQuAD (Stanford Question Answering Dataset)](https://huggingface.co/datasets/rajpurkar/squad) is a popular instruction dataset for reading comprehension. It contains questions based on Wikipedia articles along with their answers.

**SQuAD dataset structure:**

* **Context**: A paragraph of text from Wikipedia
* **Question**: A question about the context
* **Answers**: The correct answer with its position in the context

#### Create Your Own Instruction Dataset

The [`squad.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/squad.py) file contains the implementation for processing the SQuAD dataset into a format suitable for instruction tuning. It defines a dataset class and preprocessing functions that extract the context, question, and answer fields, concatenate them into a prompt-completion format, and apply tokenization, padding, and loss masking. This serves as a template for building custom instruction datasets by following a similar structure and adapting the extraction logic to your dataset's schema.

Based on the SQuAD implementation in `squad.py`, you can create your own instruction dataset using the `make_squad_dataset` pattern:

```python
from datasets import load_dataset

def make_my_instruction_dataset(
    tokenizer,
    seq_length=None,
    limit_dataset_samples=None,
    split="train",
    dataset_name="your-dataset-name",
):
    if limit_dataset_samples:
        split = f"{split}[:{limit_dataset_samples}]"

    dataset = load_dataset(dataset_name, split=split)

    return dataset.map(
        your_own_fmt_fn,  # Your formatting function
        batched=False,
        remove_columns=dataset.column_names,
    )
```

## YAML-Based Custom Dataset Configuration

NeMo AutoModel supports YAML-based dataset specification using the `_target_` key. This allows you to reference dataset-building classes or functions using either of the following methods:

1. **Python Dotted Path**

```yaml
dataset:
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
  path_or_dataset: rowan/hellaswag
  split: train
```

2. **File Path and Function Name**

```
<file-path>:<function-name>
```

where:

* `<file-path>`: The absolute path to a Python file containing your dataset function
* `<function-name>`: The name of the function to call from that file

```yaml
dataset:
  _target_: /path/to/your/custom_dataset.py:build_my_dataset
  num_blocks: 111
```

This calls `build_my_dataset()` from the specified file with the other keys (such as `num_blocks`) as arguments. This approach allows you to integrate custom datasets using configuration alone, with no need to alter the codebase or package structure.

## Packed Sequence Support in NeMo AutoModel

NeMo AutoModel supports **packed sequences**, which is a technique to optimize training with variable-length sequences (such as text) by minimizing padding.

### What Is a Packed Sequence?

Instead of padding each sequence to a fixed length, which wastes computation on `[PAD]` tokens, packed sequences perform the following actions:

* Concatenate short sequences into a single continuous sequence.
* Separate sequences with special tokens (such as `[EOS]`).
* Track lengths using an "attention mask" to prevent cross-sequence information leakage.

### Benefits

Packed sequences offer the following benefits:

* They reduce redundant computation on padding tokens, which leads to faster training.
* They enable larger effective batch sizes, which leads to better GPU utilization.
* They are especially useful for language modeling and text datasets.

### Enable Packed Sequences in NeMo AutoModel

To enable packed sequences, add these keys to your recipe's YAML config:

```
packed_sequence:
   # Set packed_sequence_size > 0 to run with packed sequences
   packed_sequence_size: 1024
   packing_strategy: thd
   # Optional: tokenize the dataset across this many processes before packing
   num_proc: 8
```

The `packed_sequence` block accepts the following settings:

* **packed\_sequence\_size**: Defines the total token length of each packed sequence. Higher values require more GPU memory.
* **packing\_strategy**: Selects `thd` (the default sequence-length-aware format) or `neat`. With `thd`, AutoModel disables packing when the model does not support `seq_lens`.
* **prepacked**: If `true`, treats the dataset as already packed and skips recipe-side packing.
* **max\_packs**: Optionally caps the number of packed samples created.
* **drop\_long\_samples**: For `neat`, drops samples longer than `packed_sequence_size` when `true`. Otherwise, the packer raises an error.
* **num\_proc**: Specifies the number of worker processes used to tokenize the dataset before packing. Packing first runs one full pass over the dataset, which tokenizes every sample serially. Setting `num_proc > 1` front-loads that tokenization across processes, producing identical, order-preserving output. This defaults to `1` (serial). The speedup grows with the per-sample tokenization cost (such as long multi-turn chats and heavy chat templates) and available CPU cores. This runs independently on each data-parallel rank, so configure it so that `local_ranks * num_proc` does not oversubscribe the node's CPUs.

## Troubleshooting Tips

* **Tokenization Mismatch?** Ensure your tokenizer aligns with the model's expected inputs.
* **Dataset too large?** Use `limit_dataset_samples` in your YAML config to load a subset, which is useful for quick debugging.
* **Loss not decreasing?** Verify that your loss mask correctly ignores prompt tokens.