{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "18dd74ce", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "import numpy as np\n", "import os\n", "base_dir = \"/app2/suno/data/christian/outputs/v3-distill-data-ctx-t2/\"\n", "\n", "# get all files in each directory\n", "\n", "#model_name = \"16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_220k\"\n", "#model_name = \"v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_2k_last\"\n", "#model_name = \"v3_flow_distill_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last\"\n", "model_name = \"4n_25hz_2b_flow_5e5_sft_t8_500k\"\n", "\n", "def process_dir(dirpath, num_files=10):\n", " \"\"\"\n", " Loads metadata and upsampled vae files for indices 0 to num_files-1.\n", " Returns:\n", " metadata_dicts: list of dicts, one per index\n", " upsampled_vae_filepaths: list of filepaths, one per index\n", " \"\"\"\n", " metadata_dicts = []\n", " upsampled_vae_filepaths = []\n", " for idx in range(num_files):\n", " metadata_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_{idx}__metadata.npz\")\n", " upsampled_vae_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_{idx}_upsampled_vae.npz\")\n", " if not os.path.exists(metadata_filepath):\n", " metadata_dicts.append(None)\n", " upsampled_vae_filepaths.append(None)\n", " continue\n", " metadata = np.load(metadata_filepath, allow_pickle=True)\n", " metadata_dict = {}\n", " for key in metadata.keys():\n", " metadata_dict[key] = metadata[key].tolist()\n", " metadata_dicts.append(metadata_dict)\n", " upsampled_vae_filepaths.append(upsampled_vae_filepath)\n", " return metadata_dicts, upsampled_vae_filepaths" ] }, { "cell_type": "code", "execution_count": null, "id": "3fac5752", "metadata": {}, "outputs": [], "source": [ "from typing import List, Dict, Any, Optional, Callable\n", "import math\n", "\n", "def _get(d: Dict[str, Any], key: str, default=None):\n", " \"\"\"Supports dotted keys like 'diffusion.steps'.\"\"\"\n", " cur = d\n", " for part in key.split('.'):\n", " if isinstance(cur, dict) and part in cur:\n", " cur = cur[part]\n", " else:\n", " return default\n", " return cur\n", "\n", "def rank_items(\n", " items: List[Dict[str, Any]],\n", " criteria: List[Dict[str, Any]],\n", ") -> List[Dict[str, Any]]:\n", " \"\"\"\n", " Rank items by a weighted sum of normalized criteria.\n", " Each criterion: {\n", " 'key': 'hoot_cer' | 'stereo_width' | 'diffusion.steps' | ...,\n", " 'mode': 'min' | 'max' | 'target',\n", " 'weight': float (default 1.0),\n", " 'target': float (required if mode='target'),\n", " 'transform': Optional[callable(value)->float]\n", " }\n", " Lower total score = better.\n", " Missing/non-finite values are treated as the WORST for all modes.\n", " \"\"\"\n", " n = len(items)\n", " if n == 0:\n", " return []\n", "\n", " # Build per-criterion raw arrays (after transform, but before normalization)\n", " raw_matrix = []\n", " for c in criteria:\n", " key = c['key']\n", " mode = c['mode']\n", " weight = float(c.get('weight', 1.0))\n", " transform: Optional[Callable[[float], float]] = c.get('transform')\n", " if mode not in ('min', 'max', 'target'):\n", " raise ValueError(f\"Unknown mode {mode!r}\")\n", "\n", " vals = []\n", " for it in items:\n", " v = _get(it, key, None)\n", " if v is None:\n", " vals.append(float('inf')) # mark missing as non-finite \"worst\"\n", " continue\n", " try:\n", " v = float(transform(v) if transform else v)\n", " except Exception:\n", " v = float('inf')\n", "\n", " if mode == 'target':\n", " tgt = c['target']\n", " v = abs(v - tgt) # smaller distance is better\n", " vals.append(v)\n", "\n", " raw_matrix.append((mode, weight, vals))\n", "\n", " # Normalize each criterion to [0,1] where 0 = best, 1 = worst\n", " norm_matrix = []\n", " for (mode, weight, vals) in raw_matrix:\n", " # consider only finite values when computing min/max\n", " finite_vals = [v for v in vals if math.isfinite(v)]\n", " if len(finite_vals) == 0:\n", " # everything missing/inf -> every item gets worst cost 1.0\n", " norm = [1.0] * n\n", " else:\n", " vmin = min(finite_vals)\n", " vmax = max(finite_vals)\n", " if vmax == vmin:\n", " # all equal among finite -> neutralize effect\n", " norm = [0.0 if math.isfinite(v) else 1.0 for v in vals]\n", " else:\n", " # min/target: smaller is better -> (v - vmin)/(vmax - vmin)\n", " # max: larger is better -> invert afterward\n", " base = []\n", " for v in vals:\n", " if not math.isfinite(v):\n", " base.append(1.0) # missing -> worst\n", " else:\n", " b = (v - vmin) / (vmax - vmin)\n", " base.append(b)\n", " if mode == 'max':\n", " base = [1.0 - b for b in base] # invert so 0 is best\n", " norm = base\n", "\n", " norm_matrix.append((weight, norm))\n", "\n", " # Weighted sum (lower total score is better)\n", " scores = []\n", " for i in range(n):\n", " s = 0.0\n", " for (weight, norm) in norm_matrix:\n", " s += weight * norm[i]\n", " scores.append(s)\n", "\n", " # Attach scores and return sorted copy (best first)\n", " ranked = []\n", " for it, sc in zip(items, scores):\n", " # don’t mutate originals; attach score in the copy\n", " ranked.append({**it, '__score__': float(sc)})\n", " ranked.sort(key=lambda d: d['__score__'])\n", " return ranked\n", "\n", "criteria = [\n", " {'key': 'hoot_cer', 'mode': 'min', 'weight': 1.0}, # or {'key': 'cer', ...} if that's your field\n", " {'key': 'stereo_width', 'mode': 'target', 'target': 0.15, 'weight': 1.0},\n", " {\"key\": \"shimmer_score\", \"mode\": \"min\", \"weight\": 1.0},\n", " {\"key\": \"lufs_db\", \"mode\": \"target\", \"target\": -16.0, \"weight\": 1.0},\n", "]" ] }, { "cell_type": "code", "execution_count": null, "id": "c5c5c1b1", "metadata": {}, "outputs": [], "source": [ "# create metas # in this case using ear score\n", "from tqdm import tqdm\n", "metas = []\n", "with_history = 0\n", "use_hoot = False\n", "\n", "# get all directories in base_dir\n", "dirs = os.listdir(base_dir)\n", "print(len(dirs))\n", "\n", "pbar = tqdm(dirs)\n", "\n", "for dirpath in pbar:\n", " # now call the function to get the metadata for each index\n", " metadata_dicts, upsampled_vae_filepaths = process_dir(dirpath)\n", " metadata_dicts = [m for m in metadata_dicts if isinstance(m, dict)] # drop None/bad rows\n", " if not metadata_dicts:\n", " continue # nothing to rank in this dir\n", " ranked = rank_items(metadata_dicts, criteria)\n", "\n", " # now lets create metas by assigning pairs of metadata as positive and negative\n", " # so the positive should always have a higher score than the negative\n", " # for simplicity lets just always use the best as positive and the others as negatives\n", " for i in range(len(ranked)):\n", " for j in range(i+1, len(ranked)):\n", " if ranked[i]['__score__'] > ranked[j]['__score__']:\n", " pos_vae_latents_filepath = upsampled_vae_filepaths[i]\n", " neg_vae_latents_filepath = upsampled_vae_filepaths[j]\n", " else:\n", " pos_vae_latents_filepath = upsampled_vae_filepaths[j]\n", " neg_vae_latents_filepath = upsampled_vae_filepaths[i]\n", "\n", " metas.append({\n", " \"id\": dirpath,\n", " \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n", " \"neg_vae_latents_filepath\": neg_vae_latents_filepath\n", " })\n", " pbar.set_description(f\"Processed {dirpath}, {len(metas)} valid\")" ] }, { "cell_type": "code", "execution_count": null, "id": "fc34ef10", "metadata": {}, "outputs": [], "source": [ "# split metas into train and test\n", "# split into train and test\n", "train_metas = metas[:int(len(metas) * 0.95)]\n", "val_metas = metas[int(len(metas) * 0.95):]\n", "\n", "print(len(train_metas), len(val_metas))" ] }, { "cell_type": "code", "execution_count": null, "id": "06b9a106", "metadata": {}, "outputs": [], "source": [ "train_metas[0]" ] }, { "cell_type": "markdown", "id": "b8b776a9", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4b3803a7", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import write_jsonl\n", "\n", "\n", "version = \"t7\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model\"\n", "\n", "write_jsonl(train_metas, f\"metas_tr_{version}.jsonl\" )\n", "write_jsonl(val_metas, f\"metas_val_{version}.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "0f4c4c89", "metadata": {}, "outputs": [], "source": [ "import os\n", "root_dir = \"/app/suno/data/dpo/diff2_v1\"\n", "root_dir = \"/app2/suno/data/christian/outputs/v3-base-data-ctx-t0/\"\n", "\n", "# get all files in the root_dir\n", "files = os.listdir(root_dir)\n", "\n", "# print the first 10 files\n", "print(files[:10])" ] }, { "cell_type": "code", "execution_count": null, "id": "a84d4026", "metadata": {}, "outputs": [], "source": [ "metas = []\n", "from tqdm import tqdm\n", "\n", "for dirname in tqdm(os.listdir(root_dir)):\n", " # get all files in the directory\n", " files = os.listdir(os.path.join(root_dir, dirname))\n", " vae_filepaths = [file for file in files if \"vae\" in file]\n", " if len(vae_filepaths) == 2:\n", " metas.append({\n", " \"id\": dirname,\n", " \"pos_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[0]),\n", " \"neg_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[1])\n", " })\n", "\n", " if len(vae_filepaths) == 3:\n", " metas.append({\n", " \"id\": dirname,\n", " \"pos_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[0]),\n", " \"neg_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[1])\n", " })\n", " metas.append({\n", " \"id\": dirname,\n", " \"pos_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[2]),\n", " \"neg_vae_latents_filepath\": os.path.join(root_dir, dirname, vae_filepaths[2])\n", " })\n", "\n", "\n", "\n", "print(len(metas))\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b4cd3571", "metadata": {}, "outputs": [], "source": [ "# split into train and test\n", "train_metas = metas[:int(len(metas) * 0.95)]\n", "val_metas = metas[int(len(metas) * 0.95):]\n", "\n", "print(len(train_metas), len(val_metas))\n", "train_metas[0]\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1b40afea", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import write_jsonl\n", "\n", "version = \"t8\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model\"\n", "\n", "write_jsonl(train_metas, output_dir + f\"/metas_tr_{version}.jsonl\" )\n", "write_jsonl(val_metas, output_dir + f\"/metas_val_{version}.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "0f73842b", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "1c20497f", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "7742bb00", "metadata": {}, "source": [ "# From prod preference data" ] }, { "cell_type": "code", "execution_count": null, "id": "01ff8fac", "metadata": {}, "outputs": [], "source": [ "import os\n", "import pandas as pd\n", "import numpy as np\n", "from tqdm import tqdm\n", "from suno_utils.utils.text import write_jsonl\n", "\n", "version = \"t13\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model\"" ] }, { "cell_type": "code", "execution_count": null, "id": "76cd757a", "metadata": {}, "outputs": [], "source": [ "dataframes = [\n", " \"/home/tony/Data/Preference/up_v2_d3/interesting_clips_ahi_d3_20250608.pkl\",\n", " \"/home/tony/Data/Preference/up_v2_d4/interesting_clips_ahi_d4_20250714.pkl\",\n", " #\"/home/tony/Data/Preference/up_v2_d5/interesting_clips_ahi_d5_20250824.pkl\"\n", " \"/home/tony/Data/Preference/up_v2_d5/fully_merged_up_v2_d5.pkl\"\n", "]\n", "\n", "root_dirs = [\n", " \"/app2/suno/data/dpo/diff2_v2_d3\",\n", " \"/app2/suno/data/dpo/diff2_v2_d4\",\n", " \"/app2/suno/data/dpo/diff2_v2_d5\"\n", "]\n", "\n", "\n", "# load each dataframe\n", "# select the last 5% of rows as validation use the rest as training\n", "# then create two dataframes, one for training and one for validation\n", "# the odd index in the dataframe is the neative and even is the positive\n", "\n", "train_metas = []\n", "val_metas = []\n", "for root_dir, df in zip(root_dirs, dataframes):\n", " df = pd.read_pickle(df)\n", " print(len(df))\n", " # iterate over the dataframe and crate the pair metas\n", " # make indices an array of odd indices\n", " subset_metas = []\n", " indices = np.arange(len(df))\n", " indices = indices[indices % 2 == 1]\n", " # iterate over the indices and create the pair metas\n", " for index in tqdm(indices):\n", " pos_item_id = df.iloc[index][\"id\"]\n", " neg_item_id = df.iloc[index - 1][\"id\"]\n", "\n", " pos_vae_latents_filepath = os.path.join(root_dir, f\"{pos_item_id}_vae.npz\")\n", " neg_vae_latents_filepath = os.path.join(root_dir, f\"{neg_item_id}_vae.npz\")\n", "\n", " if os.path.exists(pos_vae_latents_filepath) and os.path.exists(neg_vae_latents_filepath):\n", " subset_metas.append({\n", " \"id\": pos_item_id,\n", " \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n", " \"neg_vae_latents_filepath\": neg_vae_latents_filepath\n", " })\n", " train_metas.extend(subset_metas[:int(len(subset_metas) * 0.95)])\n", " val_metas.extend(subset_metas[int(len(subset_metas) * 0.95):])" ] }, { "cell_type": "code", "execution_count": null, "id": "aa0e6496", "metadata": {}, "outputs": [], "source": [ "from joblib import Parallel, delayed\n", "\n", "dataframes = [\n", " \"/home/tony/Data/Preference/up_v2_d3/interesting_clips_ahi_d3_20250608.pkl\",\n", " \"/home/tony/Data/Preference/up_v2_d4/interesting_clips_ahi_d4_20250714.pkl\",\n", " #\"/home/tony/Data/Preference/up_v2_d5/interesting_clips_ahi_d5_20250824.pkl\"\n", " \"/home/tony/Data/Preference/up_v2_d5/fully_merged_up_v2_d5.pkl\"\n", "]\n", "\n", "root_dirs = [\n", " \"/app2/suno/data/dpo/diff2_v2_d3\",\n", " \"/app2/suno/data/dpo/diff2_v2_d4\",\n", " \"/app2/suno/data/dpo/diff2_v2_d5\"\n", "]\n", "\n", "def process_index(index, df, root_dir):\n", " pos_item_id = df.iloc[index][\"id\"]\n", " neg_item_id = df.iloc[index - 1][\"id\"]\n", "\n", " pos_vae_latents_filepath = os.path.join(root_dir, f\"{pos_item_id}_vae.npz\")\n", " neg_vae_latents_filepath = os.path.join(root_dir, f\"{neg_item_id}_vae.npz\")\n", "\n", " if os.path.exists(pos_vae_latents_filepath) and os.path.exists(neg_vae_latents_filepath):\n", " return {\n", " \"id\": pos_item_id,\n", " \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n", " \"neg_vae_latents_filepath\": neg_vae_latents_filepath\n", " }\n", " else:\n", " return None\n", "\n", "train_metas = []\n", "val_metas = []\n", "for root_dir, df_path in zip(root_dirs, dataframes):\n", " df = pd.read_pickle(df_path)\n", " print(len(df))\n", " indices = np.arange(len(df))\n", " indices = indices[indices % 2 == 1]\n", "\n", " # Parallelize the processing of indices\n", " results = Parallel(n_jobs=-1, backend=\"loky\")(\n", " delayed(process_index)(index, df, root_dir) for index in tqdm(indices)\n", " )\n", " # Filter out None results\n", " subset_metas = [meta for meta in results if meta is not None]\n", "\n", " train_metas.extend(subset_metas[:int(len(subset_metas) * 0.95)])\n", " val_metas.extend(subset_metas[int(len(subset_metas) * 0.95):])" ] }, { "cell_type": "code", "execution_count": null, "id": "1bd20b71", "metadata": {}, "outputs": [], "source": [ "print(len(train_metas))\n", "print(len(val_metas))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5535a7f1", "metadata": {}, "outputs": [], "source": [ "write_jsonl(train_metas, output_dir + f\"/metas_tr_{version}.jsonl\")\n", "write_jsonl(val_metas, output_dir + f\"/metas_val_{version}.jsonl\")" ] }, { "cell_type": "markdown", "id": "08030881", "metadata": {}, "source": [ "# From prod preference data (for pretraining)" ] }, { "cell_type": "code", "execution_count": null, "id": "0b709212", "metadata": {}, "outputs": [], "source": [ "version = \"t0\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model/pretrain\"" ] }, { "cell_type": "code", "execution_count": null, "id": "3c38aeb7", "metadata": {}, "outputs": [], "source": [ "root_dirs = [\n", " \"/app2/suno/data/dpo/diff2_v2_d3\",\n", " \"/app2/suno/data/dpo/diff2_v2_d4\",\n", " \"/app2/suno/data/dpo/diff2_v2_d5\",\n", " #\"/app2/suno/data/dpo/bluejay_t1_npz\",\n", "]\n", "\n", "train_metas = []\n", "val_metas = []\n", "\n", "for root_dir in root_dirs:\n", " subset_metas = []\n", " files = os.listdir(root_dir)\n", " for file in tqdm(files):\n", " if \"vae\" in file:\n", " pos_vae_latents_filepath = os.path.join(root_dir, file)\n", " neg_vae_latents_filepath = os.path.join(root_dir, file)\n", " if os.path.exists(pos_vae_latents_filepath) and os.path.exists(neg_vae_latents_filepath):\n", " subset_metas.append({\n", " \"id\": file,\n", " \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n", " \"neg_vae_latents_filepath\": neg_vae_latents_filepath\n", " })\n", " print(len(subset_metas))\n", " train_metas.extend(subset_metas[:int(len(subset_metas) * 0.95)])\n", " val_metas.extend(subset_metas[int(len(subset_metas) * 0.95):])\n", "\n", "print(len(train_metas))\n", "print(len(val_metas))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c3eb164c", "metadata": {}, "outputs": [], "source": [ "write_jsonl(train_metas, output_dir + f\"/metas_tr_{version}.jsonl\")\n", "write_jsonl(val_metas, output_dir + f\"/metas_val_{version}.jsonl\")" ] }, { "cell_type": "markdown", "id": "4b877b11", "metadata": {}, "source": [ "# From labelmaker" ] }, { "cell_type": "code", "execution_count": null, "id": "4f615752", "metadata": {}, "outputs": [], "source": [ "import os\n", "import pandas as pd\n", "import numpy as np\n", "from tqdm import tqdm\n", "from suno_utils.utils.text import write_jsonl\n", "\n", "\n", "model_name = \"4n_25hz_2b_flow_5e5_sft_t8_500k\"" ] }, { "cell_type": "code", "execution_count": null, "id": "2bb30b66", "metadata": {}, "outputs": [], "source": [ "df" ] }, { "cell_type": "code", "execution_count": null, "id": "00e7c89c", "metadata": {}, "outputs": [], "source": [ "\n", "filepath = \"/home/christian/code/christian/metadata/labelmaker/t5/dpo_annotations_export_0_7.csv\"\n", "df = pd.read_csv(filepath)\n", "\n", "# convert this to a dict lookup that goes from clip_id to the label\n", "label_lookup = {}\n", "for index, row in df.iterrows():\n", " clip_id = row[\"clip_id\"]\n", " chosen_file_index = row[\"chosen_file_index\"]\n", " unchosen_file_index = row[\"unchosen_file_index\"]\n", " label_lookup[clip_id] = {\n", " \"chosen_file_index\": chosen_file_index,\n", " \"unchosen_file_index\": unchosen_file_index,\n", " \"agreement\": row.get(\"agreement\", None),\n", " \"num_ratings\": row.get(\"num_ratings\", None),\n", " }\n", "\n", "print(len(label_lookup))" ] }, { "cell_type": "code", "execution_count": null, "id": "a038526b", "metadata": {}, "outputs": [], "source": [ "ids = list(label_lookup.keys())\n", "label_lookup[ids[100]]" ] }, { "cell_type": "code", "execution_count": null, "id": "14e4979d", "metadata": {}, "outputs": [], "source": [ "def process_dir(base_dir, dirname):\n", " results = []\n", " for n in range(10):\n", " upsampled_vae_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}_upsampled_vae.npz\")\n", " results.append(upsampled_vae_filepath)\n", " return results\n" ] }, { "cell_type": "code", "execution_count": null, "id": "52d180d8", "metadata": {}, "outputs": [], "source": [ "base_dir = \"/app2/suno/data/christian/outputs/v3-base-data-ctx-t2/\"\n", "# get all the dirs in the base_dir\n", "dirs = os.listdir(base_dir)\n", "\n", "metas = []\n", "for dirname in tqdm(dirs):\n", " # get the label from the label_lookup\n", " label = label_lookup.get(dirname, None)\n", " if label is None:\n", " continue\n", "\n", " results = process_dir(base_dir, dirname)\n", " \n", " pos_vae_latents_filepath = results[label[\"chosen_file_index\"]]\n", " neg_vae_latents_filepath = results[label[\"unchosen_file_index\"]]\n", "\n", " # confirm both files exist\n", " if not os.path.exists(pos_vae_latents_filepath) or not os.path.exists(neg_vae_latents_filepath):\n", " continue\n", "\n", " metas.append({\n", " \"id\": dirname,\n", " \"pos_vae_latents_filepath\": results[label[\"chosen_file_index\"]],\n", " \"neg_vae_latents_filepath\": results[label[\"unchosen_file_index\"]],\n", " \"agreement\": label[\"agreement\"],\n", " \"num_ratings\": label[\"num_ratings\"],\n", " })\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "4d5b9409", "metadata": {}, "outputs": [], "source": [ "\n", "import random\n", "\n", "# Shuffle metas before splitting\n", "random.shuffle(metas)\n", "\n", "train_metas = metas[:int(len(metas) * 0.95)]\n", "val_metas = metas[int(len(metas) * 0.95):]\n", "\n", "print(len(train_metas))\n", "print(len(val_metas))\n", "\n", "version = \"t15\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model\"\n", "\n", "write_jsonl(train_metas, output_dir + f\"/metas_tr_{version}.jsonl\")\n", "write_jsonl(val_metas, output_dir + f\"/metas_val_{version}.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "c6bea30f", "metadata": {}, "outputs": [], "source": [ "val_metas[10]" ] }, { "cell_type": "code", "execution_count": null, "id": "cfcd73af", "metadata": {}, "outputs": [], "source": [ "label_lookup[val_metas[10][\"id\"]]" ] }, { "cell_type": "code", "execution_count": null, "id": "e9f88259", "metadata": {}, "outputs": [], "source": [ "# test some of the local data\n", "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n", "\n", "from suno_utils.tasks.dac_vae_fixed_25hz import (\n", " preload_models as preload_codec_models,\n", " decode as codec_decode,\n", " encode as codec_encode,\n", " decode_stream_to_full_audio,\n", ")\n", "\n", "_ = preload_codec_models(CODEC_FILEPATH)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "863d39a2", "metadata": {}, "outputs": [], "source": [ "# test metas are correct\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "\n", "filepath = \"/home/christian/code/christian/metadata/reward_model/metas_val_t15.jsonl\"\n", "metas = read_jsonl(filepath)\n", "print(len(metas))" ] }, { "cell_type": "code", "execution_count": null, "id": "f9109558", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "# pick a random meta, then decode the positive and negative vae_latents\n", "meta = metas[np.random.randint(0, len(metas))]\n", "\n", "label = label_lookup[meta[\"id\"]]\n", "print(label)\n", "\n", "# decode the positive and negative vae_latents\n", "pos_vae_latents = np.load(meta[\"pos_vae_latents_filepath\"])[\"vae_latents\"]\n", "neg_vae_latents = np.load(meta[\"neg_vae_latents_filepath\"])[\"vae_latents\"]\n", "\n", "# decode the positive and negative vae_latents\n", "pos_audio = codec_decode(pos_vae_latents)\n", "neg_audio = codec_decode(neg_vae_latents)\n", "\n", "# play the positive and negative audios\n", "print(\"pos_audio\")\n", "pos_audio.play()\n", "print(\"neg_audio\")\n", "neg_audio.play()" ] }, { "cell_type": "code", "execution_count": null, "id": "17a02d55", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "9d73d189", "metadata": {}, "source": [ "# From synthetic (modal)" ] }, { "cell_type": "code", "execution_count": null, "id": "1d8570cc", "metadata": {}, "outputs": [], "source": [ "import os\n", "import pandas as pd\n", "import numpy as np\n", "from tqdm import tqdm\n", "from suno_utils.utils.text import write_jsonl\n", "\n", "version = \"t14\"\n", "output_dir = \"/home/christian/code/christian/metadata/reward_model\"\n", "base_dir = \"/app2/suno/data/christian/outputs/v3-base-data-ctx-discogs-subset-t0/\"\n", "model_name = \"16n_25hz_v45_infill_shared_flow_resume_1_75m\"\n", "\n", "# aws s3 sync s3://suno-data/christian/outputs/v3-base-data-ctx-discogs-subset-t0/ /app2/suno/data/christian/outputs/v3-base-data-ctx-discogs-subset-t0" ] }, { "cell_type": "code", "execution_count": null, "id": "4eecd1c9", "metadata": {}, "outputs": [], "source": [ "def process_dir(base_dir, dirname):\n", " upsampled_vae_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_0_upsampled_vae.npz\")\n", " original_vae_filepath = os.path.join(base_dir, dirname, f\"{dirname}_original_vae.npz\")\n", " if os.path.exists(upsampled_vae_filepath) and os.path.exists(original_vae_filepath):\n", " return upsampled_vae_filepath, original_vae_filepath\n", " else:\n", " return None, None\n", "\n", "# get all the dirs in the base_dir\n", "dirs = os.listdir(base_dir)\n", "print(len(dirs))\n", "\n", "RATE_HZ = 25\n", "total_hours = 0\n", "\n", "metas = []\n", "for dirname in tqdm(dirs):\n", "\n", " neg_vae_latents_filepath, pos_vae_latents_filepath = process_dir(base_dir, dirname)\n", " if neg_vae_latents_filepath is None or pos_vae_latents_filepath is None:\n", " continue\n", "\n", " # load both files\n", " pos_vae_latents = np.load(pos_vae_latents_filepath)[\"vae_latents\"]\n", " neg_vae_latents = np.load(neg_vae_latents_filepath)[\"vae_latents\"]\n", "\n", " n_tokens = pos_vae_latents.shape[0]\n", " duration_s = n_tokens / RATE_HZ\n", " duration_hours = duration_s / 3600\n", " total_hours += duration_hours\n", "\n", " if pos_vae_latents.shape != neg_vae_latents.shape:\n", " min_length = min(pos_vae_latents.shape[0], neg_vae_latents.shape[0])\n", " pos_vae_latents = pos_vae_latents[:min_length]\n", " neg_vae_latents = neg_vae_latents[:min_length]\n", "\n", " # check that the shapes are the same\n", " if pos_vae_latents.shape != neg_vae_latents.shape:\n", " print(f\"Shape mismatch for {dirname}\")\n", " print(pos_vae_latents.shape)\n", " print(neg_vae_latents.shape)\n", " continue\n", "\n", " metas.append({ \n", " \"id\": dirname,\n", " \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n", " \"neg_vae_latents_filepath\": neg_vae_latents_filepath\n", " })\n", "\n", "\n", "train_metas = metas[:int(len(metas) * 0.95)]\n", "val_metas = metas[int(len(metas) * 0.95):]\n", "\n", "print(len(train_metas))\n", "print(len(val_metas))\n", "\n", "write_jsonl(train_metas, output_dir + f\"/metas_tr_{version}.jsonl\")\n", "write_jsonl(val_metas, output_dir + f\"/metas_val_{version}.jsonl\")\n", "\n", "print(f\"Total hours: {total_hours:0.2f}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "540945a3", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_diff", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.9" } }, "nbformat": 4, "nbformat_minor": 5 }