bridge.data.collators.sft

Shared text-only collators for direct SFT.

Module Contents

Classes

_TokenizedTextSFTRow

Backend-neutral tokenized text row consumed by direct-HF batching.

Functions

_pad_tokenized_rows	Right-pad shared row preprocessing outputs for direct-HF batching.
_metadata_from_example
_build_text_sft_batch	Build one padded or in-batch-packed direct-HF text batch.
text_chat_collate_fn	Collate text-only HF chat examples using the shared assistant-mask path.
text_prompt_completion_collate_fn	Collate prompt-completion rows without applying a chat template.

Data

_CONVERSATION_KEYS

API

bridge.data.collators.sft._CONVERSATION_KEYS

(‘conversation’, ‘messages’, ‘conversations’)

class bridge.data.collators.sft._TokenizedTextSFTRow

Backend-neutral tokenized text row consumed by direct-HF batching.

input_ids: torch.Tensor

None

loss_mask: torch.Tensor

None

bridge.data.collators.sft._pad_tokenized_rows(

rows: list[bridge.data.collators.sft._TokenizedTextSFTRow],

*,

pad_token_id: int,

max_length: int | None,

pad_to_max_length: bool,

pad_to_multiple_of: int,

) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Right-pad shared row preprocessing outputs for direct-HF batching.

bridge.data.collators.sft._metadata_from_example(

example: collections.abc.Mapping[str, Any],

) → dict[str, Any]

bridge.data.collators.sft._build_text_sft_batch(

examples: list[collections.abc.Mapping[str, Any]],

tokenized_rows: list[bridge.data.collators.sft._TokenizedTextSFTRow],

tokenizer: Any,

skipped_tokens: torch.Tensor,

metadata: list[dict[str, Any]],

*,

max_length: int | None,

pad_to_max_length: bool,

pad_to_multiple_of: int,

ignore_index: int,

enable_in_batch_packing: bool,

in_batch_packing_pad_to_multiple_of: int,

) → dict[str, Any]

Build one padded or in-batch-packed direct-HF text batch.

Packed batches use their emergent aggregate width; pad_to_max_length applies only to the non-packed path.

bridge.data.collators.sft.text_chat_collate_fn(

examples: list[collections.abc.Mapping[str, Any]],

processor: Any,

*,

max_length: int | None = None,

sequence_length: int | None = None,

pad_to_max_length: bool = False,

pad_to_multiple_of: int = 1,

warn_on_all_masked: bool = True,

loss_mode: Literal[assistant, last_turn, full] = 'assistant',

ignore_index: int = IGNORE_INDEX,

enable_in_batch_packing: bool = False,

in_batch_packing_pad_to_multiple_of: int = 1,

**kwargs: Any,

) → dict[str, Any]

Collate text-only HF chat examples using the shared assistant-mask path.

Parameters:

• examples – HF-style chat rows containing messages, conversation, or legacy conversations. Optional top-level tools are forwarded to the chat template for rendering and assistant masks.

• processor – A HF tokenizer or processor. It must expose apply_chat_template directly or through processor.tokenizer.

• max_length – Optional tokenizer truncation length.

• sequence_length – Optional tokenizer truncation length used by Direct Hugging Face SFT builders.

• pad_to_max_length – On non-packed batches, pad every row to max_length instead of the longest row. Packed batches use emergent width.

• pad_to_multiple_of – Optional non-packed padding multiple. The HF SFT builder uses this to keep CP/SP slices shape-compatible.

• warn_on_all_masked – Forwarded to assistant-mask construction.

• loss_mode – Chat tokens that contribute to loss.

• ignore_index – Label ignore value for masked targets.

• enable_in_batch_packing – If True, flatten the padded microbatch and emit packed-sequence metadata for GPT-style training steps.

• in_batch_packing_pad_to_multiple_of – Optional per-sequence length multiple used when enable_in_batch_packing inserts padding for CP/SP constraints.

• **kwargs – Additional common collate kwargs accepted for parity with VLM collate functions and ignored by the text-only path.

Returns:

Batch dictionary with VLM-style input_ids and GPT-style tokens aliases, shifted labels and loss_mask, position_ids, and optional tokenizer fields such as attention_mask.

bridge.data.collators.sft.text_prompt_completion_collate_fn(

examples: list[collections.abc.Mapping[str, Any]],

processor: Any,

*,

preprocessing: megatron.bridge.data.sft_processing.PromptCompletionSFTPreprocessingConfig,

max_length: int | None = None,

sequence_length: int | None = None,

pad_to_max_length: bool = False,

pad_to_multiple_of: int = 1,

ignore_index: int = IGNORE_INDEX,

enable_in_batch_packing: bool = False,

in_batch_packing_pad_to_multiple_of: int = 1,

**kwargs: Any,

) → dict[str, Any]

Collate prompt-completion rows without applying a chat template.