`bridge.models.bagel.data.energon`#

Module Contents#

Classes#

`BagelSample`	BAGEL sample ready for packing.
`BagelT2ITaskEncoder`	Register configured BAGEL T2I sample processing.
`BagelEditingTaskEncoder`	Register configured BAGEL Editing sample processing.
`BagelVLMTaskEncoder`	Register configured BAGEL VLM sample processing.

Functions#

`_load_metadata`	Load one WebDataset JSON member.
`_decode_rgb`	Decode image bytes with BAGEL’s transparent-image handling.
`cook_bagel_t2i_sample`	Apply BAGEL’s T2I image, caption, and sequence-plan processing.
`cook_bagel_editing_sample`	Apply BAGEL’s Editing path choice, transforms, tokens, and sequence plan.
`cook_bagel_vlm_sample`	Apply BAGEL’s VLM image, conversation, token, and sequence-plan processing.

API#

class bridge.models.bagel.data.energon.BagelSample#

Bases: megatron.energon.Sample

BAGEL sample ready for packing.

image_tensor_list: list[torch.Tensor]#: None

text_ids_list: list[list[int]]#: None

num_tokens: int#: None

sequence_plan: list[dict[str, object]]#: None

metadata: dict[str, object]#: None

bridge.models.bagel.data.energon._load_metadata(crude_sample: dict[str, object]) → dict[str, object]#: Load one WebDataset JSON member.

bridge.models.bagel.data.energon._decode_rgb(image_bytes: bytes) → PIL.Image.Image#: Decode image bytes with BAGEL’s transparent-image handling.

bridge.models.bagel.data.energon.cook_bagel_t2i_sample( crude_sample: dict[str, object], *, tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], image_stride: int, ) → bridge.models.bagel.data.energon.BagelSample#: Apply BAGEL’s T2I image, caption, and sequence-plan processing.

bridge.models.bagel.data.energon.cook_bagel_editing_sample( crude_sample: dict[str, object], *, tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], vit_transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], image_stride: int, vit_image_stride: int, ) → bridge.models.bagel.data.energon.BagelSample#: Apply BAGEL’s Editing path choice, transforms, tokens, and sequence plan.

bridge.models.bagel.data.energon.cook_bagel_vlm_sample( crude_sample: dict[str, object], *, tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image, int], torch.Tensor], image_stride: int, ) → bridge.models.bagel.data.energon.BagelSample#: Apply BAGEL’s VLM image, conversation, token, and sequence-plan processing.

class bridge.models.bagel.data.energon.BagelT2ITaskEncoder( tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], image_stride: int, )#

Bases: megatron.energon.TaskEncoder

Initialization

Configure the official tokenizer and image transform.

decoder#: None

class bridge.models.bagel.data.energon.BagelEditingTaskEncoder( tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], vit_transform: collections.abc.Callable[[PIL.Image.Image], torch.Tensor], image_stride: int, vit_image_stride: int, )#

Bases: megatron.energon.TaskEncoder

Initialization

Configure the official tokenizer and image transforms.

decoder#: None

class bridge.models.bagel.data.energon.BagelVLMTaskEncoder( tokenizer: transformers.PreTrainedTokenizerBase, transform: collections.abc.Callable[[PIL.Image.Image, int], torch.Tensor], image_stride: int, )#

Bases: megatron.energon.TaskEncoder

Initialization

Configure the official tokenizer and image transform.

decoder#: None

bridge.models.bagel.data.energon#

Module Contents#

Classes#

Functions#

API#

`bridge.models.bagel.data.energon`#