import numpy as np import os import tqdm from collections import defaultdict, Counter import pandas as pd from suno_utils.utils.text import write_jsonl, write_json SEMANTIC_N_CODEBOOKS = 1 COARSE_RATE_HZ = 25 SEMANTIC_CODEBOOK_SIZE = 4000 SEMANTIC_PAD_TOKEN = SEMANTIC_CODEBOOK_SIZE SEMANTIC_RATE_HZ = 25 COARSE_CODEBOOK_SIZE = 2048 COARSE_N_CODEBOOKS = 12 COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE N_TOKENS_AUDIO = 25 * 8 * 60 # max 8 mins of audio def make_dataset( input_df, output_data_dir: str, is_val=False, npz_dir="/app2/suno/data/dpo/auk_t1_npz", t_data_memmap=N_TOKENS_AUDIO, original_npz_dir="/app2/suno/data/dpo/auk_t1_npz", # this is the npz directory of the non-cycled ): print(f"t_data_memmap is set to: {t_data_memmap}") dset_type = "val" if is_val else "tr" out_mmap_path = os.path.join(output_data_dir, f"data_{dset_type}.bin") out_metas_path = os.path.join(output_data_dir, f"meta_{dset_type}.jsonl") out_info_filepath = os.path.join(output_data_dir, f"info_{dset_type}.json") def reshift(arr): sem_start_idx = 0 sem_end_idx = len(arr) - 1 semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS] # get array segments arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy() assert arr_s.max() <= SEMANTIC_PAD_TOKEN # concat and stack arr_s = np.pad( arr_s, ((0, t_data_memmap - len(arr_s)), (0, 0)), constant_values=SEMANTIC_PAD_TOKEN, mode="constant", ) arr = arr_s arr = arr.astype(np.uint16) assert arr.shape == (t_data_memmap, SEMANTIC_N_CODEBOOKS) return arr # gather the data _ = np.memmap(out_mmap_path, dtype=np.uint16, mode="w+", shape=(1,)) n_offs = 0 curr_idx = 0 tot_duration_dict = defaultdict(int) datasets_info = defaultdict(dict) total_different_prompts = 0 total_different_tags = 0 total_different_neg_tags = 0 negative_prompt = "" negative_tag = "" negative_neg_tag = "" total_task_counter = Counter() total_error_task_counter = Counter() gender_counter = Counter() control_slider_counter = Counter() neg_tags_counter = Counter() prev_skip = None for i, row in tqdm.tqdm(input_df.iterrows(), total=len(input_df)): # we need to alternate between preference: neg, pos if i - 1 == prev_skip: print( f"WTF --> {i}, skip, preference: {row['preference']}, {row['s3_id']}, task: {row.get('task', '')}." ) continue # print(i, row) assert row["preference"] == (i % 2 == 1) is_positive = row["preference"] # need to use prompt_text -- prompt could be edited; # but this only works with recent data? since Aug 2024? current_prompt = row["prompt_text"] if not pd.isna(row["prompt_text"]) else "" current_tag = row["tags"] if not pd.isna(row["tags"]) else "" current_neg_tag = ( row["negative_tags"] if not pd.isna(row["negative_tags"]) else "" ) current_vocal_gender = ( row["vocal_gender"] if not pd.isna(row["vocal_gender"]) else "" ) # let's take this with a grain of salt # cause param_experiment could mask_control_slider current_control_sliders = ( row["control_sliders"] if not pd.isna(row["control_sliders"]) else {} ) control_slider_dict = {} if isinstance(current_control_sliders, dict): tag_weight = current_control_sliders.get("style_weight", None) audio_weight = current_control_sliders.get("audio_weight", None) mixed_weirdness = current_control_sliders.get("weirdness_constraint", None) if tag_weight is not None: control_slider_dict["tag_strength"] = tag_weight if audio_weight is not None: control_slider_dict["audio_strength"] = audio_weight if mixed_weirdness is not None: control_slider_dict["weird_strength"] = mixed_weirdness # audio_weight_strength = "" # if audio_weight is not None and audio_weight >= 0.75: # audio_weight_strength = "high" # elif audio_weight is not None and audio_weight <= 0.2: # audio_weight_strength = "low" # weirdness_strength = "" # if mixed_weirdness is not None and mixed_weirdness >= 0.75: # weirdness_strength = "high" # elif mixed_weirdness is not None and mixed_weirdness <= 0.2: # weirdness_strength = "low" # tag_strength = "" # if tag_weight is not None and tag_weight >= 0.75: # tag_strength = "high" # elif tag_weight is not None and tag_weight <= 0.2: # tag_strength = "low" # control_slider_str = "" # if audio_weight_strength != "": # control_slider_str += f"audio_strength:{audio_weight_strength};" # if weirdness_strength != "": # control_slider_str += f"weird_strength:{weirdness_strength};" # if tag_strength != "": # control_slider_str += f"tag_strength:{tag_strength};" # control_slider_str = control_slider_str.rstrip(";") # this is a necessary augmentation if current_vocal_gender == "m": current_vocal_gender = "male" elif current_vocal_gender == "f": current_vocal_gender = "female" gender_counter[ current_vocal_gender if current_vocal_gender else "unspecified" ] += 1 control_slider_counter[ "has_control_slider" if control_slider_dict else "no_control_slider" ] += 1 neg_tags_counter["has_neg_tags" if current_neg_tag else "no_neg_tags"] += 1 if is_positive: if current_prompt != negative_prompt: # print("different prompts", negative_prompt, current_prompt) # ppl probably won't change negative... current_prompt = negative_prompt total_different_prompts += 1 if current_tag != negative_tag: # print("different tags", negative_tag, current_tag) # ppl probably won't change negative... current_tag = negative_tag total_different_tags += 1 if current_neg_tag != negative_neg_tag: current_neg_tag = negative_neg_tag total_different_neg_tags += 1 else: negative_prompt = current_prompt negative_tag = current_tag negative_neg_tag = current_neg_tag is_cycled = "cycle" in npz_dir # make mmap -- two different paths if "npz_path" in row: local_path = row["npz_path"] else: local_path = f"{npz_dir}/{row['s3_id'] + ('_gen_cycle' if (is_cycled and 'cycle' not in row['s3_id']) else '')}.npz" original_path = ( f"{original_npz_dir}/{row['s3_id'].replace('_gen_cycle', '')}.npz" ) if not os.path.exists(local_path): # pass raise ValueError(f"File does not exist: {local_path}") try: temp_npz = np.load(local_path) if "v5.0_raw" in temp_npz: arr = temp_npz["v5.0_raw"] elif "v4.0_raw" in temp_npz: if "cycle" not in local_path: print(f"weird, {local_path}, with only v4.0") arr = temp_npz["v4.0_raw"] elif "v3.5_raw" in temp_npz: if "cycle" not in local_path: print(f"weird, {local_path}, with only v3.5") arr = temp_npz["v3.5_raw"] elif "v3.0_raw" in temp_npz: if "cycle" not in local_path: print(f"weird, {local_path}, with only v3.0") arr = temp_npz["v3.0_raw"] else: raise ValueError() except Exception as e: print(local_path) raise e if arr.shape[0] > t_data_memmap: print(f"Overflow at {i}, {row['s3_id']}: {arr.shape[0]} > {t_data_memmap}") arr = arr[:t_data_memmap, :] assert arr.shape[0] <= t_data_memmap # can't be longer than 4 mins :) if arr.shape[1] != 1: print(f"weird, {local_path}, with {arr.shape[1]} channels") arr = arr[:, :1] current_task = row.get("task", "") add_meta = { "dataset": f"perference_{int(row['preference'])}", "id": row["s3_id"], # this is the row s3_id "start_s": row["total_start_s"] if row["total_start_s"] >= 0 else 0, "vocal_start_s": None, # these are unfortunately missing for now "vocal_end_s": None, # these are unfortunately missing for now "tags": [current_tag], # tags is a list, do you know :) "neg_tags": current_neg_tag, # but we save negative tags as a string "control_tags": "" if pd.isna(row.get("control_tags", "")) else row.get("control_tags", ""), # save the control tags "gender": current_vocal_gender, # save the gender "control": control_slider_dict, # save the OG inputs, augment in training "text": current_prompt, # replace this with the fixed prompt "generated_start_index": 0, "user_id": row["user_id"], } # default gens, arr is what we look for. However... # all the conditioning shit # order is now: future + cover + artist known_dpo_tasks = [ "cover", "artist_consistency", "infill", "extend", "upload_extend", "cover_extend", "artist_extend", "artist_cover", "playlist_condition", "overpainting", "underpainting", "stem_condition", "", ] try: if current_task == "cover": cover_prompt_arr = temp_npz["cover_arr"] assert cover_prompt_arr.shape[0] <= t_data_memmap assert cover_prompt_arr.shape[1] == 1 arr = np.concatenate([cover_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = cover_prompt_arr.shape[0] add_meta["task"] = "cover" elif current_task == "artist_consistency": artist_prompt_arr = temp_npz["artist_arr"] assert artist_prompt_arr.shape[0] <= t_data_memmap assert artist_prompt_arr.shape[1] == 1 arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "artist_consistency" elif current_task == "overpainting": artist_prompt_arr = temp_npz["overpainting_arr"] assert artist_prompt_arr.shape[0] <= t_data_memmap assert artist_prompt_arr.shape[1] == 1 arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "overpainting" elif current_task == "underpainting": artist_prompt_arr = temp_npz["underpainting_arr"] assert artist_prompt_arr.shape[0] <= t_data_memmap assert artist_prompt_arr.shape[1] == 1 arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "underpainting" elif current_task == "stem_condition": stem_prompt_arr = temp_npz["stem_arr"] assert stem_prompt_arr.shape[0] <= t_data_memmap assert stem_prompt_arr.shape[1] == 1 arr = np.concatenate([stem_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = stem_prompt_arr.shape[0] add_meta["task"] = "stem_condition" elif current_task == "playlist_condition": playlist_prompt_arr = temp_npz["playlist_arr"] assert playlist_prompt_arr.shape[0] <= t_data_memmap assert playlist_prompt_arr.shape[1] == 1 arr = np.concatenate([playlist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = playlist_prompt_arr.shape[0] add_meta["task"] = "playlist_condition" add_meta["playlist_arr_len"] = [ int(playlist_arr_len) for playlist_arr_len in temp_npz["playlist_arr_len"] ] if ( sum(add_meta["playlist_arr_len"]) != add_meta["generated_start_index"] ): print( f"Playlist array lengths {sum(add_meta['playlist_arr_len'])} don't match generation_start_index {add_meta['generated_start_index']}" ) elif current_task == "artist_cover": artist_prompt_arr = temp_npz["artist_arr"] assert artist_prompt_arr.shape[0] <= t_data_memmap assert artist_prompt_arr.shape[1] == 1 cover_prompt_arr = temp_npz["cover_arr"] assert cover_prompt_arr.shape[0] <= t_data_memmap assert cover_prompt_arr.shape[1] == 1 # the order is now: artist, cover, generated arr = np.concatenate([artist_prompt_arr, cover_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = ( artist_prompt_arr.shape[0] + cover_prompt_arr.shape[0] ) add_meta["cover_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "artist_cover" elif ( current_task == "infill" or current_task == "infill_outro" or current_task == "infill_intro" ): # print("FOUND INFILL", current_task) # infilling is special data loading... arr = temp_npz["full_arr"] assert arr.shape[1] == 1 generated_start_index = temp_npz.get("generated_start_index", 0) if isinstance(generated_start_index, np.ndarray): # print(generated_start_index, generated_start_index.shape) generated_start_index = generated_start_index.item() assert isinstance(generated_start_index, int) future_start_index = temp_npz.get("future_start_index", arr.shape[0]) if isinstance(future_start_index, np.ndarray): future_start_index = future_start_index.item() assert isinstance(future_start_index, int) # these are the context window start / ends history_start_index = temp_npz.get("history_start_index", 0) if isinstance(history_start_index, np.ndarray): history_start_index = history_start_index.item() post_future_start_index = temp_npz.get( "post_future_start_index", arr.shape[0] ) if isinstance(post_future_start_index, np.ndarray): post_future_start_index = post_future_start_index.item() # crop the array to the history start and post future start arr = arr[history_start_index:post_future_start_index, :] # now we need to offset everything to the history start add_meta["generated_start_index"] = ( generated_start_index - history_start_index ) add_meta["future_start_index"] = ( future_start_index - history_start_index ) # crop to max duration to avoid overflow context window if arr.shape[0] > t_data_memmap: print(f"Overflow at {i}: {arr.shape[0]} > {t_data_memmap}") arr = arr[:t_data_memmap, :] add_meta["task"] = "infill" add_meta["infill_lyrics"] = ( row["metadata"].get("infill_lyrics", "") or "" ) # print(add_meta["text"]) # # disable extend concat for now... don't think we understand it elif ( current_task == "extend" or current_task == "upload_extend" or current_task == "artist_extend" or current_task == "cover_extend" ): add_meta["task"] = "extend" if is_cycled and "history_arr" not in temp_npz: # replace with original npz temp_npz = np.load(original_path) history_arr = temp_npz["history_arr"] assert history_arr.shape[1] == 1 arr = np.concatenate([history_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = history_arr.shape[0] # if history_text is not empty, we need to add it to the meta # print(temp_npz["history_text"]) history_text = temp_npz.get("history_text", "") add_meta["text"] = str(history_text) + "\n" + add_meta["text"] if current_task == "artist_extend": add_meta["task"] = "artist_extend" artist_prompt_arr = temp_npz["artist_arr"] assert artist_prompt_arr.shape[0] <= t_data_memmap assert artist_prompt_arr.shape[1] == 1 arr = np.concatenate([artist_prompt_arr, arr], axis=0) add_meta["history_start_index"] = artist_prompt_arr.shape[0] if current_task == "cover_extend": add_meta["task"] = "cover_extend" cover_prompt_arr = temp_npz["cover_arr"] assert cover_prompt_arr.shape[0] <= t_data_memmap assert cover_prompt_arr.shape[1] == 1 arr = np.concatenate([cover_prompt_arr, arr], axis=0) add_meta["history_start_index"] = cover_prompt_arr.shape[0] arr = arr[:t_data_memmap, :] except Exception as E: # there are a bunch of sth wrong with some data... print(f"{i}, {E}, {current_task}, {local_path}.") # don't keep the junks # they tend to be pairs anyways... # if not we wll find them and kick them out... if i % 2 == 1 and prev_skip != i - 1: print(f"WTF --> {i}, {E}, {current_task}, {local_path}.") prev_skip = i total_error_task_counter[current_task] += 1 continue # raise E assert add_meta.get("task", "") in known_dpo_tasks arr_duration = arr.shape[0] / 25 arr = reshift(arr) # print("after shift and pad", arr.shape) arr = arr.reshape( -1, ) # print(arr.shape) out_mm = np.memmap( out_mmap_path, dtype=np.uint16, mode="r+", shape=(n_offs + arr.size,), ) out_mm[n_offs : n_offs + arr.size] = arr # print(f"offset is: {n_offs}") # break # write it once out_mm.flush() del out_mm add_metas = [] add_metas.append(add_meta) tot_duration_dict[row["preference"]] += arr_duration total_task_counter[add_meta.get("task", "gen")] += arr_duration # print(add_metas) write_jsonl( add_metas, os.path.join(out_metas_path), do_append=bool(n_offs != 0), ) if "idx_list" not in datasets_info[add_meta["dataset"]]: datasets_info[add_meta["dataset"]]["idx_list"] = [curr_idx] else: datasets_info[add_meta["dataset"]]["idx_list"].append(curr_idx) curr_idx += 1 n_offs += arr.size write_json(datasets_info, out_info_filepath) print( f"Total {curr_idx} clips, {total_different_prompts} different prompts, {total_different_tags} different tags, {total_different_neg_tags} different negative tags" ) for k, v in tot_duration_dict.items(): print(f"{round(v / 60 / 60):,} hours of {k}") for k, v in total_task_counter.items(): print(f"{k}: {round(v / 60 / 60, 1)} hours") for k, v in total_error_task_counter.items(): print(f"🚨 Error {k}: {v}") print("\n--- Gender Distribution ---") total_gender_count = sum(gender_counter.values()) for gender, count in sorted(gender_counter.items()): percentage = (count / total_gender_count * 100) if total_gender_count > 0 else 0 print(f" {gender}: {count:,} ({percentage:.1f}%)") print("\n--- Negative Tags Usage ---") total_neg_tags = sum(neg_tags_counter.values()) for tag_status, count in sorted(neg_tags_counter.items()): percentage = (count / total_neg_tags * 100) if total_neg_tags > 0 else 0 print(f" {tag_status}: {count:,} ({percentage:.1f}%)") print("\n--- Control Slider Usage ---") for slider, count in sorted(control_slider_counter.items()): percentage = (count / curr_idx * 100) if curr_idx > 0 else 0 print(f" {slider}: {count:,} ({percentage:.1f}% of clips)") print("Done")