nemo_automodel.components.datasets.llm.formatting_utils

Module Contents

Functions

Name	Description
`_add_pad_token`	Add pad token to tokenizer if not present.
`_appended_generation_prompt_length`	Return how many leading tokens of `turn` a generation prompt render supplies, or `0`.
`_build_generation_prompt_mask`	Mark the tokens of each assistant turn that the generation prompt supplies.
`_build_multiturn_assistant_mask`	Build a fallback loss mask that supervises every assistant turn.
`_build_reasoning_mask`	Build a token mask for reasoning_content spans inside assistant turns.
`_common_prefix_length`	Return the number of leading elements `left` and `right` share.
`_find_reasoning_span`	Locate the contiguous token span attributable to reasoning content.
`_get_right_trailing_pad_mask`	Boolean mask identifying right-trailing padding positions.
`_has_chat_template`	Check if the tokenizer supports a chat template.
`_is_consistent_render_prefix`	Check that a prefix render matches the start of the full-conversation render.
`_mask_labels_to_last_turn`	Restrict supervision to the final assistant turn (`mask_history`).
`_masked_reasoning_message`	Return a copy of a message with reasoning_content removed.
`_maybe_shift_mask_for_left_padding`	Shift a token-level mask right when the tokenizer uses left padding.
`_package_tokenized_example`	Package a tokenized example with proper masking and padding.
`_pad_to_seq_length`	Pad a sample to a specific sequence length.
`_resolve_chat_template`	Resolve a chat template string that may be a file path.
`_subtract_mask`	Zero every position of `mask` that `removed` marks.
`_tokenize_chat`	Tokenize chat messages without padding and return input ids.
`_tokenized_chat_length`	Return the tokenized chat length for a message prefix without padding.
`_truncation_window_offset`	Return where the retained (truncated) tokens start inside the untruncated render.
`format_chat_template`	Format a chat template style example.
`format_prompt_completion`	Format a prompt-completion style example (without chat template).

Data

GENERATION_REGEX

_warned_add_pad_token

_warned_generation_prompt

logger

API

nemo_automodel.components.datasets.llm.formatting_utils._add_pad_token(
    tokenizer
)

Add pad token to tokenizer if not present.

nemo_automodel.components.datasets.llm.formatting_utils._appended_generation_prompt_length(
    base_ids: list[int],
    generation_ids: list[int],
    turn: list[int]
) -> int

Return how many leading tokens of turn a generation prompt render supplies, or 0.

base_ids is the conversation prefix rendered plain, generation_ids the same prefix rendered with add_generation_prompt=True and turn the tokens the full render has after base_ids. The prompt counts only when the prompt render is the plain render plus an appended block and the turn opens with that whole block: the block is then exactly what the template emits at inference with no assistant message present, so it holds no token of the message. A prompt render that rewrites the prefix, or a turn that reproduces the block only partly, proves nothing and gives 0 (see :func:_build_generation_prompt_mask for why).

nemo_automodel.components.datasets.llm.formatting_utils._build_generation_prompt_mask(
    tokenizer: transformers.PreTrainedTokenizer,
    formatted_text: list[dict[str, typing.Any]],
    input_ids: list[int],
    tools: list[dict] | None = None,
    truncation: str | bool = 'do_not_truncate',
    seq_length: int | None = None,
    unpadded_full_ids: list[int] | None = None,
    prefix_cache: dict[int, list[int]] | None = None
) -> list[int]

Mark the tokens of each assistant turn that the generation prompt supplies.

At inference the chat template, not the model, emits the assistant role header plus whatever it inserts ahead of the first generated token. For a turn without reasoning that is typically an empty reasoning block such as <think></think> (Nemotron), <think>\n\n</think>\n\n (Qwen3) or </think> (DeepSeek V3.1); for a thinking turn it may be an opening <think>\n. The model never produces these tokens, so supervising them only teaches template boilerplate, and an empty reasoning block also teaches an immediate <think> -> </think> transition.

The span is located without knowing any tag string, from what a generation prompt is by definition: the tokens add_generation_prompt=True appends to the rendered conversation prefix. For every assistant turn, messages[:idx] is rendered plain and with the generation prompt; when the template reads enable_thinking, the prompt is rendered in both modes (True and False), because the data alone cannot tell them apart: a thinking turn may carry its reasoning in reasoning_content or inline in content. A template that never reads the variable renders one prompt, without the keyword (a backend without Jinja templates may reject it). A mode counts only when its prompt render appends a block to the unchanged plain render and the turn in the full render opens with that whole block (see :func:_appended_generation_prompt_length); the longest such block is marked. No token of the message is compared or inferred, so assistant content is never reached, and everything else fails closed:

