# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared conversation rendering, masking, and multimodal normalization."""
from __future__ import annotations
import copy
import inspect
import json
import re
import warnings
from collections.abc import Mapping, MutableMapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Literal
import torch
from megatron.bridge.data.datasets.utils import GENERATION_REGEX, IGNORE_INDEX
from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs
CHATML_ASSISTANT_ROLE = "assistant"
_LEGACY_CHAT_ROLE_ALIASES = {"human": "user", "gpt": CHATML_ASSISTANT_ROLE}
# Upper bound on the rendered form of a non-tokenizable payload that the boundary-token scan
# will tokenize. Sized to admit payloads that print themselves in full — a few hundred frame
# reprs, or a float array below numpy's summarization threshold — and reject the rest.
_MAX_SCANNED_REPR_CHARS = 65536
_PIPELINE_CONTROLLED_CHAT_TEMPLATE_KWARGS = frozenset(
{
"add_generation_prompt",
"chat_template",
"continue_final_message",
"max_length",
"padding",
"return_assistant_tokens_mask",
"return_dict",
"return_tensors",
"tokenize",
"tools",
"truncation",
}
)
@dataclass(frozen=True)
class AssistantMaskBoundaryConfig:
"""Token boundary fallback configuration for assistant loss masks.
Expected input format:
``role_start_tokens`` maps each chat role to the token-id sequence that
immediately precedes that role's content in already-rendered
``input_ids``. ``role_end_tokens`` maps each chat role to the token-id
sequence that terminates that role's turn. ``loss_roles`` names
top-level chat roles that should contribute to loss. This follows the
common SFT convention used by chat-template based pipelines: once a role
contributes to loss, the role's full rendered content contributes,
including formatted tool-call or reasoning spans inside that role.
``include_*_tokens_for_roles`` optionally include boundary tokens in the
loss. By default, assistant turn end tokens are included because they
are generated by the assistant and teach the model when to stop.
``trim_leading_token_ids`` removes repeated single-token delimiters from
the start of loss-role content. ``trim_leading_token_sequences`` removes
exact leading token sequences, for example an empty ``reasoning`` span.
``role_end_token_variants`` provides ordered fallback end sequences for
formats whose final turn may omit a trailing delimiter such as a
newline. The primary ``role_end_tokens`` sequence is always preferred.
``rendered_turn_end_roles`` covers formats such as GLM-5.2 whose
assistant turn has no explicit end token. For those roles, the
conversation-prefix render determines the exclusive content boundary.
Example::
AssistantMaskBoundaryConfig(
role_start_tokens={"user": [10], "assistant": [11]},
role_end_tokens={"user": [12], "assistant": [12]},
loss_roles=("assistant",),
include_end_tokens_for_roles=("assistant",),
trim_leading_token_ids=(198,),
trim_leading_token_sequences=([55, 56],),
)
Output format:
When passed to :func:`build_assistant_loss_mask`, the fallback scans
rendered ``input_ids`` for role boundaries and returns an
unshifted float mask with ``1.0`` only on role spans whose role appears
in ``loss_roles``. Boundary marker tokens are included only when
requested by the corresponding ``include_*`` fields.
"""
role_start_tokens: Mapping[str, Sequence[int]]
role_end_tokens: Mapping[str, Sequence[int]]
loss_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,)
include_start_tokens_for_roles: Sequence[str] = ()
include_end_tokens_for_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,)
trim_leading_token_ids: Sequence[int] = ()
trim_leading_token_sequences: Sequence[Sequence[int]] = ()
role_end_token_variants: Mapping[str, Sequence[Sequence[int]]] = field(default_factory=dict)
rendered_turn_end_roles: Sequence[str] = ()
@dataclass(frozen=True)
class NormalizedVLMSample:
"""Source-normalized VLM sample consumed by shared processing.
Expected input format:
Instances are produced by source adapters such as
:func:`normalize_energon_vlm_sample` and
:func:`normalize_hf_vlm_example`. ``conversation`` must be a structured
list of chat turns in HF processor format, for example::
[
{"role": "user", "content": "describe "},
{"role": "assistant", "content": "a red car"},
]
``content`` may also be a multimodal content list accepted by
``processor.apply_chat_template``, for example::
[{"type": "image", "image": image_obj}, {"type": "text", "text": "describe"}]
``images`` and ``videos`` are optional processor-ready modality payloads.
``tools`` contains optional top-level tool definitions passed to the HF
chat template. Message-level ``tool_calls`` remain in ``conversation``.
Energon adapters convert WDS tensors to PIL objects before populating
these fields; HF adapters may leave them ``None`` when media already lives
inline in ``conversation`` content.
Output format:
Shared processing treats this as the single boundary contract before
model-specific tokenization and vision preprocessing. It does not
contain batched tensors; model-specific collators convert it into model
input tensors.
"""
conversation: list[dict[str, Any]]
images: list[Any] | None = None
videos: list[Any] | None = None
audio: Any | None = None
tools: list[dict[str, Any]] | None = None
@dataclass(frozen=True)
class TokenizedConversation:
"""Token ids and an unshifted assistant-only loss mask for one chat row."""
input_ids: torch.Tensor
assistant_mask: torch.Tensor
conversation: list[dict[str, Any]]
final_assistant_start: int | None = None
def normalize_energon_vlm_sample(sample: Any) -> NormalizedVLMSample:
"""Normalize an Energon ``ChatMLSample`` into the shared VLM sample contract.
Expected input format:
``sample`` is expected to expose the Energon ``ChatMLSample`` fields:
- ``conversation``: JSON string accepted by ``cook_chatml_sample``. The
JSON may use either ``{"role": ..., "content": ...}`` turns or
``{"from": ..., "value": ...}`` turns.
- ``imgs``: optional WDS decoded image tensor/list payload.
- ``videos``: optional WDS decoded video tensor/list payload.
- ``audio``: optional audio payload, passed through unchanged.
- ``tools``: optional top-level tool-definition list. Wrapped
``conversation`` JSON may also carry this field.
Output format:
Returns ``NormalizedVLMSample`` where ``conversation`` is a list of
``{"role": str, "content": str | list[dict]}`` turns, ``images`` are
PIL/list processor inputs or ``None``, ``videos`` are nested PIL/list
processor inputs or ``None``, and ``audio`` is copied from the source
sample when present.
"""
from megatron.bridge.data.energon.task_encoder_utils import (
_images_to_pil,
_videos_to_pil,
cook_chatml_sample,
extract_chatml_tools,
)
imgs = getattr(sample, "imgs", None)
videos = getattr(sample, "videos", None)
tools = getattr(sample, "tools", None)
if tools is None:
tools = extract_chatml_tools(sample.conversation)
if isinstance(tools, (str, bytes)):
tools = json.loads(tools)
if tools is not None and (not isinstance(tools, list) or not all(isinstance(tool, Mapping) for tool in tools)):
raise ValueError("Energon ChatML tools must be a list of tool-definition dictionaries.")
if isinstance(videos, (bytes, bytearray)):
normalized_videos = [videos]
else:
normalized_videos = _videos_to_pil(videos) if videos is not None and len(videos) > 0 else None
return NormalizedVLMSample(
conversation=cook_chatml_sample(sample.conversation),
tools=copy.deepcopy(tools),
images=_images_to_pil(imgs) if imgs is not None and len(imgs) > 0 else None,
videos=normalized_videos,
audio=getattr(sample, "audio", None),
)
def normalize_hf_vlm_example(example: Mapping[str, Any]) -> NormalizedVLMSample:
"""Normalize a HF-style VLM dataset example into the shared sample contract.
Expected input format:
``example`` must contain ``"conversation"`` as a structured list of chat
turns already produced by an HF schema adapter, for example::
{
"conversation": [
{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": "Q"}]},
{"role": "assistant", "content": [{"type": "text", "text": "A"}]},
],
"audio": optional_audio,
}
Optional top-level ``images``/``image`` and ``videos``/``video`` fields
are accepted for adapter variants that do not embed media inline in the
conversation. Optional top-level ``tools`` are copied for chat-template
rendering.
Output format:
Returns ``NormalizedVLMSample`` with a deep-copied structured
``conversation`` list, optional list-valued ``images`` and ``videos``
payloads, and optional ``audio``. The adapter does not call
``cook_chatml_sample`` because HF adapters have already normalized the chat
schema.
Raises:
ValueError: If ``example["conversation"]`` is missing or is not a list.
"""
conversation = example.get("conversation")
if not isinstance(conversation, list):
raise ValueError("HF VLM examples must contain a list-valued 'conversation' field.")
images = example.get("images")
if images is None and "image" in example:
images = [example["image"]]
videos = example.get("videos")
if videos is None and "video" in example:
videos = [example["video"]]
return NormalizedVLMSample(
conversation=copy.deepcopy(conversation),
tools=copy.deepcopy(example.get("tools")),
images=images,
videos=videos,
audio=example.get("audio"),
)
def normalized_vlm_sample_to_hf_example(
sample: NormalizedVLMSample,
*,
media_first: bool = False,
) -> dict[str, Any]:
"""Convert a normalized VLM sample into the HF-style collate example schema.
Expected input format:
``sample`` follows ``NormalizedVLMSample``: ``conversation`` is a list of
chat turns, and optional ``images``/``videos`` contain processor-ready
media payloads such as PIL images or decoded video frame lists. Text
turns may contain literal ```` / ```` placeholders.
Output format:
Returns a dictionary suitable for VLM HF collate functions::
{
"conversation": [
{
"role": "user",
"content": [
{"type": "image", "image": image_obj},
{"type": "text", "text": "describe"},
],
},
{"role": "assistant", "content": [{"type": "text", "text": "answer"}]},
],
"images": [image_obj], # present when sample.images is not None
"videos": [video_obj], # present when sample.videos is not None
"audio": audio_obj, # present when sample.audio is not None
"tools": [...], # present when sample.tools is not None
}
Inline media parts are populated from ``sample.images`` and
``sample.videos`` in placeholder order. When ``media_first=True``,
media parts are moved before text parts within each turn to preserve
Qwen Energon's legacy media-before-text ordering while still using the
shared HF collate function.
"""
images = list(sample.images) if sample.images is not None else []
videos = list(sample.videos) if sample.videos is not None else []
image_idx = 0
video_idx = 0
def _media_part(media_type: str) -> dict[str, Any]:
nonlocal image_idx, video_idx
part: dict[str, Any] = {"type": media_type}
if media_type == "image":
if image_idx < len(images):
part["image"] = images[image_idx]
image_idx += 1
elif media_type == "video" and video_idx < len(videos):
part["video"] = videos[video_idx]
video_idx += 1
return part
def _normalize_content(content: Any) -> Any:
nonlocal image_idx, video_idx
if isinstance(content, list):
parts: list[dict[str, Any] | str] = []
for item in content:
if not isinstance(item, Mapping):
parts.append(item)
continue
item_copy = dict(item)
item_type = item_copy.get("type")
if item_type == "image" and "image" not in item_copy and image_idx < len(images):
item_copy["image"] = images[image_idx]
image_idx += 1
elif item_type == "video" and "video" not in item_copy and video_idx < len(videos):
item_copy["video"] = videos[video_idx]
video_idx += 1
parts.append(item_copy)
if media_first:
media_parts = [
part for part in parts if isinstance(part, Mapping) and part.get("type") in ("image", "video")
]
other_parts = [part for part in parts if part not in media_parts]
return media_parts + other_parts
return parts
if not isinstance(content, str) or ("" not in content and "" not in content):
return copy.deepcopy(content)
pieces = re.split(r"(|)", content)
parts = []
for piece in pieces:
if piece == "":
parts.append(_media_part("image"))
elif piece == "":
parts.append(_media_part("video"))
elif piece:
parts.append({"type": "text", "text": piece.strip(" ")})
if media_first:
media_parts = [part for part in parts if part.get("type") in ("image", "video")]
text_parts = [part for part in parts if part.get("type") not in ("image", "video")]
return media_parts + text_parts
return parts
conversation = []
for turn in sample.conversation:
turn_copy = dict(turn)
turn_copy["content"] = _normalize_content(turn_copy.get("content", ""))
conversation.append(turn_copy)
example: dict[str, Any] = {"conversation": conversation}
if sample.tools is not None:
example["tools"] = copy.deepcopy(sample.tools)
if sample.images is not None:
example["images"] = images
if sample.videos is not None:
example["videos"] = videos
if sample.audio is not None:
example["audio"] = sample.audio
return example
def _get_explicit_attribute(obj: Any, name: str) -> Any | None:
"""Read a declared attribute without triggering dynamic ``__getattr__`` fabrication."""
try:
inspect.getattr_static(obj, name)
except AttributeError:
# ``unittest.mock`` stores attributes that a test explicitly configured
# in ``_mock_children`` even though static inspection cannot see them.
# Reading only an existing child preserves those fixtures without ever
# probing an absent dynamic attribute.
try:
namespace = object.__getattribute__(obj, "__dict__")
except AttributeError:
namespace = {}
mock_children = namespace.get("_mock_children", {})
if name not in mock_children:
return None
try:
return getattr(obj, name)
except AttributeError:
return None
def get_processor_tokenizer(processor: Any) -> Any:
"""Return the HF-style tokenizer attached to a processor or Megatron tokenizer wrapper."""
current = processor
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
tokenizer = _get_explicit_attribute(current, "tokenizer")
if tokenizer is not None and tokenizer is not current:
current = tokenizer
continue
# HF fast tokenizers also expose a private ``_tokenizer`` backend. Once
# an object already looks like the callable HF tokenizer, keep it.
if callable(current) and _get_explicit_attribute(current, "added_tokens_decoder") is not None:
return current
tokenizer = _get_explicit_attribute(current, "_tokenizer")
if tokenizer is not None and tokenizer is not current:
current = tokenizer
continue
break
return current
def _as_token_id_list(token_ids: Any) -> list[int]:
if token_ids is None:
return []
if isinstance(token_ids, torch.Tensor):
token_ids = token_ids.detach().cpu().tolist()
if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], torch.Tensor):
token_ids = [item.detach().cpu().tolist() for item in token_ids]
if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], (list, tuple)):
if len(token_ids) != 1:
return []
token_ids = token_ids[0]
return [int(token_id) for token_id in token_ids]
def _conversation_from_example(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
) -> Sequence[Mapping[str, Any]]:
if isinstance(example_or_conversation, Mapping):
conversation = example_or_conversation.get("conversation", [])
return conversation if isinstance(conversation, Sequence) and not isinstance(conversation, str) else []
return example_or_conversation
def chat_template_kwargs_from_example(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
) -> dict[str, Any]:
"""Return optional HF chat-template kwargs stored alongside a conversation.
Rows may provide model-template controls such as ``enable_thinking`` or
``truncate_history_thinking`` under ``chat_template_kwargs``. Rendering
and tokenization controls remain owned by the data pipeline, while tool
schemas use the established top-level ``tools`` field.
Raises:
ValueError: If ``chat_template_kwargs`` is not a string-keyed mapping or
attempts to override a pipeline-controlled argument.
"""
if not isinstance(example_or_conversation, Mapping):
return {}
raw_kwargs = example_or_conversation.get("chat_template_kwargs")
if raw_kwargs is None:
kwargs: dict[str, Any] = {}
elif not isinstance(raw_kwargs, Mapping) or not all(isinstance(key, str) for key in raw_kwargs):
raise ValueError("chat_template_kwargs must be a string-keyed mapping.")
else:
controlled_keys = sorted(_PIPELINE_CONTROLLED_CHAT_TEMPLATE_KWARGS.intersection(raw_kwargs))
if controlled_keys:
joined_keys = ", ".join(controlled_keys)
raise ValueError(f"chat_template_kwargs cannot override pipeline-controlled arguments: {joined_keys}.")
kwargs = dict(raw_kwargs)
truncate_history = kwargs.get("truncate_history_thinking")
if truncate_history is not None and not isinstance(truncate_history, bool):
raise ValueError("chat_template_kwargs.truncate_history_thinking must be a boolean.")
tools = example_or_conversation.get("tools")
if tools is not None:
kwargs["tools"] = tools
return kwargs
def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[str, Any]) -> dict[str, Any]:
"""Translate canonical chat controls to a template's native parameter names.
MBridge exposes ``truncate_history_thinking`` consistently. Templates that
implement the inverse ``preserve_thinking`` control receive the equivalent
negated value only at the rendering boundary, while GLM-style
``clear_thinking`` receives the same value. Custom tokenizers such as
Kimi-K3 use ``thinking`` instead of ``enable_thinking``.
Args:
template_owner: Processor or tokenizer that owns the chat template.
template_kwargs: Canonical chat-template keyword arguments.
Returns:
A copied mapping adapted to the selected template.
"""
kwargs = dict(template_kwargs)
tokenizer = get_processor_tokenizer(template_owner)
template = _get_chat_template(template_owner, tokenizer)
if "truncate_history_thinking" in kwargs and (template is None or "truncate_history_thinking" not in template):
truncate_history = kwargs["truncate_history_thinking"]
if template is not None and "preserve_thinking" in template:
kwargs.pop("truncate_history_thinking")
kwargs["preserve_thinking"] = not truncate_history
elif template is not None and "clear_thinking" in template:
kwargs.pop("truncate_history_thinking")
kwargs["clear_thinking"] = truncate_history
if "enable_thinking" in kwargs and (template is None or "enable_thinking" not in template):
parameters: set[str] = set()
for owner in (template_owner, tokenizer):
apply_chat_template = _get_explicit_attribute(owner, "apply_chat_template")
try:
parameters.update(inspect.signature(apply_chat_template).parameters if apply_chat_template else {})
except (TypeError, ValueError):
continue
if "thinking" in parameters and "enable_thinking" not in parameters:
enable_thinking = kwargs.pop("enable_thinking")
kwargs.setdefault("thinking", enable_thinking)
return kwargs
def shared_chat_template_kwargs_from_examples(examples: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
"""Return dataset-wide chat-template kwargs shared by a processor batch.
Processor-batched VLM collators accept one set of template kwargs for the
whole microbatch. Datasets using these collators must therefore keep those
controls and any tool registry identical across rows.
Raises:
ValueError: If rows in one processor batch use different template kwargs.
"""
if not examples:
return {}
kwargs = chat_template_kwargs_from_example(examples[0])
if any(chat_template_kwargs_from_example(example) != kwargs for example in examples[1:]):
raise ValueError("All examples in one processor batch must use the same chat-template kwargs and tools.")
return kwargs
def _normalize_assistant_tool_call_arguments(turn: dict[str, Any]) -> None:
"""Parse OpenAI JSON-string tool arguments for chat-template consumption."""
if turn.get("role") != CHATML_ASSISTANT_ROLE or not isinstance(turn.get("tool_calls"), list):
return
for tool_call in turn["tool_calls"]:
function = tool_call.get("function") if isinstance(tool_call, Mapping) else None
if not isinstance(function, MutableMapping):
continue
arguments = function.get("arguments")
if not isinstance(arguments, str):
continue
try:
parsed_arguments = json.loads(arguments)
except json.JSONDecodeError as exc:
raise ValueError("Assistant tool call function.arguments must be valid JSON.") from exc
if not isinstance(parsed_arguments, dict):
raise ValueError("Assistant tool call function.arguments must be a JSON object.")
function["arguments"] = parsed_arguments
def normalize_chat_conversation(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
) -> list[dict[str, Any]]:
"""Normalize supported text-chat schemas to OpenAI-style messages.
Rows may use ``messages``, singular ``conversation``, or legacy plural
``conversations`` with ``from``/``value`` keys. Null message content is
converted to an empty string before chat-template rendering, in either
supported schema.
"""
if isinstance(example_or_conversation, Mapping):
source = example_or_conversation
if source.get("messages") is not None:
conversation = source["messages"]
elif source.get("conversation") is not None:
conversation = source["conversation"]
elif source.get("conversations") is not None:
conversation = source["conversations"]
else:
raise ValueError("Chat rows must contain 'messages', 'conversation', or 'conversations'.")
else:
source = None
conversation = example_or_conversation
if not isinstance(conversation, Sequence) or isinstance(conversation, (str, bytes)):
raise ValueError("Chat conversation must be a sequence of message dictionaries.")
normalized: list[dict[str, Any]] = []
for raw_turn in copy.deepcopy(list(conversation)):
if not isinstance(raw_turn, Mapping):
raise ValueError("Chat conversation turns must be dictionaries.")
turn = dict(raw_turn)
if "role" in turn and "content" in turn:
if turn["content"] is None:
turn["content"] = ""
_normalize_assistant_tool_call_arguments(turn)
normalized.append(turn)
continue
if "from" in turn and "value" in turn:
role = str(turn["from"]).lower()
value = turn["value"]
normalized.append(
{
"role": _LEGACY_CHAT_ROLE_ALIASES.get(role, role),
"content": "" if value is None else value,
}
)
continue
raise ValueError("Chat turns must contain role/content or legacy from/value fields.")
if source is not None and source.get("system") and not any(turn.get("role") == "system" for turn in normalized):
normalized.insert(0, {"role": "system", "content": source["system"]})
return normalized
def is_text_only_chat_example(example: Mapping[str, Any]) -> bool:
"""Return whether a supported chat row contains no media content parts."""
top_level_media_keys = (
"image",
"images",
"image_path",
"image_paths",
"video",
"videos",
"video_path",
"video_paths",
"audio",
"audio_path",
"audio_paths",
)
if any(example.get(key) is not None for key in top_level_media_keys):
return False
try:
conversation = normalize_chat_conversation(example)
except ValueError:
return False
for turn in conversation:
content = turn.get("content")
if content is None:
continue
if isinstance(content, str):
if any(marker in content for marker in ("", "", "")):
return False
continue
if not isinstance(content, Sequence) or isinstance(content, (str, bytes)):
return False
for part in content:
if isinstance(part, str):
continue
if not isinstance(part, Mapping) or part.get("type") != "text":
return False
return True
def _chat_template_input_ids(tokenized_chat: Any) -> list[int]:
input_ids = tokenized_chat.get("input_ids") if isinstance(tokenized_chat, Mapping) else tokenized_chat
return _as_token_id_list(input_ids)
def _apply_tokenized_chat_template(
template_owner: Any,
conversation: Sequence[Mapping[str, Any]],
tokenize_kwargs: Mapping[str, Any],
) -> Any:
"""Apply one chat template with compatibility fallbacks."""
apply_chat_template = _get_explicit_attribute(template_owner, "apply_chat_template")
if apply_chat_template is None:
raise ValueError("Chat preprocessing requires a processor or tokenizer with apply_chat_template.")
kwargs = resolve_chat_template_kwargs(template_owner, tokenize_kwargs)
try:
return apply_chat_template(conversation, **kwargs)
except TypeError:
kwargs.pop("add_generation_prompt", None)
try:
return apply_chat_template(conversation, **kwargs)
except TypeError:
kwargs.pop("return_dict", None)
return apply_chat_template(conversation, **kwargs)
def _common_token_prefix_length(left: Sequence[int], right: Sequence[int]) -> int:
"""Return the number of identical leading token IDs."""
return next(
(index for index, (left_id, right_id) in enumerate(zip(left, right)) if left_id != right_id),
min(len(left), len(right)),
)
def _rendered_window_start(full_ids: Sequence[int], rendered_ids: Sequence[int], template_owner: Any) -> int | None:
"""Locate a template-truncated token window in its untruncated rendering."""
if not rendered_ids:
return 0
if len(rendered_ids) > len(full_ids):
return None
tokenizer = get_processor_tokenizer(template_owner)
truncation_side = getattr(tokenizer, "truncation_side", "right")
window_start = len(full_ids) - len(rendered_ids) if truncation_side == "left" else 0
if all(full_ids[window_start + index] == token_id for index, token_id in enumerate(rendered_ids)):
return window_start
return None
def _infer_final_assistant_start(
template_owner: Any,
conversation: Sequence[Mapping[str, Any]],
tokenize_kwargs: Mapping[str, Any],
rendered_ids: Sequence[int],
*,
terminal_only: bool,
) -> int | None:
"""Map the final assistant-turn boundary into a possibly truncated render."""
if not conversation:
return None
if terminal_only and conversation[-1].get("role") != "assistant":
return len(rendered_ids)
final_assistant_index = next(
(index for index in range(len(conversation) - 1, -1, -1) if conversation[index].get("role") == "assistant"),
None,
)
if final_assistant_index is None:
return len(rendered_ids) if terminal_only else None
if final_assistant_index == 0:
return 0
untruncated_kwargs = dict(tokenize_kwargs)
untruncated_kwargs.pop("truncation", None)
untruncated_kwargs.pop("max_length", None)
if "truncation" in tokenize_kwargs:
full_chat = _apply_tokenized_chat_template(template_owner, conversation, untruncated_kwargs)
full_ids = _chat_template_input_ids(full_chat)
else:
full_ids = list(rendered_ids)
prefix_chat = _apply_tokenized_chat_template(
template_owner,
conversation[:final_assistant_index],
untruncated_kwargs,
)
prefix_ids = _chat_template_input_ids(prefix_chat)
untruncated_boundary = _common_token_prefix_length(prefix_ids, full_ids)
window_start = _rendered_window_start(full_ids, rendered_ids, template_owner)
if window_start is None:
# A template-specific truncation transform prevented exact alignment.
# Do not fall back to an earlier assistant turn.
return len(rendered_ids)
return min(max(untruncated_boundary - window_start, 0), len(rendered_ids))
def _infer_terminal_assistant_content_start(
template_owner: Any,
conversation: Sequence[Mapping[str, Any]],
tokenize_kwargs: Mapping[str, Any],
rendered_ids: Sequence[int],
) -> int | None:
"""Map terminal assistant content, excluding its rendered role header."""
if not conversation:
return None
if conversation[-1].get("role") != "assistant":
return len(rendered_ids)
untruncated_kwargs = dict(tokenize_kwargs)
untruncated_kwargs.pop("truncation", None)
untruncated_kwargs.pop("max_length", None)
if "truncation" in tokenize_kwargs:
full_chat = _apply_tokenized_chat_template(template_owner, conversation, untruncated_kwargs)
full_ids = _chat_template_input_ids(full_chat)
else:
full_ids = list(rendered_ids)
empty_conversation = [dict(turn) for turn in conversation]
empty_conversation[-1]["content"] = ""
empty_conversation[-1].pop("tool_calls", None)
empty_chat = _apply_tokenized_chat_template(template_owner, empty_conversation, untruncated_kwargs)
empty_ids = _chat_template_input_ids(empty_chat)
untruncated_content_start = _common_token_prefix_length(empty_ids, full_ids)
window_start = _rendered_window_start(full_ids, rendered_ids, template_owner)
if window_start is None:
return len(rendered_ids)
return min(max(untruncated_content_start - window_start, 0), len(rendered_ids))
def tokenize_chat_example(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
processor: Any,
*,
tool_schemas: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
max_length: int | None = None,
skipped_tokens: torch.Tensor | None = None,
boundary_config: AssistantMaskBoundaryConfig | None = None,
warn_on_all_masked: bool = True,
loss_mode: Literal["assistant", "last_turn", "full"] = "assistant",
return_final_assistant_start: bool = False,
final_assistant_span: Literal["content", "turn"] = "content",
) -> TokenizedConversation:
"""Render, tokenize, and construct the assistant mask for one chat row.
This is the canonical preprocessing entry point shared by ``GPTSFTDataset``
and direct Hugging Face SFT. Padding, label shifting, and sequence packing
remain the responsibility of their respective dataset/collate layers.
GPT dataset callers may request the final assistant boundary for generation
metadata; ``last_turn`` loss infers it automatically.
"""
if loss_mode not in {"assistant", "last_turn", "full"}:
raise ValueError("Chat SFT loss_mode must be assistant, last_turn, or full.")
if final_assistant_span not in {"content", "turn"}:
raise ValueError("final_assistant_span must be content or turn.")
conversation = normalize_chat_conversation(example_or_conversation)
template_kwargs = chat_template_kwargs_from_example(example_or_conversation)
if not template_kwargs.get("tools") and tool_schemas is not None:
# Preserve both the historical single-schema dict form and the common
# list form. Per-row non-empty tools take precedence over this fallback.
template_kwargs["tools"] = tool_schemas
tokenizer = get_processor_tokenizer(processor)
tokenized_chat = None
selected_template_owner = None
selected_tokenize_kwargs = None
seen: set[int] = set()
for template_owner in (tokenizer, processor):
if template_owner is None or id(template_owner) in seen:
continue
seen.add(id(template_owner))
if _get_explicit_attribute(template_owner, "apply_chat_template") is None:
continue
tokenize_kwargs: dict[str, Any] = {
"tokenize": True,
"add_generation_prompt": False,
"return_dict": True,
**template_kwargs,
}
template = _get_chat_template(template_owner)
if template is not None and GENERATION_REGEX.search(template) is not None:
tokenize_kwargs["return_assistant_tokens_mask"] = True
if max_length is not None:
tokenize_kwargs.update({"truncation": True, "max_length": max_length})
tokenized_chat = _apply_tokenized_chat_template(template_owner, conversation, tokenize_kwargs)
selected_template_owner = template_owner
selected_tokenize_kwargs = tokenize_kwargs
break
if tokenized_chat is None:
raise ValueError("Chat preprocessing requires a processor or tokenizer with apply_chat_template.")
input_ids = torch.tensor(_chat_template_input_ids(tokenized_chat), dtype=torch.long)
generation_assistant_start = None
last_turn_start = None
if return_final_assistant_start and selected_template_owner is not None and selected_tokenize_kwargs is not None:
if final_assistant_span == "turn":
generation_assistant_start = _infer_final_assistant_start(
selected_template_owner,
conversation,
selected_tokenize_kwargs,
input_ids.tolist(),
terminal_only=True,
)
else:
generation_assistant_start = _infer_terminal_assistant_content_start(
selected_template_owner,
conversation,
selected_tokenize_kwargs,
input_ids.tolist(),
)
if loss_mode == "last_turn" and selected_template_owner is not None and selected_tokenize_kwargs is not None:
last_turn_start = _infer_final_assistant_start(
selected_template_owner,
conversation,
selected_tokenize_kwargs,
input_ids.tolist(),
terminal_only=False,
)
returned_mask = None
if isinstance(tokenized_chat, Mapping) and tokenized_chat.get("assistant_masks") is not None:
candidate_mask = _as_token_id_list(tokenized_chat["assistant_masks"])
if len(candidate_mask) == input_ids.numel():
returned_mask = torch.tensor(candidate_mask, dtype=torch.bool)
if loss_mode == "full":
full_mask = torch.ones_like(input_ids, dtype=torch.bool)
if skipped_tokens is not None and skipped_tokens.numel() > 0:
full_mask &= ~torch.isin(input_ids, skipped_tokens.to(device=input_ids.device, dtype=torch.long))
return TokenizedConversation(
input_ids=input_ids,
assistant_mask=full_mask,
conversation=conversation,
final_assistant_start=generation_assistant_start,
)
chat_example: dict[str, Any] = {"conversation": conversation}
mask_template_kwargs = dict(template_kwargs)
mask_tools = mask_template_kwargs.pop("tools", None)
if mask_template_kwargs:
chat_example["chat_template_kwargs"] = mask_template_kwargs
if mask_tools is not None:
chat_example["tools"] = mask_tools
effective_boundary_config = boundary_config or infer_assistant_mask_boundary_config(processor)
if returned_mask is not None and effective_boundary_config is None:
assistant_mask = returned_mask
else:
try:
assistant_mask = build_assistant_loss_mask(
chat_example,
input_ids,
processor,
None,
boundary_config=effective_boundary_config,
warn_on_all_masked=False,
).to(dtype=torch.bool)
except ValueError:
if returned_mask is None:
raise
assistant_mask = returned_mask
if last_turn_start is not None:
positions = torch.arange(assistant_mask.numel(), device=assistant_mask.device)
last_turn_positions = torch.nonzero(
assistant_mask.to(dtype=torch.bool) & (positions >= last_turn_start),
as_tuple=False,
).flatten()
if last_turn_positions.numel() > 0:
last_turn_start = int(last_turn_positions[0].item())
elif "truncation" not in (selected_tokenize_kwargs or {}) and bool(assistant_mask.any()):
# Some compatibility tokenizers and simple test doubles ignore the
# conversation argument when rendering a prefix. Without actual
# truncation, retain the historical final visible assistant span.
boolean_mask = assistant_mask.to(dtype=torch.bool)
span_starts = boolean_mask & ~torch.cat(
[torch.zeros(1, dtype=torch.bool, device=boolean_mask.device), boolean_mask[:-1]]
)
last_turn_start = int(torch.nonzero(span_starts, as_tuple=False).flatten()[-1].item())
if (
generation_assistant_start == input_ids.numel()
and conversation
and conversation[-1].get("role") == "assistant"
and "truncation" not in (selected_tokenize_kwargs or {})
and bool(assistant_mask.any())
):
boolean_mask = assistant_mask.to(dtype=torch.bool)
span_starts = boolean_mask & ~torch.cat(
[torch.zeros(1, dtype=torch.bool, device=boolean_mask.device), boolean_mask[:-1]]
)
generation_assistant_start = int(torch.nonzero(span_starts, as_tuple=False).flatten()[-1].item())
assistant_mask = apply_chat_loss_mode(
assistant_mask,
input_ids,
loss_mode=loss_mode,
skipped_tokens=skipped_tokens,
last_turn_start=last_turn_start,
)
_warn_if_all_masked(
assistant_mask,
example_or_conversation,
warn_on_all_masked=warn_on_all_masked,
)
return TokenizedConversation(
input_ids=input_ids,
assistant_mask=assistant_mask,
conversation=conversation,
final_assistant_start=generation_assistant_start,
)
def apply_chat_loss_mode(
assistant_mask: torch.Tensor,
input_ids: torch.Tensor,
*,
loss_mode: Literal["assistant", "last_turn", "full"],
skipped_tokens: torch.Tensor | None = None,
last_turn_start: int | None = None,
) -> torch.Tensor:
"""Apply an explicit chat loss policy to one unpadded token sequence.
``last_turn`` retains assistant tokens after the final assistant-turn
boundary inferred from a prefix render. If a template cannot expose that
boundary, it falls back to the final contiguous assistant-mask span.
Args:
assistant_mask: Assistant-only mask produced by the chat template or
boundary inference.
input_ids: Rendered token IDs aligned with ``assistant_mask``.
loss_mode: Assistant-only, final-assistant-turn, or full-sequence loss.
skipped_tokens: Optional token IDs that must never contribute to loss.
last_turn_start: Optional token boundary inferred by rendering the
conversation prefix before its final assistant turn.
Returns:
Boolean loss mask aligned with ``input_ids``.
"""
if loss_mode == "assistant":
loss_mask = assistant_mask.to(dtype=torch.bool)
elif loss_mode == "full":
loss_mask = torch.ones_like(assistant_mask, dtype=torch.bool)
elif loss_mode == "last_turn":
assistant_mask = assistant_mask.to(dtype=torch.bool)
if last_turn_start is not None:
positions = torch.arange(assistant_mask.numel(), device=assistant_mask.device)
loss_mask = assistant_mask & (positions >= last_turn_start)
else:
span_starts = assistant_mask & ~torch.cat(
[torch.zeros(1, dtype=torch.bool, device=assistant_mask.device), assistant_mask[:-1]]
)
starts = torch.nonzero(span_starts, as_tuple=False).flatten()
if starts.numel() == 0:
loss_mask = assistant_mask
else:
positions = torch.arange(assistant_mask.numel(), device=assistant_mask.device)
loss_mask = assistant_mask & (positions >= starts[-1])
else:
raise ValueError("Chat SFT loss_mode must be assistant, last_turn, or full.")
if skipped_tokens is not None and skipped_tokens.numel() > 0:
loss_mask = loss_mask & ~torch.isin(input_ids, skipped_tokens.to(device=input_ids.device, dtype=torch.long))
return loss_mask
def find_token_span(sequence: Sequence[int] | torch.Tensor, pattern: Sequence[int], start: int = 0) -> tuple[int, int]:
"""Find the first ``[start, end)`` token span matching ``pattern``.
Args:
sequence: Token id sequence to search.
pattern: Token id pattern to locate.
start: Index to begin searching from.
Returns:
``(start, end)`` for the first match, or ``(-1, -1)`` when no match exists.
"""
ids = _as_token_id_list(sequence)
pat = list(pattern)
if not pat or start >= len(ids):
return -1, -1
end_limit = len(ids) - len(pat) + 1
for idx in range(start, max(end_limit, start)):
if ids[idx : idx + len(pat)] == pat:
return idx, idx + len(pat)
return -1, -1
def gather_assistant_text_segments(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
) -> list[str]:
"""Extract assistant text segments from a structured VLM conversation."""
conversation = _conversation_from_example(example_or_conversation)
texts: list[str] = []
for turn in conversation:
if turn.get("role") != "assistant":
continue
content = turn.get("content", "")
text_parts: list[str] = []
if isinstance(content, list):
for part in content:
if isinstance(part, Mapping) and part.get("type") == "text" and isinstance(part.get("text"), str):
text_parts.append(part["text"])
elif isinstance(part, str):
text_parts.append(part)
elif isinstance(content, str):
text_parts.append(content)
if text_parts:
texts.append("".join(text_parts))
return texts
def tokenize_text_without_special_tokens(tokenizer: Any, text: str) -> list[int]:
"""Tokenize text using a HF-like tokenizer without adding special tokens."""
try:
if hasattr(tokenizer, "encode"):
tokens = tokenizer.encode(text, add_special_tokens=False)
elif callable(tokenizer):
tokens = tokenizer(text, add_special_tokens=False)["input_ids"]
else:
return []
except (KeyError, TypeError, ValueError):
return []
return _as_token_id_list(tokens)
def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundaryConfig | None:
"""Infer assistant role boundaries from known HF chat-template formats."""
tokenizer = get_processor_tokenizer(processor)
template = _get_chat_template(tokenizer, processor)
if template is None:
# Kimi-K3 owns rendering in a custom tokenizer rather than a Jinja
# ``chat_template``. Detect its XTML structural vocabulary instead of a
# Python class name so local/remote-code wrappers retain the same behavior.
apply_chat_template = _get_explicit_attribute(tokenizer, "apply_chat_template")
try:
template_parameters = inspect.signature(apply_chat_template).parameters if apply_chat_template else {}
except (TypeError, ValueError):
template_parameters = {}
if "thinking" in template_parameters:
get_vocab = _get_explicit_attribute(tokenizer, "get_vocab")
try:
vocab = get_vocab() if callable(get_vocab) else {}
except (AttributeError, TypeError, ValueError):
vocab = {}
kimi_markers = ("<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>")
if isinstance(vocab, Mapping) and all(marker in vocab for marker in kimi_markers):
return assistant_mask_boundary_config_from_markers(
processor,
assistant_start='<|open|>message role="assistant"<|sep|>',
assistant_end="<|close|>message<|sep|><|end_of_msg|>",
)
return None
# GLM-5.2 does not terminate completed assistant turns with a structural
# token. Prefix rendering is therefore the only loss boundary that remains
# exact for both terminal and historical turns.
if all(marker in template for marker in ("<|assistant|>", "clear_thinking", "")):
return assistant_mask_boundary_config_from_markers(
processor,
assistant_start="<|assistant|>",
assistant_end=None,
role_start_markers={
"system": "<|system|>",
"user": "<|user|>",
"tool": "<|observation|>",
},
rendered_turn_end_roles=(CHATML_ASSISTANT_ROLE,),
)
candidates: tuple[tuple[tuple[str, ...], str, str, tuple[str, ...], Mapping[str, str]], ...] = (
(
("<|im_start|>assistant", "<|im_end|>"),
"<|im_start|>assistant\n",
"<|im_end|>\n",
("<|im_end|>",),
{role: f"<|im_start|>{role}\n" for role in ("system", "developer", "user", "tool")},
),
(
("<|im_assistant|>", "<|im_user|>", "<|im_system|>", "<|im_middle|>", "<|im_end|>"),
"<|im_assistant|>assistant<|im_middle|>",
"<|im_end|>",
(),
{
"system": "<|im_system|>system<|im_middle|>",
"user": "<|im_user|>user<|im_middle|>",
},
),
(("<|turn>model", ""), "<|turn>model\n", "", ("<|tool_response>",), {}),
(("model", ""), "model\n", "", (), {}),
(
("<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"),
"<|start_header_id|>assistant<|end_header_id|>\n\n",
"<|eot_id|>",
(),
{
role: f"<|start_header_id|>{role}<|end_header_id|>\n\n"
for role in ("system", "developer", "user", "tool")
},
),
# OpenAI Harmony (gpt-oss). One assistant turn renders as one or more
# channel segments, e.g. ``<|start|>assistant<|channel|>analysis<|message|>...
# <|end|><|start|>assistant<|channel|>final<|message|>...<|return|>``. The
# start marker stops at ``<|start|>assistant`` so the loss span begins at
# ``<|channel|>`` and covers every channel of the turn, analysis included.
# The template emits ``<|start|>`` and the role separately, so the required
# markers use the standalone structural specials rather than a concatenated
# ``<|start|>assistant``. End variants cover message (``<|end|>``), final
# generation (``<|return|>``), and tool call (``<|call|>``) terminators.
(
("<|start|>", "<|channel|>", "<|message|>"),
"<|start|>assistant",
"<|return|>",
("<|end|>", "<|call|>"),
{role: f"<|start|>{role}" for role in ("system", "developer", "user", "tool")},
),
(
("]~!b[", "]~b]", "[e~[", ""),
"]~b]ai\n",
"[e~[\n",
(),
{
"system": "]~!b[]~b]system\n",
"developer": "]~b]developer\n",
"user": "]~b]user\n",
"tool": "]~b]tool",
},
),
)
for required_markers, assistant_start, assistant_end, assistant_end_fallbacks, other_role_starts in candidates:
if not all(marker in template for marker in required_markers):
continue
start_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_start)
end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end)
end_token_variants = [
token_ids
for fallback in assistant_end_fallbacks
if (token_ids := tokenize_text_without_special_tokens(tokenizer, fallback))
]
if start_tokens and end_tokens:
role_start_tokens = {CHATML_ASSISTANT_ROLE: start_tokens}
role_start_tokens.update(
{
role: token_ids
for role, marker in other_role_starts.items()
if (token_ids := tokenize_text_without_special_tokens(tokenizer, marker))
}
)
trim_leading_token_sequences = ()
if all(marker in template for marker in ("truncate_history_thinking", "", " ")):
empty_think_tokens = tokenize_text_without_special_tokens(tokenizer, "\n")
if empty_think_tokens and think_open_tokens:
trim_leading_token_sequences = (empty_think_tokens, think_open_tokens)
return AssistantMaskBoundaryConfig(
role_start_tokens=role_start_tokens,
role_end_tokens={role: end_tokens for role in role_start_tokens},
role_end_token_variants={role: end_token_variants for role in role_start_tokens},
trim_leading_token_sequences=trim_leading_token_sequences,
)
return None
def assistant_mask_boundary_config_from_markers(
processor: Any,
*,
assistant_start: str,
assistant_end: str | None,
assistant_end_fallbacks: Sequence[str] = (),
role_start_markers: Mapping[str, str] | None = None,
loss_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
include_start_tokens_for_roles: Sequence[str] = (),
include_end_tokens_for_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
trim_leading_token_ids: Sequence[int] = (),
trim_leading_token_sequences: Sequence[Sequence[int]] = (),
rendered_turn_end_roles: Sequence[str] = (),
) -> AssistantMaskBoundaryConfig:
"""Build explicit assistant boundary config from model-declared marker strings."""
tokenizer = get_processor_tokenizer(processor)
start_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_start)
end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end) if assistant_end is not None else []
end_token_variants = [
token_ids
for fallback in assistant_end_fallbacks
if (token_ids := tokenize_text_without_special_tokens(tokenizer, fallback))
]
if not start_tokens or (not end_tokens and CHATML_ASSISTANT_ROLE not in rendered_turn_end_roles):
raise ValueError(
"Unable to tokenize assistant loss-mask boundary markers. "
"Check the model collator's assistant_start and assistant_end markers."
)
role_start_tokens = {CHATML_ASSISTANT_ROLE: start_tokens}
for role, marker in (role_start_markers or {}).items():
marker_tokens = tokenize_text_without_special_tokens(tokenizer, marker)
if not marker_tokens:
raise ValueError(f"Unable to tokenize loss-mask start marker for chat role '{role}'.")
role_start_tokens[role] = marker_tokens
return AssistantMaskBoundaryConfig(
role_start_tokens=role_start_tokens,
role_end_tokens={role: end_tokens for role in role_start_tokens if role not in rendered_turn_end_roles},
loss_roles=loss_roles,
include_start_tokens_for_roles=include_start_tokens_for_roles,
include_end_tokens_for_roles=include_end_tokens_for_roles,
trim_leading_token_ids=trim_leading_token_ids,
trim_leading_token_sequences=trim_leading_token_sequences,
role_end_token_variants={role: end_token_variants for role in role_start_tokens},
rendered_turn_end_roles=rendered_turn_end_roles,
)
def _apply_skipped_tokens(mask: torch.Tensor, ids: Sequence[int], skipped_tokens: torch.Tensor | None) -> None:
if skipped_tokens is None or skipped_tokens.numel() == 0:
return
skipped = set(skipped_tokens.detach().cpu().tolist())
for idx, token_id in enumerate(ids):
if token_id in skipped:
mask[idx] = 0.0
def _warn_if_all_masked(
mask: torch.Tensor,
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
*,
warn_on_all_masked: bool,
) -> None:
if warn_on_all_masked and len(mask) > 0 and float(mask.sum().item()) == 0.0:
warnings.warn("*" * 100, stacklevel=3)
warnings.warn(f"All tokens are masked for example:\n{example_or_conversation}.", stacklevel=3)
warnings.warn("*" * 100, stacklevel=3)
def _get_chat_template(*objects: Any) -> str | None:
for obj in objects:
template = _get_explicit_attribute(obj, "chat_template")
if isinstance(template, str) and template:
return template
return None
def _assistant_mask_from_hf_chat_template(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
ids: Sequence[int],
processor: Any,
tokenizer: Any,
) -> torch.Tensor | None:
"""Return HF generation-block assistant mask when it aligns to ``ids``."""
conversation = _conversation_from_example(example_or_conversation)
chat_template_kwargs = chat_template_kwargs_from_example(example_or_conversation)
seen: set[int] = set()
for template_owner in (tokenizer, processor):
if id(template_owner) in seen:
continue
seen.add(id(template_owner))
apply_chat_template = _get_explicit_attribute(template_owner, "apply_chat_template")
template = _get_chat_template(template_owner)
if apply_chat_template is None or template is None or GENERATION_REGEX.search(template) is None:
continue
try:
tokenized_chat = _apply_tokenized_chat_template(
template_owner,
conversation,
{
"tokenize": True,
"add_generation_prompt": False,
"return_dict": True,
"return_assistant_tokens_mask": True,
**chat_template_kwargs,
},
)
except (AttributeError, KeyError, TypeError, ValueError):
continue
if not isinstance(tokenized_chat, Mapping):
continue
if "assistant_masks" not in tokenized_chat:
continue
candidate_ids = _as_token_id_list(tokenized_chat.get("input_ids"))
assistant_mask = _as_token_id_list(tokenized_chat["assistant_masks"])
aligned_mask = _align_assistant_mask_to_padded_ids(candidate_ids, assistant_mask, ids, tokenizer)
if aligned_mask is None:
continue
return torch.tensor(aligned_mask, dtype=torch.float32)
return None
def _align_assistant_mask_to_padded_ids(
candidate_ids: Sequence[int],
assistant_mask: Sequence[int],
ids: Sequence[int],
tokenizer: Any,
) -> list[int] | None:
"""Align an HF assistant mask to possibly padded processor batch ids."""
target_ids = list(ids)
source_ids = list(candidate_ids)
source_mask = [int(mask_value) for mask_value in assistant_mask]
if len(source_ids) != len(source_mask):
return None
if source_ids == target_ids:
return source_mask
if len(source_ids) > len(target_ids):
window_start = _rendered_window_start(source_ids, target_ids, tokenizer)
if window_start is None:
return None
return source_mask[window_start : window_start + len(target_ids)]
pad_token_id = getattr(tokenizer, "pad_token_id", None)
if pad_token_id is None:
return None
pad_len = len(target_ids) - len(source_ids)
if target_ids[: len(source_ids)] == source_ids and all(
token_id == pad_token_id for token_id in target_ids[-pad_len:]
):
return source_mask + [0] * pad_len
if target_ids[pad_len:] == source_ids and all(token_id == pad_token_id for token_id in target_ids[:pad_len]):
return [0] * pad_len + source_mask
return None
def _assistant_mask_from_conversation_turns(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
ids: Sequence[int],
template_owner: Any,
tokenizer: Any,
boundary_config: AssistantMaskBoundaryConfig,
*,
include_content: bool = True,
) -> torch.Tensor | None:
"""Build a role-safe mask by locating turns through conversation-prefix renders."""
conversation = _conversation_from_example(example_or_conversation)
loss_roles = set(boundary_config.loss_roles)
loss_turns = [(turn_index, turn) for turn_index, turn in enumerate(conversation) if turn.get("role") in loss_roles]
if not loss_turns:
return None
tokenize_kwargs = {
"tokenize": True,
"add_generation_prompt": False,
"return_dict": True,
**chat_template_kwargs_from_example(example_or_conversation),
}
template = _get_chat_template(template_owner)
if template is not None and GENERATION_REGEX.search(template) is not None:
tokenize_kwargs["return_assistant_tokens_mask"] = True
try:
rendered = _apply_tokenized_chat_template(template_owner, conversation, tokenize_kwargs)
rendered_ids = _chat_template_input_ids(rendered)
except (AssertionError, AttributeError, KeyError, TypeError, ValueError):
return None
role_start_tokens = _token_map_from_boundary_config(boundary_config.role_start_tokens)
role_end_tokens = _token_map_from_boundary_config(boundary_config.role_end_tokens)
role_end_token_variants = {
role: [token_ids for raw_token_ids in raw_variants if (token_ids := _as_token_id_list(raw_token_ids))]
for role, raw_variants in boundary_config.role_end_token_variants.items()
}
include_start_roles = set(boundary_config.include_start_tokens_for_roles)
include_end_roles = set(boundary_config.include_end_tokens_for_roles)
rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
rendered_mask = torch.zeros(len(rendered_ids), dtype=torch.float32)
for turn_index, turn in loss_turns:
role = turn["role"]
start_tokens = role_start_tokens.get(role)
end_patterns = [role_end_tokens.get(role, []), *role_end_token_variants.get(role, [])]
if not start_tokens or (not any(end_patterns) and role not in rendered_turn_end_roles):
return None
try:
before_ids = (
_chat_template_input_ids(
_apply_tokenized_chat_template(template_owner, conversation[:turn_index], tokenize_kwargs)
)
if turn_index > 0
else []
)
through_ids = (
rendered_ids
if turn_index + 1 == len(conversation)
else _chat_template_input_ids(
_apply_tokenized_chat_template(template_owner, conversation[: turn_index + 1], tokenize_kwargs)
)
)
except (AssertionError, AttributeError, KeyError, TypeError, ValueError):
return None
turn_search_start = _common_token_prefix_length(before_ids, rendered_ids)
turn_limit = _common_token_prefix_length(through_ids, rendered_ids)
if turn_search_start >= turn_limit:
return None
role_start, content_start = find_token_span(rendered_ids, start_tokens, turn_search_start)
matched_start_length = len(start_tokens)
if role_start < 0:
if turn_index == 0 or conversation[turn_index - 1].get("role") != "tool":
return None
# Some tool templates continue the same model turn after an inline tool response,
# so the next assistant payload begins at the prefix boundary without another role header.
role_start = content_start = turn_search_start
matched_start_length = 0
if content_start > turn_limit:
return None
if role in rendered_turn_end_roles:
content_end = after_end = turn_limit
if turn_index + 1 < len(conversation):
# Only this turn's payload renders between ``content_start`` and the next role
# header, so a marker quoted in any other turn cannot shift that boundary.
boundary_tokens_in_payload = _conversation_contains_boundary_tokens(
example_or_conversation, tokenizer, boundary_config, turns=[turn]
)
if boundary_tokens_in_payload is not False:
return None
next_role = conversation[turn_index + 1].get("role")
next_role_start_tokens = role_start_tokens.get(next_role, [])
next_role_start, _ = find_token_span(rendered_ids, next_role_start_tokens, content_start)
if next_role_start < 0:
return None
content_end = after_end = next_role_start
else:
candidate_ends: list[tuple[int, int, int, list[int]]] = []
for priority, pattern in enumerate(end_patterns):
if not pattern:
continue
search_start = content_start
while search_start < turn_limit:
end_start, after_end = find_token_span(rendered_ids, pattern, search_start)
if end_start < 0 or after_end > turn_limit:
break
candidate_ends.append((end_start, -priority, after_end, pattern))
search_start = end_start + 1
if not candidate_ends:
return None
content_end, _, after_end, _ = max(candidate_ends)
loss_start = _trim_leading_token_sequences(
rendered_ids,
content_start,
content_end,
boundary_config.trim_leading_token_sequences,
)
loss_start = _trim_leading_token_ids(
rendered_ids,
loss_start,
content_end,
_as_token_id_list(boundary_config.trim_leading_token_ids),
)
if include_content:
rendered_mask[loss_start:content_end] = 1.0
if role in include_start_roles:
rendered_mask[role_start : role_start + matched_start_length] = 1.0
if role in include_end_roles and role not in rendered_turn_end_roles:
rendered_mask[content_end:after_end] = 1.0
aligned_mask = _align_assistant_mask_to_padded_ids(rendered_ids, rendered_mask.tolist(), ids, tokenizer)
if aligned_mask is None:
return None
return torch.tensor(aligned_mask, dtype=torch.float32)
def _token_map_from_boundary_config(token_map: Mapping[str, Sequence[int]] | None) -> dict[str, list[int]]:
if token_map is None:
return {}
return {
role: token_ids for role, raw_token_ids in token_map.items() if (token_ids := _as_token_id_list(raw_token_ids))
}
def _value_contains_token_sequence(
value: Any,
tokenizer: Any,
token_sequences: Sequence[Sequence[int]],
) -> bool | None:
"""Check recursively tokenizable payload values for configured boundary sequences."""
if isinstance(value, (bytes, bytearray, memoryview)):
# A raw buffer reaches the prompt only as an escaped repr literal, so escaping hides
# any marker inside it and this scan can never clear one. Stringifying a decoded image
# would also cost megabytes per example.
return None
if isinstance(value, str):
try:
token_ids = tokenize_text_without_special_tokens(tokenizer, value)
except (AssertionError, AttributeError, KeyError, TypeError, ValueError):
return None
return any(find_token_span(token_ids, token_sequence)[0] >= 0 for token_sequence in token_sequences)
if isinstance(value, Mapping):
results = [_value_contains_token_sequence(item, tokenizer, token_sequences) for item in value.values()]
elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
results = [_value_contains_token_sequence(item, tokenizer, token_sequences) for item in value]
elif value is None or isinstance(value, (bool, int, float)):
return False
else:
# Leaves we cannot tokenize directly — PIL images, arrays, tensors, paths. Reporting
# unknown here would disable the boundary-config fallback for every multimodal example,
# so scan the text the payload renders as instead. Jinja emits an object through str()
# standalone but repr() inside a container, so both forms must look clean to report
# clean. This is a heuristic: a payload whose rendered text its own str()/repr() does
# not reveal can still slip through — an array that abbreviates its printout, or an
# opaque container whose repr escapes the newline inside a marker, the same escaping
# that makes raw buffers above undetectable. Reaching either needs a template that
# iterates or attribute-accesses the payload rather than rendering it whole.
try:
if callable(value) or (type(value).__str__ is object.__str__ and type(value).__repr__ is object.__repr__):
# Templates expand callables by introspection (HF renders `tools` entries through
# get_json_schema, pulling in names and docstrings), and an object with neither
# hook has no rendered text at all. Neither can be cleared by stringifying it.
return None
rendered = {str(value), repr(value)}
# str.__len__ rather than len(): __str__ may return a str subclass with its own.
oversized = any(str.__len__(form) > _MAX_SCANNED_REPR_CHARS for form in rendered)
except Exception:
# Rendering an arbitrary payload can raise anything, and callers rely on None to
# fail closed — tokenize_chat_example only recovers from ValueError, so anything
# escaping here kills the dataloader worker instead.
return None
if oversized:
# A payload that prints itself in full can run to megabytes; tokenizing that once
# per example would stall the dataloader. Treat it as unknown provenance instead.
return None
results = [_value_contains_token_sequence(form, tokenizer, token_sequences) for form in rendered]
if any(result is True for result in results):
return True
return None if any(result is None for result in results) else False
def _conversation_contains_boundary_tokens(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
tokenizer: Any,
boundary_config: AssistantMaskBoundaryConfig,
*,
turns: Sequence[Mapping[str, Any]] | None = None,
) -> bool | None:
"""Return whether rendered conversation payloads contain structural token sequences.
``turns`` restricts the scan to those turns instead of the whole conversation. Template
kwargs are always scanned.
"""
token_sequences = [
*(_token_map_from_boundary_config(boundary_config.role_start_tokens).values()),
*(_token_map_from_boundary_config(boundary_config.role_end_tokens).values()),
*(
token_ids
for raw_variants in boundary_config.role_end_token_variants.values()
for raw_token_ids in raw_variants
if (token_ids := _as_token_id_list(raw_token_ids))
),
]
scanned_turns = _conversation_from_example(example_or_conversation) if turns is None else turns
payloads: list[Any] = [{key: value for key, value in turn.items() if key != "role"} for turn in scanned_turns]
payloads.append(chat_template_kwargs_from_example(example_or_conversation))
return _value_contains_token_sequence(payloads, tokenizer, token_sequences)
def _trim_leading_token_ids(
ids: Sequence[int],
start: int,
end: int,
trim_token_ids: Sequence[int],
) -> int:
trim_tokens = set(trim_token_ids)
idx = start
while idx < end and ids[idx] in trim_tokens:
idx += 1
return idx
def _trim_leading_token_sequences(
ids: Sequence[int],
start: int,
end: int,
trim_token_sequences: Sequence[Sequence[int]],
) -> int:
"""Skip complete token sequences that appear contiguously at a span start.
This is intentionally sequence-based instead of token-set based. For Kimi,
it can remove a leading empty ``reasoning... `` block as supervised content.
Args:
ids: Rendered token ids to inspect.
start: Inclusive span start.
end: Exclusive span end.
trim_token_sequences: Candidate token-id sequences to skip when they
match exactly at the current span start.
Returns:
The first index after all matched leading token sequences.
"""
idx = start
while idx < end:
for raw_token_sequence in trim_token_sequences:
token_sequence = _as_token_id_list(raw_token_sequence)
sequence_end = idx + len(token_sequence)
if token_sequence and sequence_end <= end and list(ids[idx:sequence_end]) == token_sequence:
idx = sequence_end
break
else:
return idx
return idx
def _assistant_mask_from_boundary_config(
ids: Sequence[int],
boundary_config: AssistantMaskBoundaryConfig,
*,
include_content: bool = True,
conversation_roles: Sequence[str] | None = None,
) -> torch.Tensor | None:
role_start_tokens = _token_map_from_boundary_config(boundary_config.role_start_tokens)
role_end_tokens = _token_map_from_boundary_config(boundary_config.role_end_tokens)
role_end_token_variants = {
role: [token_ids for raw_token_ids in raw_variants if (token_ids := _as_token_id_list(raw_token_ids))]
for role, raw_variants in boundary_config.role_end_token_variants.items()
}
loss_roles = set(boundary_config.loss_roles)
rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
if rendered_turn_end_roles:
# A raw token scan cannot distinguish the end of an unterminated turn
# from following non-loss roles. Only the conversation-prefix renderer
# can establish this boundary safely.
return None
if not role_start_tokens or not role_end_tokens or not loss_roles:
raise ValueError(
"AssistantMaskBoundaryConfig must provide non-empty role_start_tokens, role_end_tokens, and loss_roles."
)
missing_role_end_tokens = set(role_start_tokens) - set(role_end_tokens)
if missing_role_end_tokens:
raise ValueError(
"AssistantMaskBoundaryConfig role_end_tokens is missing entries for roles: "
f"{sorted(missing_role_end_tokens)}."
)
markers: list[tuple[int, str, int]] = []
for role, start_tokens in role_start_tokens.items():
search_start = 0
while search_start < len(ids):
start_pos, after_start = find_token_span(ids, start_tokens, search_start)
if start_pos < 0:
break
markers.append((start_pos, role, len(start_tokens)))
search_start = after_start
if not markers:
return None
markers.sort(key=lambda marker: (marker[0], -marker[2]))
deduped_markers: list[tuple[int, str, int]] = []
for marker in markers:
if deduped_markers and marker[0] == deduped_markers[-1][0]:
continue
deduped_markers.append(marker)
if conversation_roles is not None:
expected_loss_roles = [role for role in conversation_roles if role in loss_roles]
observed_loss_roles = [role for _, role, _ in deduped_markers if role in loss_roles]
if observed_loss_roles != expected_loss_roles:
return None
include_start_tokens_for_roles = set(boundary_config.include_start_tokens_for_roles)
include_end_tokens_for_roles = set(boundary_config.include_end_tokens_for_roles)
trim_token_ids = _as_token_id_list(boundary_config.trim_leading_token_ids)
trim_token_sequences = boundary_config.trim_leading_token_sequences
mask = torch.zeros(len(ids), dtype=torch.float32)
found_loss_segment = False
for marker_idx, (start_pos, role, marker_len) in enumerate(deduped_markers):
content_start = start_pos + marker_len
next_marker_start = deduped_markers[marker_idx + 1][0] if marker_idx + 1 < len(deduped_markers) else len(ids)
if role not in loss_roles:
continue
candidate_end_spans = []
for priority, role_end_pattern in enumerate([role_end_tokens[role], *role_end_token_variants.get(role, [])]):
search_start = content_start
while search_start < next_marker_start:
end_start, after_end = find_token_span(ids, role_end_pattern, search_start)
if end_start < 0 or after_end > next_marker_start:
break
candidate_end_spans.append((end_start, priority, after_end))
search_start = end_start + 1
if not candidate_end_spans:
continue
if conversation_roles is not None and len({end_start for end_start, _, _ in candidate_end_spans}) != 1:
return None
end_start, _, after_end = min(candidate_end_spans)
role_end_span = (end_start, after_end)
found_loss_segment = True
content_end = role_end_span[0]
if include_content:
segment_start = _trim_leading_token_sequences(ids, content_start, content_end, trim_token_sequences)
segment_start = _trim_leading_token_ids(ids, segment_start, content_end, trim_token_ids)
if segment_start < content_end:
mask[segment_start:content_end] = 1.0
if role in include_start_tokens_for_roles:
mask[start_pos:content_start] = 1.0
if role in include_end_tokens_for_roles and role_end_span is not None:
mask[role_end_span[0] : role_end_span[1]] = 1.0
return mask if found_loss_segment else None
def build_assistant_loss_mask(
example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]],
input_ids: Sequence[int] | torch.Tensor,
processor: Any,
skipped_tokens: torch.Tensor | None = None,
*,
boundary_config: AssistantMaskBoundaryConfig | None = None,
warn_on_all_masked: bool = True,
) -> torch.Tensor:
"""Build an unshifted assistant-only loss mask.
The preferred path uses HF chat templates with ``{% generation %}`` blocks.
When those are unavailable, callers may provide ``boundary_config`` to scan
explicit role token boundaries from rendered ``input_ids``.
Raises:
ValueError: If neither a HF generation mask nor a valid explicit
boundary mask can be built.
"""
tokenizer = get_processor_tokenizer(processor)
ids = _as_token_id_list(input_ids)
conversation = _conversation_from_example(example_or_conversation)
conversation_roles = [turn.get("role") for turn in conversation] if conversation else None
mask = _assistant_mask_from_hf_chat_template(example_or_conversation, ids, processor, tokenizer)
if boundary_config is not None:
boundary_scan_is_safe = not boundary_config.rendered_turn_end_roles and (
not conversation
or _conversation_contains_boundary_tokens(example_or_conversation, tokenizer, boundary_config) is False
)
if mask is None or float(mask.sum().item()) == 0.0:
boundary_mask = _assistant_mask_from_conversation_turns(
example_or_conversation,
ids,
tokenizer,
tokenizer,
boundary_config,
)
if boundary_mask is None and processor is not tokenizer:
boundary_mask = _assistant_mask_from_conversation_turns(
example_or_conversation,
ids,
processor,
tokenizer,
boundary_config,
)
if boundary_mask is None and boundary_scan_is_safe:
boundary_mask = _assistant_mask_from_boundary_config(
ids,
boundary_config,
conversation_roles=conversation_roles,
)
if boundary_mask is None and not boundary_scan_is_safe and mask is not None:
mask = None
if mask is None or (boundary_mask is not None and float(boundary_mask.sum().item()) > 0.0):
mask = boundary_mask
else:
boundary_token_mask = _assistant_mask_from_conversation_turns(
example_or_conversation,
ids,
tokenizer,
tokenizer,
boundary_config,
include_content=False,
)
if boundary_token_mask is None and processor is not tokenizer:
boundary_token_mask = _assistant_mask_from_conversation_turns(
example_or_conversation,
ids,
processor,
tokenizer,
boundary_config,
include_content=False,
)
if boundary_token_mask is None and boundary_scan_is_safe:
boundary_token_mask = _assistant_mask_from_boundary_config(
ids,
boundary_config,
include_content=False,
conversation_roles=conversation_roles,
)
if boundary_token_mask is not None:
mask = torch.maximum(mask, boundary_token_mask)
if mask is None:
if boundary_config is None:
raise ValueError(
"Unable to build assistant loss mask. Provide a processor/tokenizer chat_template with "
"{% generation %} blocks that returns assistant_masks, or pass AssistantMaskBoundaryConfig "
"with explicit role token boundaries."
)
raise ValueError("AssistantMaskBoundaryConfig did not match any loss-contributing spans in input_ids.")
_apply_skipped_tokens(mask, ids, skipped_tokens)
_warn_if_all_masked(mask, example_or_conversation, warn_on_all_masked=warn_on_all_masked)
return mask
def build_shifted_labels_and_loss_mask(
input_ids: torch.Tensor,
assistant_loss_mask: torch.Tensor,
skipped_tokens: torch.Tensor | None = None,
*,
ignore_index: int = IGNORE_INDEX,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build next-token labels and shifted loss mask for Megatron training."""
squeeze = input_ids.dim() == 1
ids = input_ids.unsqueeze(0) if squeeze else input_ids
mask = assistant_loss_mask.unsqueeze(0) if assistant_loss_mask.dim() == 1 else assistant_loss_mask
mask = mask.to(device=ids.device, dtype=torch.float32)
labels = ids.clone()[:, 1:].contiguous()
labels = torch.cat([labels, ignore_index * torch.ones_like(labels[:, :1])], dim=1)
if skipped_tokens is not None and skipped_tokens.numel() > 0:
labels = labels.masked_fill(torch.isin(labels, skipped_tokens.to(device=labels.device)), ignore_index)
shifted_loss_mask = torch.cat([mask[:, 1:], torch.zeros_like(mask[:, :1])], dim=1)
labels = labels.masked_fill(shifted_loss_mask == 0, ignore_index)
shifted_loss_mask = shifted_loss_mask.masked_fill(labels == ignore_index, 0)
if squeeze:
return labels[0].contiguous(), shifted_loss_mask[0].contiguous()
return labels.contiguous(), shifted_loss_mask.contiguous()
def apply_assistant_labels_to_batch(
batch: MutableMapping[str, Any],
examples: Sequence[Mapping[str, Any]],
processor: Any,
skipped_tokens: torch.Tensor,
*,
boundary_config: AssistantMaskBoundaryConfig | None = None,
unmask_last_token: bool = False,
) -> None:
"""Attach ``labels`` and ``loss_mask`` to a collated HF VLM batch."""
loss_masks = [
build_assistant_loss_mask(
example,
input_ids,
processor,
skipped_tokens,
boundary_config=boundary_config,
)
for example, input_ids in zip(examples, batch["input_ids"])
]
loss_mask_t = torch.stack(loss_masks).to(device=batch["input_ids"].device, dtype=torch.float32)
if unmask_last_token and loss_mask_t.numel() > 0:
loss_mask_t[:, -1] = 1.0
labels, shifted_loss_mask = build_shifted_labels_and_loss_mask(batch["input_ids"], loss_mask_t, skipped_tokens)
batch["labels"] = labels
batch["loss_mask"] = shifted_loss_mask
def ensure_position_ids(batch: MutableMapping[str, Any]) -> None:
"""Ensure a collated batch has 2D position IDs."""
if "position_ids" in batch or "input_ids" not in batch:
return
batch_size, seq_len = batch["input_ids"].shape
batch["position_ids"] = (
torch.arange(seq_len, device=batch["input_ids"].device)
.unsqueeze(0)
.expand(batch_size, -1)
.clone()
.contiguous()
)
def pop_generic_visual_inputs(
batch: MutableMapping[str, Any],
visual_keys: Sequence[str],
) -> GenericVisualInputs | None:
"""Move processor visual tensor keys into ``GenericVisualInputs``."""
visual_kwargs = {}
for key in visual_keys:
if key in batch:
visual_kwargs[key] = batch.pop(key)
return GenericVisualInputs(**visual_kwargs) if visual_kwargs else None