content: # 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}")
  