a plain prefix render that is not an exact prefix of the full render (the template rewrites history as the conversation grows, or replaces its terminator) leaves the turn supervised, warned once per process;
a prompt render that rewrites the prefix instead of appending to it (Gemma 4 prepends a thinking system block, GLM 4.5 appends /nothink to the user turn) is ignored, since its added tokens cannot be told from rendered history;
a turn that reproduces the appended block only partly (SmolLM3, or Qwen3-Thinking without reasoning text, render fewer tokens than the prompt emits) is left supervised. A partial match says nothing about where the template stops and the message starts, because a message token can equal the prompt token at the same position: Gemma 4’s <|channel>thought\n is generated in thinking mode but appended by the non-thinking prompt.

Truncation is handled like :func:_build_reasoning_mask: spans are located in the untruncated render and mapped back through the retained window, which must be a contiguous prefix or suffix of it. A leading assistant turn whose template cannot render an empty conversation is left untouched (warned once per process); any other render error propagates. Positions are computed from unpadded (left-aligned) ids, like :func:_build_multiturn_assistant_mask, whose prefix renders are read from prefix_cache (nothing is written back). One or two extra apply_chat_template renders per assistant turn are the price, which is why this is opt-in.

nemo_automodel.components.datasets.llm.formatting_utils._build_multiturn_assistant_mask(
    tokenizer: transformers.PreTrainedTokenizer,
    formatted_text: typing.List[typing.Dict[str, typing.Any]],
    input_ids: typing.List[int],
    tools: typing.List[typing.Dict] | None = None,
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    seq_length: int | None = None,
    unpadded_full_ids: list[int] | None = None,
    prefix_cache: dict[int, list[int]] | None = None
) -> typing.List[int]

Build a fallback loss mask that supervises every assistant turn.

Each assistant span is located by tokenizing the conversation prefixes before and after the turn, which is O(turns) apply_chat_template calls. Two reductions keep that from re-doing work:

unpadded_full_ids is the caller’s already-known unpadded tokenization of the whole conversation. When the dialogue ends on an assistant turn its closing boundary is the full conversation, so passing it skips re-tokenizing the entire prefix (the single most expensive call in the loop). When omitted, the full conversation is tokenized once here.
Each prefix length is memoized so a boundary shared by adjacent turns (a turn’s end and the next turn’s start) is tokenized at most once. When the caller passes a prefix_cache dict (k -> ids of formatted_text[:k]) the rendered ids are kept in it as well, so :func:_build_generation_prompt_mask can reuse these renders instead of repeating them; without one only the lengths are retained, since holding every prefix render would cost O(turns x sequence length) per sample.

Every tokenized prefix is validated against unpadded_full_ids (see :func:_is_consistent_render_prefix). A known trailing EOS may differ when the prefix is rendered alone. Any other mismatch raises :class:ValueError because prefix arithmetic cannot place the spans safely.

nemo_automodel.components.datasets.llm.formatting_utils._build_reasoning_mask(
    tokenizer: transformers.PreTrainedTokenizer,
    formatted_text: typing.List[typing.Dict[str, typing.Any]],
    input_ids: typing.List[int],
    tools: typing.List[typing.Dict] | None = None,
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    seq_length: int | None = None,
    unpadded_full_ids: list[int] | None = None
) -> typing.List[int]

Build a token mask for reasoning_content spans inside assistant turns.

Each span is isolated by comparing the full conversation render with a second full render where only that message’s reasoning_content is cleared. This remains correct when the template rewrites earlier turns based on later messages. If clearing the field does not change the render, that turn’s reasoning is not present and no tokens need to be masked.

nemo_automodel.components.datasets.llm.formatting_utils._common_prefix_length(
    left: list[int],
    right: list[int]
) -> int

Return the number of leading elements left and right share.

nemo_automodel.components.datasets.llm.formatting_utils._find_reasoning_span(
    full_segment: typing.List[int],
    masked_segment: typing.List[int]
) -> tuple[int, int] | None

Locate the contiguous token span attributable to reasoning content.

nemo_automodel.components.datasets.llm.formatting_utils._get_right_trailing_pad_mask(
    sequence: torch.Tensor,
    pad_token_id: int,
    eos_token_id: int
) -> torch.Tensor

Boolean mask identifying right-trailing padding positions.

When pad_token_id != eos_token_id, it is simply sequence == pad_token_id.

When the two IDs collide, a plain equality check would also match real EOS tokens inside the content. In that case the function locates the trailing contiguous run of the shared token and treats all positions after the first one in that run as padding. The first token in the trailing run is the real EOS and is kept unmasked so the model still learns to predict end-of-sequence.

Parameters:

sequence

torch.Tensor

1-D token id tensor.

