from suno_utils.utils.s3 import read_from_s3, open_from_s3, _read_npz import json from tqdm import tqdm import numpy as np import os import gc from dataclasses import dataclass from collections import defaultdict from joblib import Parallel, delayed import funcy import copy import itertools from suno_utils.utils.s3 import ( read_from_s3, check_s3_file_exists, list_s3_dir, download_s3_file_if_needed, ) from suno_utils.utils.text import write_jsonl, read_jsonl import polars as pl @dataclass class Bundle: name: str version: str = "v4" codec: str = "dac_vae_tuned_25hz" def __post_init__(self): self.check_s3_paths() @property def s3_dir(self): return f"s3://suno-data/datasets/bundles/{self.version}/{self.name}" @property def s3_metas_path(self): return f"{self.s3_dir}/metas.jsonl" @property def s3_parts_path(self): return f"{self.s3_dir}/{self.codec}" def s3_part_path(self, part_idx): return f"{self.s3_parts_path}/part_{part_idx}.npz" def s3_metas_part_path(self, part_idx): return f"{self.s3_parts_path}/metas/part_{part_idx}.jsonl" def check_s3_paths(self): assert check_s3_file_exists(self.s3_metas_path) assert check_s3_file_exists(self.s3_parts_path + "/part_0.npz") assert check_s3_file_exists(self.s3_part_path(0)) assert check_s3_file_exists(self.s3_metas_part_path(0)) def num_parts(self): print("Checking number of parts, this may take a while...") return len(list(list_s3_dir(f"{self.s3_parts_path}/metas"))) def get_part_metas(self, part_idx): with open_from_s3(self.s3_metas_part_path(part_idx), as_binary=False) as f: return [json.loads(line) for line in f] def get_metas(self, n_lines=None): f = download_s3_file_if_needed(self.s3_metas_path) df = pl.read_ndjson(f) return df.to_dicts() def get_part(self, part_idx): return read_from_s3(self.s3_part_path(part_idx), read_f=_read_npz) @dataclass class DatasetConfig: bundle: Bundle start_idx: int end_idx: int n_vae: int = 1 sort_key: str = None # sort by this key to make sure they are grouped (eg data sharding) def _parse_arrays(self, meta_info, vae_arr): """ This method should be implemented by the subclass Returns a list of tuples, where each tuple contains an array and a metadata dictionary corresponding to a single training sample Metas must have a "duration_s" key """ raise NotImplementedError("This method should be implemented by the subclass") @dataclass class MemmapMaker: out_data_dir: str duration_s: float = 30 token_hz: int = 25 vae_dim: int = 128 @property def vae_n_tokens_memmap(self): return int(self.duration_s * self.token_hz) def _process_archives(self, dataset: DatasetConfig, s3_vae_archive_filepaths, relevant_metas): # print(len(relevant_metas)) vae_archive = {} s3_vae_archive_filepaths = set(s3_vae_archive_filepaths) for s3_vae_archive_filepath in s3_vae_archive_filepaths: if "s3://" in s3_vae_archive_filepath: if not check_s3_file_exists(s3_vae_archive_filepath): print(f"missing {s3_vae_archive_filepath}") continue try: archive = { k: v for k, v in read_from_s3(s3_vae_archive_filepath, read_f=np.load).items() } except Exception as e: # corrupt archive print(e) print(f"corrupt {s3_vae_archive_filepath}") continue else: # check if file exists if not os.path.exists(s3_vae_archive_filepath): print(f"missing {s3_vae_archive_filepath}") continue try: archive = {k: v for k, v in np.load(s3_vae_archive_filepath).items()} except: # corrupt archive print(f"corrupt {s3_vae_archive_filepath}") continue for k, v in archive.items(): vae_archive[k] = v vae_uids = set(vae_archive.keys()) arr_list = [] for uid in vae_uids: if uid not in relevant_metas: # print(f"skipping {uid} not in relevant_metas") continue vae_arr = vae_archive[uid] arr_list.extend(dataset._parse_arrays(relevant_metas[uid], vae_arr)) del vae_archive gc.collect() return arr_list def _collect_uids(self, s3_vae_metas_filepaths): vae_uids = [] for fp in s3_vae_metas_filepaths: try: if "s3://" in fp: metas = read_from_s3(fp, read_f=read_jsonl) else: metas = read_jsonl(fp, progress=False) except: print(f"failed on metas for fp: {fp}") continue vae_uids.extend([m["id"] for m in metas]) return set(vae_uids) def _prep_data( self, dataset: DatasetConfig, meta_info_map: dict, njobs=5, chunksize=10, is_val=False, n_offs_v=0, ): dset_name = dataset.bundle.name start_idx = dataset.start_idx end_idx = dataset.end_idx n_vae = dataset.n_vae dset_type = "val" if is_val else "tr" out_mm_vae_filepath = os.path.join(self.out_data_dir, f"data_vae_{dset_type}.bin") out_metas_filepath = os.path.join(self.out_data_dir, f"metas_{dset_type}.jsonl") tot_duration_dict = defaultdict(int) n_chunks = int(np.ceil((end_idx - start_idx) / chunksize)) for idx_chunk in tqdm(funcy.chunks(chunksize, list(range(start_idx, end_idx))), total=n_chunks): n_jobs = np.min([njobs, chunksize, len(idx_chunk)]) # collect relevant parts of meta file to avoid copying all to subprocesses tmp_uid_chunks = Parallel(n_jobs=n_jobs, prefer="processes")( delayed(self._collect_uids)( [ dataset.bundle.s3_metas_part_path(idx_idx) for idx_idx in range(idx * n_vae, (idx + 1) * n_vae) ] ) for idx in idx_chunk ) ## PART A: takes ~40% of loop time uids_per_part = {idx: tmp_uid_chunks[n] for n, idx in enumerate(idx_chunk)} # print(len(uids_per_part)) # collect data encoded_arrays_list = Parallel(n_jobs=n_jobs, prefer="processes")( delayed(self._process_archives)( dataset, [ dataset.bundle.s3_part_path(idx_idx) for idx_idx in range(idx * n_vae, (idx + 1) * n_vae) ], {uid: meta_info_map[uid] for uid in uids_per_part[idx] if uid in meta_info_map}, ) for idx in idx_chunk ) if dataset.sort_key is not None: # sort to make sure they are grouped (eg data sharding) # print(f"Sorting by {dataset.sort_key}") encoded_arrays_list = [ sorted(encoded_arrays, key=lambda x: x[1][dataset.sort_key]) for encoded_arrays in encoded_arrays_list ] ## end Part A ## PART B: takes ~40% of loop time add_metas = [] for encoded_arrays in encoded_arrays_list: to_write_len_v = np.sum([arr.size for arr, _ in encoded_arrays]) if to_write_len_v == 0: print(f"skipping {len(encoded_arrays)}") continue out_mm_vae = np.memmap( out_mm_vae_filepath, dtype=np.float16, mode="r+", shape=(n_offs_v + to_write_len_v,), ) for arr_v, arr_meta in encoded_arrays: if arr_v.size != self.vae_n_tokens_memmap * self.vae_dim: print(f"skipping {arr_v.size} != {self.vae_n_tokens_memmap * self.vae_dim}") continue out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape( -1, ) n_offs_v += arr_v.size add_meta = copy.deepcopy(arr_meta) add_meta["dataset"] = dset_name tot_duration_dict[dset_name] += arr_meta["duration_s"] add_metas.append(add_meta) # write it once out_mm_vae.flush() del out_mm_vae ## end Part B write_jsonl( add_metas, os.path.join(out_metas_filepath), do_append=bool(n_offs_v != 0), ) del encoded_arrays_list # TODO: this gc collect takes super long but maybe ok outside of loop. somehow needed sometimes gc.collect() for k, v in tot_duration_dict.items(): print(f"{round(v / 60 / 60):,} hours of {k}") return n_offs_v def prep_data( self, datasets: list[tuple[DatasetConfig, dict]], is_val=False, njobs=5, chunksize=10, ): os.makedirs(self.out_data_dir, exist_ok=True) n_offs_v = 0 dset_type = "val" if is_val else "tr" out_mm_vae_filepath = os.path.join(self.out_data_dir, f"data_vae_{dset_type}.bin") out_metas_filepath = os.path.join(self.out_data_dir, f"metas_{dset_type}.jsonl") out_mm_vae = np.memmap(out_mm_vae_filepath, dtype=np.float16, mode="w+", shape=(1,)) with open(out_metas_filepath, "w") as f: f.write("") print("start prepare data") for dataset, meta_info_map in datasets: n_offs_v = self._prep_data( dataset, meta_info_map, njobs=njobs, chunksize=chunksize, is_val=is_val, n_offs_v=n_offs_v, )