`bridge.models.bagel.data.dataset`#

Declarative BAGEL WDS/Energon training dataset.

Module Contents#

Classes#

BagelDatasetConfig

Build the validated raw-data→WDS→Energon→BAGEL packing chain.

Data#

API#

bridge.models.bagel.data.dataset.logger#: ‘getLogger(…)’

class bridge.models.bagel.data.dataset.BagelDatasetConfig#

Bases: megatron.bridge.data.base.DatasetProvider

Build the validated raw-data→WDS→Energon→BAGEL packing chain.

dataset_root: str | None#: None

bagel_repo: str | None#: None

tokenizer_model: str | None#: None

seed: int#: 42

data_seed: int#: 42

t2i_num_used_data: int#: 10

editing_num_used_data: int#: 10

vlm_num_used_data: int#: 1000

expected_num_tokens: int#: 32768

max_num_tokens_per_sample: int#: 16384

max_num_tokens: int#: 36864

prefer_buffer_before: int#: 16384

max_buffer_size: int#: 50

max_latent_size: int#: 64

text_cond_dropout_prob: float#: 0.1

vit_cond_dropout_prob: float#: 0.4

vae_cond_dropout_prob: float#: 0.1

dataloader_type: Literal[external]#: ‘external’

num_workers: int#: 0

pin_memory: bool#: False

persistent_workers: bool#: False

dataloader_save: str | None#: None

dataloader_load: str | None#: None

static _raw_dataset( path: pathlib.Path, task_encoder: object, worker_config: megatron.energon.WorkerConfig, ) → object#: Build an Energon dataset used through deterministic restore keys.

build_datasets( context: megatron.bridge.data.base.DatasetBuildContext, ) → tuple[megatron.bridge.models.bagel.data.external.BagelExternalLoader | None, None, None]#: Build one packed external loader for this DP rank.