bridge.data.energon.hf_task_encoder#

Generic HF VLM task encoder for Energon dataloading.

Normalizes Energon ChatMLSample objects into HF-style multimodal examples and delegates tokenization, vision preprocessing, masking, and padding to the selected HF VLM collate function.

Module Contents#

Classes#

HFEnergonSample

HF-style VLM example produced from an Energon ChatMLSample.

HFEnergonBatch

Batched format for a generic HF VLM.

HFTaskEncoder

Task encoder for HF VLMs that rely on processor() for tokenization + vision.

API#

class bridge.data.energon.hf_task_encoder.HFEnergonSample#

HF-style VLM example produced from an Energon ChatMLSample.

__key__: str#

None

__subflavors__: Dict#

None

example: Dict[str, Any]#

None

class bridge.data.energon.hf_task_encoder.HFEnergonBatch#

Bases: megatron.energon.Batch

Batched format for a generic HF VLM.

__keys__: List[str]#

‘field(…)’

__subflavors__: List[Dict]#

‘field(…)’

input_ids: torch.Tensor#

‘field(…)’

labels: torch.Tensor#

‘field(…)’

loss_mask: torch.Tensor#

‘field(…)’

position_ids: torch.Tensor#

‘field(…)’

visual_inputs: megatron.bridge.training.utils.visual_inputs.GenericVisualInputs | None#

None

attention_mask: torch.Tensor | None#

None

cu_seqlens_q: torch.Tensor | None#

None

cu_seqlens_kv: torch.Tensor | None#

None

cu_seqlens_q_padded: torch.Tensor | None#

None

cu_seqlens_kv_padded: torch.Tensor | None#

None

max_seqlen_q: torch.Tensor | None#

None

max_seqlen_kv: torch.Tensor | None#

None

class bridge.data.energon.hf_task_encoder.HFTaskEncoder(
processor,
seq_length: int = 4096,
visual_keys: Sequence[str] = ('pixel_values',),
min_pixels: Optional[int] = None,
max_pixels: Optional[int] = None,
collate_fn: Callable[..., dict[str, Any]] | None = None,
pad_to_max_length: bool = False,
pad_to_multiple_of: int = 128,
enable_in_batch_packing: bool = False,
in_batch_packing_pad_to_multiple_of: int = 1,
)#

Bases: megatron.energon.DefaultTaskEncoder[megatron.bridge.data.energon.task_encoder_utils.ChatMLSample, bridge.data.energon.hf_task_encoder.HFEnergonSample, bridge.data.energon.hf_task_encoder.HFEnergonBatch, dict]

Task encoder for HF VLMs that rely on processor() for tokenization + vision.

Parameters:

  • processor – HF AutoProcessor instance passed to the selected collate function.

  • seq_length – Maximum sequence length accepted after collation.

  • visual_keys – Processor output keys to retain when the selected collate function supports configurable visual input selection.

  • min_pixels – Optional min pixel constraint forwarded when supported by the selected collate function.

  • max_pixels – Optional max pixel constraint forwarded when supported by the selected collate function.

  • collate_fn – Optional collate implementation override. If omitted, the implementation is selected from the processor type.

  • pad_to_max_length – Whether collate-time padding should pad non-packed batches to seq_length when the selected collate supports it.

  • pad_to_multiple_of – Non-packed collate-time padding multiple used when pad_to_max_length is false and the selected collate supports it.

  • enable_in_batch_packing – Whether the selected collate should do in-batch sequence packing.

  • in_batch_packing_pad_to_multiple_of – Per-sample padding multiple used only by the in-batch packed path, typically to satisfy CP/SP divisibility.

Initialization

encode_sample(
sample: megatron.bridge.data.energon.task_encoder_utils.ChatMLSample,
) → bridge.data.energon.hf_task_encoder.HFEnergonSample#

Normalize a single ChatML sample into a HF-style collate example.

Expected input format: sample is an Energon ChatMLSample with JSON string conversation plus optional WDS-decoded imgs and videos.

Output format: Returns HFEnergonSample whose example follows the same dictionary schema consumed by HF VLM dataset collate functions. Tokenization, processor calls, label construction, and visual tensor batching are deferred to self.collate_fn.

collate_fn(
examples: list[dict[str, Any]],
) → dict[str, Any]#

Collate HF-style examples with this encoder’s model collator.

Expected input format: List of HF-style VLM example dictionaries with conversation and optional modality fields.

Output format: The exact batch dictionary returned by the selected HF collate function for this processor type.

batch(
samples: List[bridge.data.energon.hf_task_encoder.HFEnergonSample],
) → bridge.data.energon.hf_task_encoder.HFEnergonBatch#

Collate normalized samples with the selected HF VLM collator.

encode_batch(
batch: bridge.data.energon.hf_task_encoder.HFEnergonBatch,
) → dict#

Convert batch dataclass to dict without expanding visual_inputs.