conversation_processing.py

# 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 ``<think></think>``
        wrapper without masking a non-empty ``<think>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 <image>"},
                {"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 ``<image>`` / ``<video>`` 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 ("<image>" not in content and "<video>" not in content):
            return copy.deepcopy(content)

        pieces = re.split(r"(<image>|<video>)", content)
        parts = []
        for piece in pieces:
            if piece == "<image>":
                parts.append(_media_part("image"))
            elif piece == "<video>":
                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 OpenAI message content
    is converted to an empty string before chat-template rendering.
    """
    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()
            normalized.append({"role": _LEGACY_CHAT_ROLE_ALIASES.get(role, role), "content": turn["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 ("<image>", "<video>", "<audio>")):
                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,
) -> 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.")
    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:
        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", "<think>")):
        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|>"), "<|turn>model\n", "<turn|>", ("<|tool_response>",), {}),
        (("<start_of_turn>model", "<end_of_turn>"), "<start_of_turn>model\n", "<end_of_turn>", (), {}),
        (
            ("<|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~[", "<mm:think>"),
            "]~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", "<think>", "</think>")):
                empty_think_tokens = tokenize_text_without_special_tokens(tokenizer, "<think></think>")
                think_open_tokens = tokenize_text_without_special_tokens(tokenizer, "<think>\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 ``<think></think>`` block while preserving a
    non-empty ``<think>reasoning...</think>`` 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