import os import numpy as np import pandas as pd from collections import defaultdict from tqdm import tqdm from suno_utils.utils.text import write_jsonl import traceback # -------------------------------------------------------------------------- # Constants VAE_MEMMAP_SIZE = 750 SEMANTIC_MEMMAP_SIZE = 750 VAE_DIM = 128 CHUNK_SIZE = 100 ENCODE_HZ = 25 def validate_positive_and_negative_quality( positive_quality_score, negative_quality_score, previous_positive_quality_score, previous_negative_quality_score, ): # positive_quality_score["pair_quality"] = float(positive_quality_score["pref"]) positive_quality_score["total_shimmer_score"] = float( positive_quality_score["shimmer_score"] ) negative_quality_score["total_shimmer_score"] = float( negative_quality_score["shimmer_score"] ) positive_quality_score["shimmer_score_diff"] = float( positive_quality_score["shimmer_score"] - negative_quality_score["shimmer_score"] ) positive_quality_score["stereo_width_diff"] = float( float(positive_quality_score["stereo_width"]) - float(negative_quality_score["stereo_width"]) ) positive_quality_score["spectral_centroid_diff"] = float( float(positive_quality_score["spectral_centroid"]) - float(negative_quality_score["spectral_centroid"]) ) / (float(positive_quality_score["spectral_centroid"]) + 0.01) spectral_decay_diff = float( positive_quality_score["spectrum_decay"] - negative_quality_score["spectrum_decay"] ) / ( float(positive_quality_score["spectrum_decay"]) + float(negative_quality_score["spectrum_decay"]) + 0.0001 ) if ( previous_positive_quality_score is not None and previous_negative_quality_score is not None ): positive_spectral_centroid_drift_delta = float( positive_quality_score["spectral_centroid"] ) - float(previous_positive_quality_score["spectral_centroid"]) negative_spectral_centroid_drift_delta = float( negative_quality_score["spectral_centroid"] ) - float(previous_negative_quality_score["spectral_centroid"]) relative_spectral_centroid_drift_delta = ( positive_spectral_centroid_drift_delta - negative_spectral_centroid_drift_delta ) / ( ( float(previous_positive_quality_score["spectral_centroid"]) + float(previous_negative_quality_score["spectral_centroid"]) + 1 ) / 2 ) else: relative_spectral_centroid_drift_delta = 0 negative_ear_average_score = np.mean( negative_quality_score["ear_v2_quality_scores"] ) positive_ear_average_score = np.mean( positive_quality_score["ear_v2_quality_scores"] ) ear_score_relative_diff = ( positive_ear_average_score - negative_ear_average_score ) / (positive_ear_average_score + 0.001) previous_negative_ear_average_score = ( np.mean(previous_negative_quality_score["ear_v2_quality_scores"]) if previous_negative_quality_score is not None else 0 ) previous_positive_ear_average_score = ( np.mean(previous_positive_quality_score["ear_v2_quality_scores"]) if previous_positive_quality_score is not None else 0 ) prev_ear_score_relative_diff = ( previous_positive_ear_average_score - previous_negative_ear_average_score ) / (previous_positive_ear_average_score + 0.001) # previous_ear_score_relative_diff = ( # positive_ear_average_score - previous_positive_ear_average_score # ) / (positive_ear_average_score + 0.001) - ( # negative_ear_average_score - previous_negative_ear_average_score # ) / (negative_ear_average_score + 0.001) # these cuts should be consistent with the original notebook # these are roughly bottom 10% cuts each # -0.19 is 5% -- r2 -- each round need calibration ear_is_different = abs(ear_score_relative_diff) > 0.05 # -0.15 is 5%x prev_ear_is_much_different = abs(prev_ear_score_relative_diff) > 0.1 # is_positive_no_shimer = (positive_quality_score["total_shimmer_score"] < 1) | ( # positive_quality_score["shimmer_score_diff"] < 0.4 # ) # 5% is 1.8, 2% is 2.9 is_positive_no_shimer = (positive_quality_score["total_shimmer_score"] < 1.5) | ( positive_quality_score["shimmer_score_diff"] < 0.75 ) loundess_ratio = ( positive_quality_score["abs_loudness_factor"] - negative_quality_score["abs_loudness_factor"] ) / abs(negative_quality_score["abs_loudness_factor"] + 0.001) # not too loud by itself # 0.15 is 5%, 0.18 is 2% is_positive_not_louder = loundess_ratio < 0.15 # # not too much quieter than negative # is_positive_not_much_quieter = loundess_ratio > -0.20 # more stereo width 10% cut # CHECK DATA: 5% -0.20, 2% -0.26 is_positive_more_stereo = positive_quality_score["stereo_width_diff"] > -0.25 # spectral centroid 2% cut # is_positive_not_high_freq = positive_quality_score["spectral_centroid_diff"] < 0.30 # spectral centroid drift delta 5% cut # is_positive_not_shifting_high_freq = relative_spectral_centroid_drift_delta < 0.22 # not significantly different spectrum decay is_not_much_different_spectrum_decay = spectral_decay_diff < 0.30 return ( is_positive_no_shimer and is_positive_not_louder and is_positive_more_stereo and is_not_much_different_spectrum_decay # and is_positive_not_much_quieter # and is_positive_not_louder_abs # and is_positive_not_high_freq # and is_positive_not_shifting_high_freq ) # return True def make_dataset( input_df, output_data_dir: str, is_val=False, npz_dir: str = "/app/suno/data/dpo/diff_v2", do_extend_chunks=False, clip_id_to_quality_scores=None, ): # set dataset type dset_type = "val" if is_val else "tr" # -------------------------------------------------------------------------- # create output dir os.makedirs(output_data_dir, exist_ok=True) print(f"df shape: {input_df.shape}") # create memmap files out_mm_vae_filepath = os.path.join(output_data_dir, f"data_vae_{dset_type}.bin") out_metas_filepath = os.path.join(output_data_dir, f"metas_{dset_type}.jsonl") out_mm_semantic_filepath = os.path.join( output_data_dir, f"data_semantic_{dset_type}.bin" ) if os.path.exists(out_mm_vae_filepath): os.remove(out_mm_vae_filepath) if os.path.exists(out_mm_semantic_filepath): os.remove(out_mm_semantic_filepath) if os.path.exists(out_metas_filepath): os.remove(out_metas_filepath) # initial write out_mm_semantic = np.memmap( out_mm_semantic_filepath, dtype=np.uint16, mode="w+", shape=(1), ) out_mm_vae = np.memmap( out_mm_vae_filepath, dtype=np.float16, mode="w+", shape=(1), ) n_offs_s = 0 n_offs_v = 0 # create indices for the df # we will count by 2, so take all pairs of rows starting at index 0 df_indices = list(range(0, len(input_df), 2)) # now chunk the df indices df_indices_chunks = [ df_indices[i : i + CHUNK_SIZE] for i in range(0, len(df_indices), CHUNK_SIZE) ] print(f"total chunks: {len(df_indices_chunks)}") n_total_failed_quality_check = 0 n_total_passed_quality_check = 0 n_total_recorded = 0 n_total_chunks_with_prev_chunk_as_vae_ctx = 0 n_total_seeds = 0 # two consecutive rows form the positive and negative pair for _, df_indices in enumerate(tqdm(df_indices_chunks)): arr_s_list = [] arr_v_list = [] new_metas = [] for i in df_indices: negative_row = input_df.iloc[i] positive_row = input_df.iloc[i + 1] assert positive_row["preference"] == 1 assert negative_row["preference"] == 0 positive_clip_id = positive_row["s3_id"] negative_clip_id = negative_row["s3_id"] # load npz files negative_infos = defaultdict(list) positive_infos = defaultdict(list) negative_parent_idx = negative_row["metadata"].get("upsample_clip_id", None) positive_parent_idx = positive_row["metadata"].get("upsample_clip_id", None) assert negative_parent_idx is not None assert positive_parent_idx is not None assert negative_parent_idx == positive_parent_idx parent_idx = positive_parent_idx negative_seed = None positive_seed = None for id_x in [negative_row["s3_id"], positive_row["s3_id"]]: # get semantic codes coarse_npz_path = os.path.join(npz_dir, f"{parent_idx}.npz") coarse_npz = np.load(coarse_npz_path) try: if "v5.0_raw" in coarse_npz: codes = coarse_npz["v5.0_raw"] elif "v4.0_raw" in coarse_npz: codes = coarse_npz["v4.0_raw"] elif "v3.5_raw" in coarse_npz: codes = coarse_npz["v3.5_raw"] elif "v3.0_raw" in coarse_npz: codes = coarse_npz["v3.0_raw"] elif "v2.0_raw" in coarse_npz: codes = coarse_npz["v2.0_raw"] else: raise ValueError() except Exception as e: print(coarse_npz_path) raise e semantic_codes = codes[:, 0].astype(np.uint16) # get vae latents vae_npz_path = os.path.join(npz_dir, f"{id_x}_vae.npz") vae_npz = np.load(vae_npz_path) vae_latents = vae_npz["vae_latents"].astype(np.float16) try: if id_x == negative_row["s3_id"]: negative_seed = int(vae_npz["seed"]) else: positive_seed = int(vae_npz["seed"]) n_total_seeds += 1 except Exception as e: pass # if we expand the chunks if do_extend_chunks: n_total_chunks = ( min(semantic_codes.shape[0], vae_latents.shape[0]) // VAE_MEMMAP_SIZE ) # if we don't expand the chunks # even for extend chunks, let's be conservative for now # n_total_chunks = min(n_total_chunks, 2) else: n_total_chunks = 1 can_use_prev_chunk_as_vae_ctx = False for local_chunk_idx in range(0, n_total_chunks): # process a quality check first if ( clip_id_to_quality_scores # and local_chunk_idx < 2 # others don't exist yet ): positive_quality_score = clip_id_to_quality_scores[ positive_clip_id + (f"_{local_chunk_idx}" if local_chunk_idx > 0 else "") ] negative_quality_score = clip_id_to_quality_scores[ negative_clip_id + (f"_{local_chunk_idx}" if local_chunk_idx > 0 else "") ] previous_positive_quality_score = None previous_negative_quality_score = None if local_chunk_idx > 0: previous_positive_quality_score = clip_id_to_quality_scores[ positive_clip_id + ( f"_{local_chunk_idx - 1}" if local_chunk_idx > 1 else "" ) ] previous_negative_quality_score = clip_id_to_quality_scores[ negative_clip_id + ( f"_{local_chunk_idx - 1}" if local_chunk_idx > 1 else "" ) ] assert positive_quality_score is not None, positive_clip_id assert negative_quality_score is not None, negative_clip_id # check if it passes the quality check try: if not validate_positive_and_negative_quality( positive_quality_score, negative_quality_score, previous_positive_quality_score, previous_negative_quality_score, ): n_total_failed_quality_check += 1 can_use_prev_chunk_as_vae_ctx = False # continue will keep the later chunks # break will only use the ones above continue n_total_passed_quality_check += 1 except Exception as e: print(f"Error: {e}") traceback.print_exc() break # cut off to the first 30 seconds of semantic codes arr_s = semantic_codes[ local_chunk_idx * VAE_MEMMAP_SIZE : (local_chunk_idx + 1) * VAE_MEMMAP_SIZE ] arr_v = vae_latents[ local_chunk_idx * VAE_MEMMAP_SIZE : (local_chunk_idx + 1) * VAE_MEMMAP_SIZE, :, ] # check for nan in arr_v or arr_s if not np.all(np.isfinite(arr_v)) or not np.all(np.isfinite(arr_s)): print(f"{vae_npz_path} has nan or inf") # create a new meta # always take the negative as the prompt -- cause positive can be modified # TODO: this is a bad interleave pattern # Right now -- negative will load the prev vae which is positive # This might be good since positive will be more consistent? # But this might also be too easy to learn...? meta = { "text": negative_row["prompt_text"] if not pd.isna(negative_row["prompt_text"]) else "", # note that tags are a list... "tags": [ negative_row["metadata"]["tags"] if not pd.isna(negative_row["metadata"]["tags"]) else "" ], "n_vae_tokens": VAE_MEMMAP_SIZE, "start_s": round( 0.0 + local_chunk_idx * VAE_MEMMAP_SIZE / ENCODE_HZ, 1 ), "end_s": round( 30.0 + local_chunk_idx * VAE_MEMMAP_SIZE / ENCODE_HZ, 1 ), "id_x": id_x, "seed": negative_seed if id_x == negative_row["s3_id"] else positive_seed, } # only do this for future VAE chunks # TODO: This might be why we have a noise decay issue... if ( do_extend_chunks and local_chunk_idx > 0 and can_use_prev_chunk_as_vae_ctx ): meta["prev_context_id"] = id_x meta["offset_rows"] = 2 n_total_chunks_with_prev_chunk_as_vae_ctx += 1 assert arr_s.size >= SEMANTIC_MEMMAP_SIZE assert arr_v.size >= VAE_MEMMAP_SIZE * VAE_DIM if id_x == negative_row["s3_id"]: negative_infos["meta"].append(meta) negative_infos["semantic_codes"].append(arr_s) negative_infos["vae_latents"].append(arr_v) else: positive_infos["meta"].append(meta) positive_infos["semantic_codes"].append(arr_s) positive_infos["vae_latents"].append(arr_v) can_use_prev_chunk_as_vae_ctx = True # make sure we doubled the chunks assert len(negative_infos["meta"]) == len(positive_infos["meta"]) assert len(negative_infos["semantic_codes"]) == len( positive_infos["semantic_codes"] ) assert len(negative_infos["vae_latents"]) == len( positive_infos["vae_latents"] ) try: for n_chunk, (neg_s, pos_s) in enumerate( zip( negative_infos["semantic_codes"], positive_infos["semantic_codes"], ) ): assert np.array_equal( neg_s, pos_s ), f"semantic codes should be the same {neg_s.shape} {pos_s.shape} n_chunk: {n_chunk}" except Exception as e: print( f"negative_row: {negative_row['s3_id']}, positive_row: {positive_row['s3_id']}, request_id: {positive_row['request_id']}" ) print(e) continue # Then unpack the rest for neg_meta, pos_meta in zip( negative_infos["meta"], positive_infos["meta"] ): new_metas.append(neg_meta) new_metas.append(pos_meta) for neg_s, pos_s in zip( negative_infos["semantic_codes"], positive_infos["semantic_codes"] ): assert np.array_equal( neg_s, pos_s ), f"semantic codes should be the same {neg_s.shape} {pos_s.shape}" arr_s_list.append(neg_s) arr_s_list.append(neg_s) for neg_v, pos_v in zip( negative_infos["vae_latents"], positive_infos["vae_latents"] ): arr_v_list.append(neg_v) arr_v_list.append(pos_v) # the infos should be consistent assert len(arr_s_list) == len(arr_v_list) == len(new_metas) to_write_len_s = SEMANTIC_MEMMAP_SIZE * len(arr_s_list) to_write_len_v = VAE_MEMMAP_SIZE * VAE_DIM * len(arr_v_list) out_mm_semantic = np.memmap( out_mm_semantic_filepath, dtype=np.uint16, mode="r+", shape=(n_offs_s + to_write_len_s,), ) out_mm_vae = np.memmap( out_mm_vae_filepath, dtype=np.float16, mode="r+", shape=(n_offs_v + to_write_len_v,), ) # write to memmap (has to happen sequentially) for _, arr_s, arr_v in zip(new_metas, arr_s_list, arr_v_list): out_mm_semantic[n_offs_s : n_offs_s + arr_s.size] = arr_s.reshape( -1, ) out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape( -1, ) n_offs_s += arr_s.size n_offs_v += arr_v.size n_total_recorded += 1 # write it once out_mm_semantic.flush() out_mm_vae.flush() del out_mm_semantic, out_mm_vae write_jsonl(new_metas, out_metas_filepath, do_append=True) print( f"Done! {dset_type}: wrote {n_offs_s} semantic tokens and {n_offs_v} vae latents. \n" f"Total slices of data: {n_total_recorded}. Per node: {round(n_total_recorded / 32, 1)}. \n" f"Passed quality check: {n_total_passed_quality_check}, Failed quality check: {n_total_failed_quality_check}. \n" f"Total chunks with prev chunk as vae ctx: {n_total_chunks_with_prev_chunk_as_vae_ctx}. \n" f"Total seeds: {n_total_seeds}." )