nemo_automodel.components.datasets.llm.seq2seq

Seq2seq (encoder-decoder) fine-tuning dataset for AutoModelForSeq2SeqLM.

Encoder-decoder models such as T5 and BART differ from decoder-only models in how the training batch is laid out:

The encoder reads input_ids (the source) with its own attention_mask.
labels are the target tokens, kept at full length and not shifted. The model’s loss aligns logits[i] with labels[i] directly.
decoder_input_ids is the right-shifted copy of labels (teacher forcing). HuggingFace builds this internally when labels is passed, but the training loop pops labels before calling the model, so we build it here and put it in the batch so it survives.

This is the opposite of the causal SFT path (see formatting_utils._package_tokenized_example), which concatenates prompt and answer into one stream and pre-shifts input_ids/labels by one position.

The produced per-sample dict reuses the ___PAD_TOKEN_IDS___ convention so that utils.default_collater pads each field with the right value (labels -> -100, the rest -> the pad id / 0).

Module Contents

Functions

Name	Description
`_extract_target_text`	Pull a target string out of a dataset field.
`_format_seq2seq_example`	Turn one raw example into the seq2seq batch fields.
`_shift_right`	Right-shift target tokens to form decoder inputs (teacher forcing).
`make_seq2seq_dataset`	Load and preprocess a dataset for encoder-decoder (seq2seq) fine-tuning.

Data

logger

API

nemo_automodel.components.datasets.llm.seq2seq._extract_target_text(
    value
)

Pull a target string out of a dataset field.

Supports plain strings and the SQuAD answers layout ({"text": [...], "answer_start": [...]}).

nemo_automodel.components.datasets.llm.seq2seq._format_seq2seq_example(
    example,
    tokenizer,
    source_template,
    source_key,
    target_key,
    decoder_start_token_id,
    pad_token_id,
    seq_length = None,
    truncation = True
)

Turn one raw example into the seq2seq batch fields.

Returns a dict with input_ids, attention_mask, labels (unshifted) and decoder_input_ids (right-shifted labels), plus the ___PAD_TOKEN_IDS___ metadata used by default_collater.

nemo_automodel.components.datasets.llm.seq2seq._shift_right(
    token_ids,
    decoder_start_token_id
)

Right-shift target tokens to form decoder inputs (teacher forcing).

Mirrors HuggingFace shift_tokens_right / T5 _shift_right: prepend decoder_start_token_id and drop the final token, so position i of the result is the input that should predict token_ids[i].

Parameters:

token_ids

The (unshifted) target token ids.

decoder_start_token_id

The token the decoder starts from.

Returns:

A list of the same length as token_ids.

nemo_automodel.components.datasets.llm.seq2seq.make_seq2seq_dataset(
    tokenizer,
    seq_length = None,
    limit_dataset_samples = None,
    split = 'train',
    dataset_name = 'rajpurkar/squad',
    source_template = 'question: {question}  cont...,
    source_key = 'question',
    target_key = 'answers',
    decoder_start_token_id = None,
    truncation = True
)

Load and preprocess a dataset for encoder-decoder (seq2seq) fine-tuning.

Each example is tokenized into an encoder source (input_ids + attention_mask) and a decoder target. The target becomes the unshifted labels; decoder_input_ids is its right-shifted copy. default_collater pads labels with -100 and the id fields with the pad id.

Parameters:

tokenizer

A HuggingFace tokenizer (injected by the recipe). Must support the text_target argument for target-side tokenization.

seq_length

intDefaults to None

If set, truncate source and target to this length.

limit_dataset_samples

intDefaults to None

If set, only load this many examples from the split.

split

strDefaults to 'train'

Dataset split to load (e.g. “train”, “validation”).

dataset_name

strDefaults to 'rajpurkar/squad'

HuggingFace dataset identifier. Defaults to SQuAD, framed as a question+context -> answer seq2seq task.

source_template

strDefaults to 'question: {question} context: {context}'

str.format template applied to each example to build the source text. If None, source_key is used verbatim.

source_key

strDefaults to 'question'

Field used as the source when source_template is None.

target_key

strDefaults to 'answers'

Field holding the target. Supports plain strings and the SQuAD answers dict layout.

decoder_start_token_id

intDefaults to None

Token the decoder starts from. If None, defaults to the tokenizer’s pad id (correct for T5/mT5). Models with a different convention (e.g. BART uses eos) should set this explicitly.

truncation

boolDefaults to True

Whether to truncate to seq_length when it is set.

Returns:

A LazyMappedDataset yielding the per-sample seq2seq fields.

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