pad_token_id

int

The token id used for padding.

eos_token_id

int

The token id used for end-of-sequence. When equal to pad_token_id the positional trailing-run logic is used.

Returns: torch.Tensor

Boolean tensor (same shape as sequence) where True = padding.

nemo_automodel.components.datasets.llm.formatting_utils._has_chat_template(
    tokenizer: transformers.PreTrainedTokenizer
) -> bool

Check if the tokenizer supports a chat template.

Parameters:

tokenizer

PreTrainedTokenizer

The tokenizer to check.

Returns: bool

True if the tokenizer supports a chat template, False otherwise.

nemo_automodel.components.datasets.llm.formatting_utils._is_consistent_render_prefix(
    prefix_ids: list[int],
    reference_ids: list[int],
    trailing_token_id: int | None = None
) -> bool

Check that a prefix render matches the start of the full-conversation render.

Locating spans by prefix lengths is only correct when render(messages[:k]) reproduces the first tokens of render(messages). The comparison is exact unless a multi-token prefix ends with a known standalone terminator that is replaced when the next message is appended.

nemo_automodel.components.datasets.llm.formatting_utils._mask_labels_to_last_turn(
    mask: typing.List[int],
    ignore_index: int = -100
) -> typing.List[int]

Restrict supervision to the final assistant turn (mask_history).

Operates on any per-token sequence where ignore_index marks unsupervised positions: a label list (ignore_index=-100) or a 0/1 assistant mask (ignore_index=0). Each assistant turn renders as a single contiguous supervised span, so this keeps only the last such run and rewrites every earlier supervised position to ignore_index.

Apply this to the assistant mask before any reasoning_content holes are punched into it; running it on already-holed labels would treat the reasoning gap as a turn boundary and drop in-turn content before the hole.

Parameters:

mask

List[int]

per-token labels or 0/1 mask (ignore_index marks unsupervised).

ignore_index

intDefaults to -100

the value marking unsupervised positions.

Returns: List[int]

The same list, mutated so only the final supervised run is kept.

nemo_automodel.components.datasets.llm.formatting_utils._masked_reasoning_message(
    message: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]

Return a copy of a message with reasoning_content removed.

nemo_automodel.components.datasets.llm.formatting_utils._maybe_shift_mask_for_left_padding(
    mask: typing.List[int],
    tokenizer: transformers.PreTrainedTokenizer,
    attention_mask: typing.List[int] | None
) -> typing.List[int]

Shift a token-level mask right when the tokenizer uses left padding.

_build_multiturn_assistant_mask and _build_reasoning_mask compute span indices from unpadded (left-aligned) tokenizations. When the tokenizer pads on the left, actual content is right-aligned in input_ids, so the mask must be shifted right by the padding offset to keep positions aligned.

For right-padding tokenizers (the majority) this is a no-op.

nemo_automodel.components.datasets.llm.formatting_utils._package_tokenized_example(
    tokenizer,
    input_ids,
    assistant_masks,
    eos_token_id,
    pad_token_id,
    seq_length,
    truncation = 'do_not_truncate',
    padding = 'do_not_pad',
    unshifted = False
)

Package a tokenized example with proper masking and padding.

Returns: A dictionary with input_ids, labels, and attention_mask. When unshifted is True, labels is replaced by loss_mask.

Parameters:

tokenizer

The tokenizer to use.

input_ids

The tokenized input ids.

assistant_masks

Boolean mask indicating which tokens are assistant/answer tokens (1) vs prompt tokens (0).

eos_token_id

The end-of-sequence token id.

pad_token_id

The padding token id.

seq_length

Optional sequence length for padding.

truncation

Defaults to 'do_not_truncate'

Optional truncation strategy.

padding

Defaults to 'do_not_pad'

Optional padding strategy.

unshifted

Defaults to False

If True, return unshifted format for dLLM training (input_ids at full length with loss_mask instead of shifted input_ids/labels).

nemo_automodel.components.datasets.llm.formatting_utils._pad_to_seq_length(
    sample,
    pad_token_id,
    seq_length
)

Pad a sample to a specific sequence length.

nemo_automodel.components.datasets.llm.formatting_utils._resolve_chat_template(
    chat_template: str | None
) -> str | None

Resolve a chat template string that may be a file path.

If chat_template points to an existing file, its contents are returned. If opening it as a file fails and the string contains Jinja-like characters ({, }, or newlines) it is treated as a literal template. Otherwise a :class:ValueError is raised so the caller knows the path was invalid.

Parameters:

chat_template

str | None

A Jinja template string or path to a template file.

Returns: str | None

The resolved template string, or None when the input is None.

