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

## Module Contents

### Functions

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

### Data

[`GENERATION_REGEX`](#nemo_automodel-components-datasets-llm-formatting_utils-GENERATION_REGEX)

[`_warned_add_pad_token`](#nemo_automodel-components-datasets-llm-formatting_utils-_warned_add_pad_token)

[`_warned_generation_prompt`](#nemo_automodel-components-datasets-llm-formatting_utils-_warned_generation_prompt)

[`logger`](#nemo_automodel-components-datasets-llm-formatting_utils-logger)

### API

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

Add pad token to tokenizer if not present.

```python
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 `_build_generation_prompt_mask` for why).

```python
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 `_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 `_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
`_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.

```python
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
  `_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
`_is_consistent_render_prefix`). A known trailing EOS may differ when
the prefix is rendered alone. Any other mismatch raises `ValueError`
because prefix arithmetic cannot place the spans safely.

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

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

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

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

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

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

```python
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`** `int` — default: -100

the value marking unsupervised positions.

---

**Returns:** `List[int]`

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

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

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

```python
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`** — default: 'do\_not\_truncate'

Optional truncation strategy.

---

**`padding`** — default: 'do\_not\_pad'

Optional padding strategy.

---

**`unshifted`** — default: False

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

---

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

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

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

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

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

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

```python
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 | None` — default: None

Optional sequence length for padding.

---

**`tools`** `List[Dict] | None` — default: None

Optional list of tool definitions for function calling.

---

**`answer_only_loss_mask`** `bool` — default: True

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

---

**`mask_reasoning_content`** `bool` — default: False

Whether to exclude rendered reasoning\_content tokens from loss.

---

**`train_on_last_turn_only`** `bool` — default: 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`** `bool` — default: 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
`_build_generation_prompt_mask`.

---

**Returns:** `Dict[str, List[int]]`

A dictionary with the formatted example.

```python
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 | None` — default: None

Optional sequence length for padding.

---

**Returns:** `Dict[str, List[int]]`

A dictionary with the formatted example.

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

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

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

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