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 = 6016 # max 240s of audio def make_dataset( input_df, output_data_dir: str, is_val=False, npz_dir="/app/suno/data/dpo/30b_npz", t_data_memmap=N_TOKENS_AUDIO, original_npz_dir="/app/suno/data/dpo/13b_s32_npz", # this is the npz directory of the non-cycled ): print(f"t_data_memmap is set to: {t_data_memmap}") def reshift(arr): sem_start_idx = 0 sem_end_idx = len(arr) - 1 semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS] coarse_arr = arr[:, SEMANTIC_N_CODEBOOKS:] coarse_start_idx = int(round(sem_start_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ)) coarse_end_idx = int(round(sem_end_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ)) assert sem_end_idx >= 0 and coarse_start_idx >= 0 assert not (sem_end_idx > len(semantic_arr) or coarse_end_idx > len(coarse_arr)) # get array segments arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy() arr_c = coarse_arr[coarse_start_idx:coarse_end_idx, :].copy() assert arr_s.max() <= SEMANTIC_PAD_TOKEN assert arr_c.max() <= COARSE_PAD_TOKEN assert len(arr_s) == len(arr_c) # concat and stack if len(arr_c) < t_data_memmap: arr_c = np.pad( arr_c, ((0, t_data_memmap - len(arr_c)), (0, 0)), constant_values=COARSE_PAD_TOKEN, mode="constant", ) arr_s = np.pad( arr_s, ((0, t_data_memmap - len(arr_s)), (0, 0)), constant_values=SEMANTIC_PAD_TOKEN, mode="constant", ) arr = np.concatenate([arr_s, arr_c], axis=-1) arr = arr.astype(np.uint16) assert arr.shape == (t_data_memmap, SEMANTIC_N_CODEBOOKS + COARSE_N_CODEBOOKS) return arr 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") # gather the data _ = np.memmap(out_mmap_path, dtype=np.uint16, mode="w+", shape=(1,)) n_offs = 0 tot_duration_dict = defaultdict(int) datasets_info = defaultdict(dict) n = 0 total_different_prompts = 0 negative_prompt = "" total_task_counter = Counter() total_error_task_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 "" 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 else: negative_prompt = current_prompt 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): raise ValueError(f"File does not exist: {local_path}") try: temp_npz = np.load(local_path) if "v4.0_raw" in temp_npz: arr = temp_npz["v4.0_raw"] elif "v3.5_raw" in temp_npz: if "cycle" not in local_path and "13b" 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 :) assert arr.shape[1] == 13 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": [ row["tags"] if not pd.isna(row["tags"]) else "" ], # tags is a list, do you know :) "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", "", ] try: if current_task == "cover": cover_prompt_arr = temp_npz["cover_arr"] assert ( cover_prompt_arr.shape[0] <= t_data_memmap ) # can't be longer than 4 mins :) assert cover_prompt_arr.shape[1] == 13 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 ) # can't be longer than 4 mins :) assert artist_prompt_arr.shape[1] == 13 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 == "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] == 13 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": 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] == 13 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"] 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 # print(arr.shape) 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 chosen_end_s = None original_duration_s = arr_duration if row["total_clip_s"] >= 0: # is not a total clip chosen_end_s = row["total_start_s"] + arr_duration original_duration_s = row["total_clip_s"] elif not ( row["continue_at"] is not None and row["continue_at"] >= 0 ): # is not a continue if arr_duration < 60 * 4 - 5: chosen_end_s = arr_duration # for full clips we do know it has an edding, other wise, we don't know add_meta["end_s"] = chosen_end_s # this needs to be... a bit more complicated, only works with concat! add_meta["original_duration_s"] = original_duration_s 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"] = [n] else: datasets_info[add_meta["dataset"]]["idx_list"].append(n) n += 1 n_offs += arr.size write_json(datasets_info, out_info_filepath) print(f"Total {n} clips, {total_different_prompts} different prompts") 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("Done")