nemo_automodel.components.datasets.llm.formatting_utils._subtract_mask(
    mask: list[int],
    removed: list[int]
) -> list[int]

Zero every position of mask that removed marks.

nemo_automodel.components.datasets.llm.formatting_utils._tokenize_chat(
    tokenizer: transformers.PreTrainedTokenizer,
    messages: typing.List[typing.Dict[str, typing.Any]],
    tools: typing.List[typing.Dict] | None = None,
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    seq_length: int | None = None,
    template_kwargs: typing.Any = {}
) -> typing.List[int]

Tokenize chat messages without padding and return input ids.

template_kwargs (for example add_generation_prompt or enable_thinking) are forwarded to apply_chat_template as given, so the default path calls the tokenizer exactly as before.

nemo_automodel.components.datasets.llm.formatting_utils._tokenized_chat_length(
    tokenizer: transformers.PreTrainedTokenizer,
    messages: typing.List[typing.Dict[str, str]],
    tools: typing.List[typing.Dict] | None = None,
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    seq_length: int | None = None
) -> int

Return the tokenized chat length for a message prefix without padding.

nemo_automodel.components.datasets.llm.formatting_utils._truncation_window_offset(
    tokenizer: transformers.PreTrainedTokenizer,
    unpadded_full_ids: list[int],
    reference_full_ids: list[int],
    what: str
) -> int

Return where the retained (truncated) tokens start inside the untruncated render.

The retained window must be a contiguous prefix (right truncation) or suffix (left truncation) of reference_full_ids. A short window can match both ends of a render that starts and ends with the same tokens, so the side the tokenizer actually truncates on (tokenizer.truncation_side) is tried first.

nemo_automodel.components.datasets.llm.formatting_utils.format_chat_template(
    tokenizer: transformers.PreTrainedTokenizer,
    formatted_text: typing.List[typing.Dict[str, typing.Any]],
    eos_token_id: int,
    pad_token_id: int,
    seq_length: int | None = None,
    padding: typing.Union[str, bool] = 'do_not_pad',
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    tools: typing.List[typing.Dict] | None = None,
    answer_only_loss_mask: bool = True,
    mask_reasoning_content: bool = False,
    train_on_last_turn_only: bool = False,
    unshifted: bool = False,
    mask_generation_prompt: bool = False
) -> typing.Dict[str, typing.List[int]]

Format a chat template style example.

Parameters:

tokenizer

PreTrainedTokenizer

The tokenizer to use.

formatted_text

List[Dict[str, Any]]

The formatted text, with role tags embedded in the content.

eos_token_id

int

The end-of-sequence token id.

pad_token_id

int

The padding token id.

seq_length

int | NoneDefaults to None

Optional sequence length for padding.

tools

List[Dict] | NoneDefaults to None

Optional list of tool definitions for function calling.

answer_only_loss_mask

boolDefaults to True

Whether to compute the loss mask only on the answer tokens.

mask_reasoning_content

boolDefaults to False

Whether to exclude rendered reasoning_content tokens from loss.

train_on_last_turn_only

boolDefaults to False

Whether to supervise only the final assistant turn, masking every earlier assistant turn (mask_history). Applied to the assistant mask before reasoning_content is masked out.

mask_generation_prompt

boolDefaults to False

Whether to 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 <think></think>). See :func:_build_generation_prompt_mask.

Returns: Dict[str, List[int]]

A dictionary with the formatted example.

nemo_automodel.components.datasets.llm.formatting_utils.format_prompt_completion(
    tokenizer: transformers.PreTrainedTokenizer,
    prompt: str,
    answer: str,
    eos_token_id: int,
    pad_token_id: int,
    seq_length: int | None = None,
    padding: typing.Union[str, bool] = 'do_not_pad',
    truncation: typing.Union[str, bool] = 'do_not_truncate',
    answer_only_loss_mask: bool = True,
    unshifted: bool = False
) -> typing.Dict[str, typing.List[int]]

Format a prompt-completion style example (without chat template).

Parameters:

tokenizer

PreTrainedTokenizer

The tokenizer to use.

prompt

str

The prompt string (e.g. context + question).

answer

str

The answer string.

eos_token_id

int

The end-of-sequence token id.

pad_token_id

int

The padding token id.

seq_length

int | NoneDefaults to None

Optional sequence length for padding.

Returns: Dict[str, List[int]]

A dictionary with the formatted example.

nemo_automodel.components.datasets.llm.formatting_utils.GENERATION_REGEX = re.compile('\\{%-?\\s+generation\\s+-?%\\}')

nemo_automodel.components.datasets.llm.formatting_utils._warned_add_pad_token = set()

nemo_automodel.components.datasets.llm.formatting_utils._warned_generation_prompt: set[str] = set()

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