bridge.data.energon.prepare#

Helpers for preparing Bridge-owned Energon WebDatasets.

Module Contents#

Functions#

_relative_tar_paths

_validate_split_patterns

prepare_webdataset

Index WebDataset shards and write deterministic Energon split metadata.

Data#

logger

API#

bridge.data.energon.prepare.logger#

‘getLogger(…)’

bridge.data.energon.prepare._relative_tar_paths(dataset_path: pathlib.Path) → list[str]#
bridge.data.energon.prepare._validate_split_patterns(
paths: list[str],
split_patterns: collections.abc.Mapping[str, str],
) → None#
bridge.data.energon.prepare.prepare_webdataset(
dataset_path: str | pathlib.Path,
split_patterns: collections.abc.Mapping[str, str],
*,
num_workers: int = 8,
) → None#

Index WebDataset shards and write deterministic Energon split metadata.

This calls Energon’s preparation API directly instead of its interactive command-line wrapper. The API is shared by supported Megatron-Energon 7.x releases even though their CLI flags differ.

Parameters:

  • dataset_path – Directory containing .tar or .tgz shards.

  • split_patterns – Mapping from split names to shard-path regular expressions.

  • num_workers – Number of parallel shard-indexing workers.

Raises:

  • FileNotFoundError – If dataset_path does not contain any tar shards.

  • ValueError – If workers or split patterns are invalid.