`bridge.data.conversation_processing`#

Shared conversation rendering, masking, and multimodal normalization.

Module Contents#

Classes#

`AssistantMaskBoundaryConfig`	Token boundary fallback configuration for assistant loss masks.
`NormalizedVLMSample`	Source-normalized VLM sample consumed by shared processing.
`TokenizedConversation`	Token ids and an unshifted assistant-only loss mask for one chat row.

Functions#

`normalize_energon_vlm_sample`	Normalize an Energon `ChatMLSample` into the shared VLM sample contract.
`normalize_hf_vlm_example`	Normalize a HF-style VLM dataset example into the shared sample contract.
`normalized_vlm_sample_to_hf_example`	Convert a normalized VLM sample into the HF-style collate example schema.
`_get_explicit_attribute`	Read a declared attribute without triggering dynamic `__getattr__` fabrication.
`get_processor_tokenizer`	Return the HF-style tokenizer attached to a processor or Megatron tokenizer wrapper.
`_as_token_id_list`
`_conversation_from_example`
`chat_template_kwargs_from_example`	Return optional HF chat-template kwargs stored alongside a conversation.
`shared_chat_template_kwargs_from_examples`	Return dataset-wide chat-template kwargs shared by a processor batch.
`normalize_chat_conversation`	Normalize supported text-chat schemas to OpenAI-style messages.
`is_text_only_chat_example`	Return whether a supported chat row contains no media content parts.
`_chat_template_input_ids`
`_apply_tokenized_chat_template`	Apply one chat template with compatibility fallbacks.
`_common_token_prefix_length`	Return the number of identical leading token IDs.
`_rendered_window_start`	Locate a template-truncated token window in its untruncated rendering.
`_infer_final_assistant_start`	Map the final assistant-turn boundary into a possibly truncated render.
`_infer_terminal_assistant_content_start`	Map terminal assistant content, excluding its rendered role header.
`tokenize_chat_example`	Render, tokenize, and construct the assistant mask for one chat row.
`apply_chat_loss_mode`	Apply an explicit chat loss policy to one unpadded token sequence.
`find_token_span`	Find the first `[start, end)` token span matching `pattern`.
`gather_assistant_text_segments`	Extract assistant text segments from a structured VLM conversation.
`tokenize_text_without_special_tokens`	Tokenize text using a HF-like tokenizer without adding special tokens.
`infer_assistant_mask_boundary_config`	Infer assistant role boundaries from known HF chat-template formats.
`assistant_mask_boundary_config_from_markers`	Build explicit assistant boundary config from model-declared marker strings.
`_apply_skipped_tokens`
`_warn_if_all_masked`
`_get_chat_template`
`_assistant_mask_from_hf_chat_template`	Return HF generation-block assistant mask when it aligns to `ids`.
`_align_assistant_mask_to_padded_ids`	Align an HF assistant mask to possibly padded processor batch ids.
`_token_map_from_boundary_config`
`_trim_leading_token_ids`
`_trim_leading_token_sequences`	Skip complete token sequences that appear contiguously at a span start.
`_assistant_mask_from_boundary_config`
`build_assistant_loss_mask`	Build an unshifted assistant-only loss mask.
`build_shifted_labels_and_loss_mask`	Build next-token labels and shifted loss mask for Megatron training.
`apply_assistant_labels_to_batch`	Attach `labels` and `loss_mask` to a collated HF VLM batch.
`ensure_position_ids`	Ensure a collated batch has 2D position IDs.
`pop_generic_visual_inputs`	Move processor visual tensor keys into `GenericVisualInputs`.

Data#

CHATML_ASSISTANT_ROLE

API#

bridge.data.conversation_processing.CHATML_ASSISTANT_ROLE#: ‘assistant’

class bridge.data.conversation_processing.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.

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: collections.abc.Mapping[str, collections.abc.Sequence[int]]#: None

role_end_tokens: collections.abc.Mapping[str, collections.abc.Sequence[int]]#: None

loss_roles: collections.abc.Sequence[str]#: ()

include_start_tokens_for_roles: collections.abc.Sequence[str]#: ()

include_end_tokens_for_roles: collections.abc.Sequence[str]#: ()

trim_leading_token_ids: collections.abc.Sequence[int]#: ()

trim_leading_token_sequences: collections.abc.Sequence[collections.abc.Sequence[int]]#: ()

role_end_token_variants: collections.abc.Mapping[str, collections.abc.Sequence[collections.abc.Sequence[int]]]#: ‘field(…)’

class bridge.data.conversation_processing.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]]#: None

images: list[Any] | None#: None

videos: list[Any] | None#: None

