update format

Author: zhaoyinglia

flagscale/train/datasets/energon/cooker.py

@@ -1,12 +1,10 @@
-
 import json
 
 from megatron.energon import stateless
 from megatron.energon.task_encoder.cooking import Cooker
 
-
-IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.webp')
-VIDEO_EXTENSIONS = ('.mp4', '.avi', '.mov', '.webm', '.mkv', '.flv')
+IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".webp")
+VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".webm", ".mkv", ".flv")
 
 
 @stateless
@@ -17,22 +15,22 @@ def cook_bagel_images_sample(sample: dict) -> dict:
       000000000.000.jpg  -> sample['000.jpg'] (PIL Image or bytes)
       000000000.json     -> sample['json'] (decoded dict)
     """
-    json_data = sample.get('json', {})
+    json_data = sample.get("json", {})
     if isinstance(json_data, (bytes, str)):
         json_data = json.loads(json_data)
 
     # Collect all image keys (e.g., '000.jpg', '000.png', etc.)
     images = []
     for key, value in sample.items():
-        if key.startswith('__'):
+        if key.startswith("__"):
             continue
         if any(key.endswith(ext) for ext in IMAGE_EXTENSIONS):
             images.append(value)
 
     result = {
-        **{k: v for k, v in sample.items() if k.startswith('__')},  # preserve meta keys
-        'json_data': json_data,
-        'images': images,
+        **{k: v for k, v in sample.items() if k.startswith("__")},  # preserve meta keys
+        "json_data": json_data,
+        "images": images,
     }
     return result
 
@@ -54,7 +52,7 @@ def cook_bagel_video_sample(sample: dict) -> dict:
         'video_bytes': bytes,   # raw video bytes for FrameSampler
       }
     """
-    json_data = sample.get('json', {})
+    json_data = sample.get("json", {})
     if isinstance(json_data, (bytes, str)):
         json_data = json.loads(json_data)
 
@@ -63,11 +61,11 @@ def cook_bagel_video_sample(sample: dict) -> dict:
     video_bytes = None
     images = []
     for key, value in sample.items():
-        if key.startswith('__'):
+        if key.startswith("__"):
             continue
         if any(key.endswith(ext) for ext in VIDEO_EXTENSIONS):
             # AVDecoder object → extract raw bytes for decord
-            if hasattr(value, 'stream'):
+            if hasattr(value, "stream"):
                 value.stream.seek(0)
                 video_bytes = value.stream.read()
                 value.stream.seek(0)
@@ -80,10 +78,10 @@ def cook_bagel_video_sample(sample: dict) -> dict:
             images.append(value)
 
     result = {
-        **{k: v for k, v in sample.items() if k.startswith('__')},  # preserve meta keys
-        'json_data': json_data,
-        'images': images,
-        'video_bytes': video_bytes,
+        **{k: v for k, v in sample.items() if k.startswith("__")},  # preserve meta keys
+        "json_data": json_data,
+        "images": images,
+        "video_bytes": video_bytes,
     }
     return result
 

flagscale/train/datasets/energon/data_utils.py

@@ -1,17 +1,13 @@
 # Copyright (c) 2025, BAAI. All rights reserved.
 # Utility functions for Bagel data processing, ported from Bagel's data_utils.py.
 
-import math
-import random
-from PIL import Image, ImageFile, PngImagePlugin
-
 import torch
-
+from PIL import Image, ImageFile, PngImagePlugin
 
 Image.MAX_IMAGE_PIXELS = 200000000
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 MaximumDecompressedSize = 1024
-MegaByte = 2 ** 20
+MegaByte = 2**20
 PngImagePlugin.MAX_TEXT_CHUNK = MaximumDecompressedSize * MegaByte
 
 
@@ -52,20 +48,22 @@ def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"):
 
     csum = 0
     for s, attn_mode in zip(split_lens, attn_modes):
-        assert attn_mode in ['causal', 'full', 'noise']
+        assert attn_mode in ["causal", "full", "noise"]
         if attn_mode == "causal":
-            attention_mask[csum:csum + s, csum:csum + s] = torch.ones((s, s), device=device).tril()
-            attention_mask[csum:csum + s, :csum] = 1
+            attention_mask[csum : csum + s, csum : csum + s] = torch.ones(
+                (s, s), device=device
+            ).tril()
+            attention_mask[csum : csum + s, :csum] = 1
         else:
