(self, datasets_metainfo, data_status)
| 122 | self.step_counter = 0 |
| 123 | |
| 124 | def build_datasets(self, datasets_metainfo, data_status): |
| 125 | datasets = [] |
| 126 | is_mandatory = [] |
| 127 | grouped_weights = [] |
| 128 | grouped_names = [] |
| 129 | for grouped_dataset_name, dataset_args in datasets_metainfo.items(): |
| 130 | is_mandatory.append(dataset_args.pop('is_mandatory', False)) |
| 131 | grouped_weights.append(dataset_args.pop('weight', 0.0)) |
| 132 | |
| 133 | if 'frame_sampler_args' in dataset_args.keys(): |
| 134 | frame_sampler = FrameSampler(**dataset_args.pop('frame_sampler_args')) |
| 135 | dataset_args['frame_sampler'] = frame_sampler |
| 136 | if 'image_transform_args' in dataset_args.keys(): |
| 137 | transform = ImageTransform(**dataset_args.pop('image_transform_args')) |
| 138 | dataset_args['transform'] = transform |
| 139 | if 'vit_image_transform_args' in dataset_args.keys(): |
| 140 | vit_transform = QwenVL2ImageTransform(**dataset_args.pop('vit_image_transform_args')) |
| 141 | dataset_args['vit_transform'] = vit_transform |
| 142 | if 'dino_image_transform_args' in dataset_args.keys(): |
| 143 | dino_transform = DinoImageTransform(**dataset_args.pop('dino_image_transform_args')) |
| 144 | dataset_args['dino_transform'] = dino_transform |
| 145 | |
| 146 | assert 'dataset_names' in dataset_args.keys() |
| 147 | dataset_names = dataset_args.pop('dataset_names') |
| 148 | dataset_args['data_dir_list'] = [] |
| 149 | for item in dataset_names: |
| 150 | if self.local_rank == 0: |
| 151 | print(f'Preparing Dataset {grouped_dataset_name}/{item}') |
| 152 | meta_info = DATASET_INFO[grouped_dataset_name][item] |
| 153 | dataset_args['data_dir_list'].append(meta_info['data_dir']) |
| 154 | |
| 155 | if "parquet_info_path" in meta_info.keys(): |
| 156 | if 'parquet_info' not in dataset_args.keys(): |
| 157 | dataset_args['parquet_info'] = {} |
| 158 | with open(meta_info['parquet_info_path'], 'r') as f: |
| 159 | parquet_info = json.load(f) |
| 160 | dataset_args['parquet_info'].update(parquet_info) |
| 161 | |
| 162 | if 'json_dir' in meta_info.keys(): |
| 163 | # parquet/tar with json |
| 164 | if 'json_dir_list' not in dataset_args.keys(): |
| 165 | dataset_args['json_dir_list'] = [meta_info['json_dir']] |
| 166 | else: |
| 167 | dataset_args['json_dir_list'].append(meta_info['json_dir']) |
| 168 | |
| 169 | if 'jsonl_path' in meta_info.keys(): |
| 170 | # jsonl with jpeg |
| 171 | if 'jsonl_path_list' not in dataset_args.keys(): |
| 172 | dataset_args['jsonl_path_list'] = [meta_info['jsonl_path']] |
| 173 | else: |
| 174 | dataset_args['jsonl_path_list'].append(meta_info['jsonl_path']) |
| 175 | |
| 176 | resume_data_status = dataset_args.pop('resume_data_status', True) |
| 177 | if data_status is not None and grouped_dataset_name in data_status.keys() and resume_data_status: |
| 178 | data_status_per_group = data_status[grouped_dataset_name] |
| 179 | else: |
| 180 | data_status_per_group = None |
| 181 | dataset = DATASET_REGISTRY[grouped_dataset_name]( |
no test coverage detected