audio: Any | None#: None

tools: list[dict[str, Any]] | None#: None

class bridge.data.conversation_processing.TokenizedConversation#

Token ids and an unshifted assistant-only loss mask for one chat row.

input_ids: torch.Tensor#: None

assistant_mask: torch.Tensor#: None

conversation: list[dict[str, Any]]#: None

final_assistant_start: int | None#: None

bridge.data.conversation_processing.normalize_energon_vlm_sample( sample: Any, ) → bridge.data.conversation_processing.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.

bridge.data.conversation_processing.normalize_hf_vlm_example( example: collections.abc.Mapping[str, Any], ) → bridge.data.conversation_processing.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.

bridge.data.conversation_processing.normalized_vlm_sample_to_hf_example( sample: bridge.data.conversation_processing.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.

bridge.data.conversation_processing._get_explicit_attribute( obj: Any, name: str, ) → Any | None#: Read a declared attribute without triggering dynamic __getattr__ fabrication.

bridge.data.conversation_processing.get_processor_tokenizer(processor: Any) → Any#: Return the HF-style tokenizer attached to a processor or Megatron tokenizer wrapper.

bridge.data.conversation_processing._as_token_id_list(token_ids: Any) → list[int]#

bridge.data.conversation_processing._conversation_from_example( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], ) → collections.abc.Sequence[collections.abc.Mapping[str, Any]]#

bridge.data.conversation_processing.chat_template_kwargs_from_example( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], ) → dict[str, Any]#: Return optional HF chat-template kwargs stored alongside a conversation.

bridge.data.conversation_processing.shared_chat_template_kwargs_from_examples( examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], ) → dict[str, Any]#

Return dataset-wide chat-template kwargs shared by a processor batch.

Processor-batched VLM collators accept one tool schema for the whole microbatch. Datasets using these collators must therefore use a shared tool registry instead of varying tool definitions per row.

Raises:: ValueError – If rows in one processor batch use different template kwargs.

bridge.data.conversation_processing.normalize_chat_conversation( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.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.

bridge.data.conversation_processing.is_text_only_chat_example( example: collections.abc.Mapping[str, Any], ) → bool#: Return whether a supported chat row contains no media content parts.

bridge.data.conversation_processing._chat_template_input_ids(tokenized_chat: Any) → list[int]#

bridge.data.conversation_processing._apply_tokenized_chat_template( template_owner: Any, conversation: collections.abc.Sequence[collections.abc.Mapping[str, Any]], tokenize_kwargs: collections.abc.Mapping[str, Any], ) → Any#: Apply one chat template with compatibility fallbacks.

bridge.data.conversation_processing._common_token_prefix_length( left: collections.abc.Sequence[int], right: collections.abc.Sequence[int], ) → int#: Return the number of identical leading token IDs.

bridge.data.conversation_processing._rendered_window_start( full_ids: collections.abc.Sequence[int], rendered_ids: collections.abc.Sequence[int], template_owner: Any, ) → int | None#: Locate a template-truncated token window in its untruncated rendering.

bridge.data.conversation_processing._infer_final_assistant_start( template_owner: Any, conversation: collections.abc.Sequence[collections.abc.Mapping[str, Any]], tokenize_kwargs: collections.abc.Mapping[str, Any], rendered_ids: collections.abc.Sequence[int], *, terminal_only: bool, ) → int | None#: Map the final assistant-turn boundary into a possibly truncated render.

bridge.data.conversation_processing._infer_terminal_assistant_content_start( template_owner: Any, conversation: collections.abc.Sequence[collections.abc.Mapping[str, Any]], tokenize_kwargs: collections.abc.Mapping[str, Any], rendered_ids: collections.abc.Sequence[int], ) → int | None#: Map terminal assistant content, excluding its rendered role header.

bridge.data.conversation_processing.tokenize_chat_example( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], processor: Any, *, tool_schemas: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]] | None = None, max_length: int | None = None, skipped_tokens: torch.Tensor | None = None, boundary_config: bridge.data.conversation_processing.AssistantMaskBoundaryConfig | None = None, warn_on_all_masked: bool = True, loss_mode: Literal[assistant, last_turn, full] = 'assistant', return_final_assistant_start: bool = False, ) → bridge.data.conversation_processing.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.

bridge.data.conversation_processing.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.

Parameters:

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.

bridge.data.conversation_processing.find_token_span( sequence: collections.abc.Sequence[int] | torch.Tensor, pattern: collections.abc.Sequence[int], start: int = 0, ) → tuple[int, int]#

Find the first [start, end) token span matching pattern.

Parameters:

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.