-            attention_mask[csum:csum + s, csum:csum + s] = torch.ones((s, s))
-            attention_mask[csum:csum + s, :csum] = 1
+            attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s))
+            attention_mask[csum : csum + s, :csum] = 1
         csum += s
 
     csum = 0
     for s, attn_mode in zip(split_lens, attn_modes):
         if attn_mode == "noise":
-            attention_mask[:, csum:csum + s] = torch.zeros((sample_len, s))
-            attention_mask[csum:csum + s, csum:csum + s] = torch.ones((s, s))
+            attention_mask[:, csum : csum + s] = torch.zeros((sample_len, s))
+            attention_mask[csum : csum + s, csum : csum + s] = torch.ones((s, s))
         csum += s
 
     attention_mask = torch.zeros_like(attention_mask, dtype=torch.float).masked_fill_(
@@ -74,15 +72,15 @@ def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"):
     return attention_mask
 
 
-def len2weight(x, loss_reduction='square'):
+def len2weight(x, loss_reduction="square"):
     if x == 0:
         return x
-    if loss_reduction == 'token':
+    if loss_reduction == "token":
         return 1
-    if loss_reduction == 'sample':
+    if loss_reduction == "sample":
         return 1 / x
-    if loss_reduction == 'square':
-        return 1 / (x ** 0.5)
+    if loss_reduction == "square":
+        return 1 / (x**0.5)
     raise NotImplementedError(loss_reduction)
 
 
@@ -107,15 +105,15 @@ def add_special_tokens(tokenizer):
             all_special_tokens += v
 
     new_tokens = []
-    for token in ['<|im_start|>', '<|im_end|>', '<|vision_start|>', '<|vision_end|>']:
+    for token in ["<|im_start|>", "<|im_end|>", "<|vision_start|>", "<|vision_end|>"]:
         if token not in all_special_tokens:
             new_tokens.append(token)
 
     num_new_tokens = tokenizer.add_tokens(new_tokens)
-    bos_token_id = tokenizer.convert_tokens_to_ids('<|im_start|>')
-    eos_token_id = tokenizer.convert_tokens_to_ids('<|im_end|>')
-    start_of_image = tokenizer.convert_tokens_to_ids('<|vision_start|>')
-    end_of_image = tokenizer.convert_tokens_to_ids('<|vision_end|>')
+    bos_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
+    eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+    start_of_image = tokenizer.convert_tokens_to_ids("<|vision_start|>")
+    end_of_image = tokenizer.convert_tokens_to_ids("<|vision_end|>")
 
     new_token_ids = dict(
         bos_token_id=bos_token_id,

flagscale/train/datasets/energon/sample_types.py

@@ -2,7 +2,7 @@
 # Sample type definitions for Bagel-Energon integration.
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 import torch
 
@@ -13,12 +13,13 @@ class BagelSample:
 
     This is what goes into the PackingDataset buffer.
     """
-    image_tensor_list: List[torch.Tensor]  # ViT/VAE preprocessed images
-    text_ids_list: List[List[int]]         # Tokenized text segments
-    sequence_plan: List[Dict[str, Any]]    # Packing instructions per segment
-    num_tokens: int                        # Total token count for this sample
-    is_mandatory: bool = False             # Whether this sample must appear in every pack
-    subflavor: str = ""                    # Task type: "t2i" / "vlm"
+
+    image_tensor_list: list[torch.Tensor]  # ViT/VAE preprocessed images
+    text_ids_list: list[list[int]]  # Tokenized text segments
+    sequence_plan: list[dict[str, Any]]  # Packing instructions per segment
+    num_tokens: int  # Total token count for this sample
+    is_mandatory: bool = False  # Whether this sample must appear in every pack
+    subflavor: str = ""  # Task type: "t2i" / "vlm"
     # Energon metadata
     __key__: str = ""
     __restore_key__: tuple = ()
@@ -27,58 +28,59 @@ class BagelSample:
 @dataclass
 class BagelPackedBatch:
     """Output of pack_selected_samples — a fully packed sequence ready for the model."""
+
     sequence_length: int
-    sample_lens: List[int]
+    sample_lens: list[int]
     packed_text_ids: torch.Tensor
     packed_text_indexes: torch.Tensor
     packed_position_ids: torch.Tensor
     # FlexAttention fields
-    split_lens: List[int] = field(default_factory=list)
-    attn_modes: List[str] = field(default_factory=list)
+    split_lens: list[int] = field(default_factory=list)
+    attn_modes: list[str] = field(default_factory=list)
     # VAE image generation (optional)
-    padded_images: Optional[torch.Tensor] = None
-    patchified_vae_latent_shapes: Optional[List] = None
-    packed_latent_position_ids: Optional[torch.Tensor] = None
-    packed_vae_token_indexes: Optional[torch.Tensor] = None
+    padded_images: torch.Tensor | None = None
+    patchified_vae_latent_shapes: list | None = None
+    packed_latent_position_ids: torch.Tensor | None = None
+    packed_vae_token_indexes: torch.Tensor | None = None
     # ViT image understanding (optional)
-    packed_vit_tokens: Optional[torch.Tensor] = None
-    packed_vit_position_ids: Optional[torch.Tensor] = None
-    packed_vit_token_indexes: Optional[torch.Tensor] = None
-    vit_token_seqlens: Optional[torch.Tensor] = None
+    packed_vit_tokens: torch.Tensor | None = None
+    packed_vit_position_ids: torch.Tensor | None = None
+    packed_vit_token_indexes: torch.Tensor | None = None
+    vit_token_seqlens: torch.Tensor | None = None
     # Diffusion timesteps (optional)
-    packed_timesteps: Optional[torch.Tensor] = None
-    mse_loss_indexes: Optional[torch.Tensor] = None
+    packed_timesteps: torch.Tensor | None = None
+    mse_loss_indexes: torch.Tensor | None = None
     # CE loss (optional)
-    packed_label_ids: Optional[torch.Tensor] = None
-    ce_loss_indexes: Optional[torch.Tensor] = None
-    ce_loss_weights: Optional[torch.Tensor] = None
+    packed_label_ids: torch.Tensor | None = None
+    ce_loss_indexes: torch.Tensor | None = None
+    ce_loss_weights: torch.Tensor | None = None
 
     def to_dict(self):
         """Convert to dict for get_batch compatibility."""
         data = {
-            'sequence_length': self.sequence_length,
-            'sample_lens': self.sample_lens,
-            'packed_text_ids': self.packed_text_ids,
-            'packed_text_indexes': self.packed_text_indexes,
-            'packed_position_ids': self.packed_position_ids,
-            'split_lens': self.split_lens,
-            'attn_modes': self.attn_modes,
+            "sequence_length": self.sequence_length,
+            "sample_lens": self.sample_lens,
+            "packed_text_ids": self.packed_text_ids,
+            "packed_text_indexes": self.packed_text_indexes,
+            "packed_position_ids": self.packed_position_ids,
+            "split_lens": self.split_lens,
+            "attn_modes": self.attn_modes,
         }
         if self.padded_images is not None:
-            data['padded_images'] = self.padded_images
-            data['patchified_vae_latent_shapes'] = self.patchified_vae_latent_shapes
-            data['packed_latent_position_ids'] = self.packed_latent_position_ids
-            data['packed_vae_token_indexes'] = self.packed_vae_token_indexes
+            data["padded_images"] = self.padded_images
+            data["patchified_vae_latent_shapes"] = self.patchified_vae_latent_shapes
+            data["packed_latent_position_ids"] = self.packed_latent_position_ids
+            data["packed_vae_token_indexes"] = self.packed_vae_token_indexes
         if self.packed_vit_tokens is not None:
-            data['packed_vit_tokens'] = self.packed_vit_tokens
-            data['packed_vit_position_ids'] = self.packed_vit_position_ids
-            data['packed_vit_token_indexes'] = self.packed_vit_token_indexes
-            data['vit_token_seqlens'] = self.vit_token_seqlens
+            data["packed_vit_tokens"] = self.packed_vit_tokens
+            data["packed_vit_position_ids"] = self.packed_vit_position_ids
+            data["packed_vit_token_indexes"] = self.packed_vit_token_indexes
+            data["vit_token_seqlens"] = self.vit_token_seqlens
         if self.packed_timesteps is not None:
-            data['packed_timesteps'] = self.packed_timesteps
-            data['mse_loss_indexes'] = self.mse_loss_indexes
+            data["packed_timesteps"] = self.packed_timesteps
+            data["mse_loss_indexes"] = self.mse_loss_indexes
         if self.packed_label_ids is not None:
-            data['packed_label_ids'] = self.packed_label_ids
-            data['ce_loss_indexes'] = self.ce_loss_indexes
-            data['ce_loss_weights'] = self.ce_loss_weights
+            data["packed_label_ids"] = self.packed_label_ids
+            data["ce_loss_indexes"] = self.ce_loss_indexes
+            data["ce_loss_weights"] = self.ce_loss_weights
         return data

flagscale/train/datasets/energon/task_handlers/base.py

@@ -1,7 +1,4 @@
-
-
 class BaseTaskHandler:
-
     def __init__(self, tokenizer, special_tokens, data_config):
         self.tokenizer = tokenizer
         self.special_tokens = special_tokens

flagscale/train/datasets/energon/task_handlers/t2i.py

@@ -1,22 +1,19 @@
-import io
 import json
 import random
-from PIL import Image
 
 from flagscale.train.datasets.energon.data_utils import pil_img2rgb
 from flagscale.train.datasets.energon.sample_types import BagelSample
 from flagscale.train.datasets.energon.task_handlers.base import BaseTaskHandler
 
 
 class T2IHandler(BaseTaskHandler):
-
     def encode(self, sample, **kwargs):
         """Encode text-to-image sample: caption + VAE image."""
         transform = kwargs.get("transform")
 
-        data_item = sample.get('json_data')
-        images = sample.get('images', [])
-        caption_dict = data_item.get('caption_dict', '')
+        data_item = sample.get("json_data")
+        images = sample.get("images", [])
+        caption_dict = data_item.get("caption_dict", "")
         # print(f"{images=}, {caption_dict=}")
 
         assert images is not None and len(images) == 1
@@ -34,7 +31,7 @@ def encode(self, sample, **kwargs):
         image_tensor = transform(raw_image)
         image_tensor_list.append(image_tensor)
         height, width = image_tensor.shape[1:]
-        num_tokens += width * height // transform_stride ** 2
+        num_tokens += width * height // transform_stride**2
 
         # Tokenize caption
         caption_dict = json.loads(caption_dict)
@@ -46,31 +43,35 @@ def encode(self, sample, **kwargs):
         if len(caption_token) > 0:
             text_ids_list.append(caption_token)
             num_tokens += len(caption_token)
-            sequence_plan.append({
-                'type': 'text',
-                'enable_cfg': 1,
-                'loss': 0,
-                'special_token_loss': 0,
-                'special_token_label': None,
-            })
+            sequence_plan.append(
+                {
+                    "type": "text",
+                    "enable_cfg": 1,
+                    "loss": 0,
+                    "special_token_loss": 0,
+                    "special_token_label": None,
+                }
+            )
 
         # VAE image plan
         if image_tensor_list:
-            sequence_plan.append({
-                'type': 'vae_image',
-                'enable_cfg': 0,
-                'loss': 1,
-                'special_token_loss': 0,
-                'special_token_label': None,
-            })
+            sequence_plan.append(
+                {
+                    "type": "vae_image",
+                    "enable_cfg": 0,
+                    "loss": 1,
+                    "special_token_loss": 0,
+                    "special_token_label": None,
+                }
+            )
 
         return BagelSample(
             image_tensor_list=image_tensor_list,
             text_ids_list=text_ids_list,
             sequence_plan=sequence_plan,
             num_tokens=num_tokens,
-            is_mandatory=sample.get('__subflavors__', {}).get('is_mandatory', False),
-            subflavor=sample.get('__subflavors__', {}).get("task", "vlm"),
-            __key__=sample.get('__key__', ''),
-            __restore_key__=sample.get('__restore_key__', ()),
+            is_mandatory=sample.get("__subflavors__", {}).get("is_mandatory", False),
+            subflavor=sample.get("__subflavors__", {}).get("task", "vlm"),
+            __key__=sample.get("__key__", ""),
+            __restore_key__=sample.get("__restore_key__", ()),
         )