Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright 2025 Bytedance Ltd. and/or its affiliates. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import pyarrow.parquet as pq | |
| from ..distributed_iterable_dataset import DistributedIterableDataset | |
| from ..parquet_utils import get_parquet_data_paths, init_arrow_pf_fs | |
| class InterleavedBaseIterableDataset(DistributedIterableDataset): | |
| def _init_data(self): | |
| data = { | |
| 'sequence_plan': [], | |
| 'text_ids_list': [], | |
| 'image_tensor_list': [], | |
| 'num_tokens': 0, | |
| } | |
| return data | |
| def _add_text(self, data, text, need_loss, enable_cfg=True): | |
| text_ids = self.tokenizer.encode(text) | |
| data['num_tokens'] += len(text_ids) | |
| data['text_ids_list'].append(text_ids) | |
| data['sequence_plan'].append( | |
| { | |
| 'type': 'text', | |
| 'enable_cfg': int(enable_cfg), | |
| 'loss': int(need_loss), | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| } | |
| ) | |
| return data | |
| def _add_image(self, data, image, need_loss, need_vae, need_vit, enable_cfg=True): | |
| assert need_loss or need_vae or need_vit | |
| if need_loss: | |
| data['sequence_plan'].append( | |
| { | |
| 'type': 'vae_image', | |
| 'enable_cfg': 0, | |
| 'loss': 1, | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| } | |
| ) | |
| image_tensor = self.transform(image) | |
| height, width = image_tensor.shape[1:] | |
| data['num_tokens'] += width * height // self.transform.stride ** 2 | |
| data['image_tensor_list'].append(image_tensor) | |
| if need_vae: | |
| data['sequence_plan'].append( | |
| { | |
| 'type': 'vae_image', | |
| 'enable_cfg': int(enable_cfg), | |
| 'loss': 0, | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| } | |
| ) | |
| image_tensor = self.transform(image) | |
| height, width = image_tensor.shape[1:] | |
| data['num_tokens'] += width * height // self.transform.stride ** 2 | |
| data['image_tensor_list'].append(image_tensor.clone()) | |
| if need_vit: | |
| data['sequence_plan'].append( | |
| { | |
| 'type': 'vit_image', | |
| 'enable_cfg': int(enable_cfg), | |
| 'loss': 0, | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| }, | |
| ) | |
| vit_image_tensor = self.vit_transform(image) | |
| height, width = vit_image_tensor.shape[1:] | |
| data['num_tokens'] += width * height // self.vit_transform.stride ** 2 | |
| data['image_tensor_list'].append(vit_image_tensor) | |
| return data | |
| def _add_video(self, data, frames, frame_indexes, need_loss, need_vae, enable_cfg=True): | |
| assert int(need_loss) + int(need_vae) == 1 | |
| if need_loss: | |
| for idx, (image, frame_idx) in enumerate(zip(frames, frame_indexes)): | |
| current_sequence_plan = { | |
| 'type': 'vae_image', | |
| 'enable_cfg': 0, | |
| 'loss': 1, | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| 'split_start': idx == 0, | |
| 'split_end': idx == len(frames) - 1, | |
| } | |
| if idx < len(frame_indexes) - 1: | |
| current_sequence_plan['frame_delta'] = frame_indexes[idx + 1] - frame_idx | |
| data['sequence_plan'].append(current_sequence_plan) | |
| image_tensor = self.transform(image) | |
| height, width = image_tensor.shape[1:] | |
| data['image_tensor_list'].append(image_tensor) | |
| data['num_tokens'] += width * height // self.transform.stride ** 2 | |
| elif need_vae: | |
| for idx, (image, frame_idx) in enumerate(zip(frames, frame_indexes)): | |
| current_sequence_plan = { | |
| 'type': 'vae_image', | |
| 'enable_cfg': int(enable_cfg), | |
| 'loss': 0, | |
| 'special_token_loss': 0, | |
| 'special_token_label': None, | |
| 'split_start': idx == 0, | |
| 'split_end': idx == len(frames) - 1, | |
| } | |
| if idx < len(frame_indexes) - 1: | |
| current_sequence_plan['frame_delta'] = frame_indexes[idx + 1] - frame_idx | |
| data['sequence_plan'].append(current_sequence_plan) | |
| image_tensor = self.transform(image) | |
| height, width = image_tensor.shape[1:] | |
| data['image_tensor_list'].append(image_tensor) | |
| data['num_tokens'] += width * height // self.transform.stride ** 2 | |
| return data | |
| class ParquetStandardIterableDataset(DistributedIterableDataset): | |
| def __init__( | |
| self, dataset_name, transform, tokenizer, vit_transform, | |
| data_dir_list, num_used_data, parquet_info, | |
| local_rank=0, world_size=1, num_workers=8, data_status=None, | |
| ): | |
| """ | |
| data_dir_list: list of data directories contains parquet files | |
| num_used_data: list of number of sampled data paths for each data directory | |
| vit_transform: input transform for vit model. | |
| """ | |
| super().__init__(dataset_name, local_rank, world_size, num_workers) | |
| self.transform = transform | |
| self.vit_transform = vit_transform | |
| self.tokenizer = tokenizer | |
| self.data_status = data_status | |
| self.data_paths = self.get_data_paths(data_dir_list, num_used_data, parquet_info) | |
| self.set_epoch() | |
| def get_data_paths(self, data_dir_list, num_used_data, parquet_info): | |
| row_groups = [] | |
| for data_dir, num_data_path in zip(data_dir_list, num_used_data): | |
| data_paths = get_parquet_data_paths([data_dir], [num_data_path]) | |
| for data_path in data_paths: | |
| if data_path in parquet_info.keys(): | |
| num_row_groups = parquet_info[data_path]['num_row_groups'] | |
| for rg_idx in range(num_row_groups): | |
| row_groups.append((data_path, rg_idx)) | |
| return row_groups | |
| def parse_row(self, row): | |
| raise NotImplementedError | |
| def __iter__(self): | |
| file_paths_per_worker, worker_id = self.get_data_paths_per_worker() | |
| if self.data_status is not None: | |
| global_row_group_start_id = self.data_status[worker_id][0] | |
| row_start_id = self.data_status[worker_id][1] + 1 | |
| else: | |
| global_row_group_start_id = 0 | |
| row_start_id = 0 | |
| print( | |
| f"rank-{self.local_rank} worker-{worker_id} dataset-{self.dataset_name}: " | |
| f"resuming data at global_rg#{global_row_group_start_id}, row#{row_start_id}" | |
| ) | |
| while True: | |
| file_paths_per_worker_ = file_paths_per_worker[global_row_group_start_id:] | |
| for global_row_group_idx, (parquet_file_path, row_group_id) in enumerate( | |
| file_paths_per_worker_, start=global_row_group_start_id | |
| ): | |
| fs = init_arrow_pf_fs(parquet_file_path) | |
| with fs.open_input_file(parquet_file_path) as f: | |
| try: | |
| fr = pq.ParquetFile(f) | |
| df = fr.read_row_group(row_group_id).to_pandas() | |
| df = df.iloc[row_start_id:] | |
| except Exception as e: | |
| print(f'Error {e} in rg#{row_group_id}, {parquet_file_path}') | |
| continue | |
| for row_idx, row in df.iterrows(): | |
| try: | |
| data = self.parse_row(row) | |
| if len(data) == 0: | |
| continue | |
| data['data_indexes'] = { | |
| "data_indexes": [global_row_group_idx, row_idx], | |
| "worker_id": worker_id, | |
| "dataset_name": self.dataset_name, | |
| } | |
| except Exception as e: | |
| print(f'Error {e} in rg#{row_group_id}, {parquet_file_path}') | |
| continue | |
| yield data | |
| row_start_id = 0 | |
| global_row_group_start_id = 0 | |
| print(f"{self.dataset_name} repeat in rank-{self.local_rank} worker-{worker_id}") | |