bridge.data.conversation_processing.gather_assistant_text_segments( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], ) → list[str]#: Extract assistant text segments from a structured VLM conversation.

bridge.data.conversation_processing.tokenize_text_without_special_tokens( tokenizer: Any, text: str, ) → list[int]#: Tokenize text using a HF-like tokenizer without adding special tokens.

bridge.data.conversation_processing.infer_assistant_mask_boundary_config( processor: Any, ) → bridge.data.conversation_processing.AssistantMaskBoundaryConfig | None#: Infer assistant role boundaries from known HF chat-template formats.

bridge.data.conversation_processing.assistant_mask_boundary_config_from_markers( processor: Any, *, assistant_start: str, assistant_end: str, assistant_end_fallbacks: collections.abc.Sequence[str] = (), role_start_markers: collections.abc.Mapping[str, str] | None = None, loss_roles: collections.abc.Sequence[str] = (CHATML_ASSISTANT_ROLE,), include_start_tokens_for_roles: collections.abc.Sequence[str] = (), include_end_tokens_for_roles: collections.abc.Sequence[str] = (CHATML_ASSISTANT_ROLE,), trim_leading_token_ids: collections.abc.Sequence[int] = (), trim_leading_token_sequences: collections.abc.Sequence[collections.abc.Sequence[int]] = (), ) → bridge.data.conversation_processing.AssistantMaskBoundaryConfig#: Build explicit assistant boundary config from model-declared marker strings.

bridge.data.conversation_processing._apply_skipped_tokens( mask: torch.Tensor, ids: collections.abc.Sequence[int], skipped_tokens: torch.Tensor | None, ) → None#

bridge.data.conversation_processing._warn_if_all_masked( mask: torch.Tensor, example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], *, warn_on_all_masked: bool, ) → None#

bridge.data.conversation_processing._get_chat_template(*objects: Any) → str | None#

bridge.data.conversation_processing._assistant_mask_from_hf_chat_template( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], ids: collections.abc.Sequence[int], processor: Any, tokenizer: Any, ) → torch.Tensor | None#: Return HF generation-block assistant mask when it aligns to ids.

bridge.data.conversation_processing._align_assistant_mask_to_padded_ids( candidate_ids: collections.abc.Sequence[int], assistant_mask: collections.abc.Sequence[int], ids: collections.abc.Sequence[int], tokenizer: Any, ) → list[int] | None#: Align an HF assistant mask to possibly padded processor batch ids.

bridge.data.conversation_processing._token_map_from_boundary_config( token_map: collections.abc.Mapping[str, collections.abc.Sequence[int]] | None, ) → dict[str, list[int]]#

bridge.data.conversation_processing._trim_leading_token_ids( ids: collections.abc.Sequence[int], start: int, end: int, trim_token_ids: collections.abc.Sequence[int], ) → int#

bridge.data.conversation_processing._trim_leading_token_sequences( ids: collections.abc.Sequence[int], start: int, end: int, trim_token_sequences: collections.abc.Sequence[collections.abc.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.

Parameters:

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.

bridge.data.conversation_processing._assistant_mask_from_boundary_config( ids: collections.abc.Sequence[int], boundary_config: bridge.data.conversation_processing.AssistantMaskBoundaryConfig, *, include_content: bool = True, ) → torch.Tensor | None#

bridge.data.conversation_processing.build_assistant_loss_mask( example_or_conversation: collections.abc.Mapping[str, Any] | collections.abc.Sequence[collections.abc.Mapping[str, Any]], input_ids: collections.abc.Sequence[int] | torch.Tensor, processor: Any, skipped_tokens: torch.Tensor | None = None, *, boundary_config: bridge.data.conversation_processing.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.

bridge.data.conversation_processing.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.

bridge.data.conversation_processing.apply_assistant_labels_to_batch( batch: collections.abc.MutableMapping[str, Any], examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], processor: Any, skipped_tokens: torch.Tensor, *, boundary_config: bridge.data.conversation_processing.AssistantMaskBoundaryConfig | None = None, unmask_last_token: bool = False, ) → None#: Attach labels and loss_mask to a collated HF VLM batch.

bridge.data.conversation_processing.ensure_position_ids( batch: collections.abc.MutableMapping[str, Any], ) → None#: Ensure a collated batch has 2D position IDs.

bridge.data.conversation_processing.pop_generic_visual_inputs( batch: collections.abc.MutableMapping[str, Any], visual_keys: collections.abc.Sequence[str], ) → megatron.bridge.training.utils.visual_inputs.GenericVisualInputs | None#: Move processor visual tensor keys into GenericVisualInputs.

bridge.data.conversation_processing#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.data.conversation_processing`#