{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "9a03b89d", "metadata": {}, "outputs": [], "source": [ "import torch\n" ] }, { "cell_type": "code", "execution_count": null, "id": "83533a33", "metadata": {}, "outputs": [], "source": [ "import os\n", "import numpy as np\n", "from tqdm import tqdm\n", "\n", "metas = []\n", "BASE_LOCAL_DIR = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t2/\"\n", "\n", "# Get up to 1000 valid subdirectories\n", "local_dirs = [f for f in os.listdir(BASE_LOCAL_DIR) if os.path.isdir(os.path.join(BASE_LOCAL_DIR, f))]\n", "#local_dirs = local_dirs[:10000]\n", "\n", "# remove \"memmap\" from the list \n", "local_dirs = [f for f in local_dirs if \"memmap\" not in f]\n", "\n", "for local_dir in tqdm(local_dirs):\n", " dir_path = os.path.join(BASE_LOCAL_DIR, local_dir)\n", " \n", " # Find metadata files\n", " metadata_files = sorted([\n", " f for f in os.listdir(dir_path) \n", " if f.endswith(\"__metadata.npz\")\n", " ])\n", " \n", " # Require at least 2 files to form a pair\n", " if len(metadata_files) < 2:\n", " continue\n", " \n", " # Only use the first 2 (sorted so higher steps come last)\n", " metadata_files = metadata_files[:2]\n", "\n", " def extract_metadata(file_path):\n", " data = np.load(file_path, allow_pickle=True)\n", " md = {}\n", " for key, value in data.items():\n", " if key == \"diffusion\":\n", " if isinstance(value, np.ndarray) and value.dtype == object:\n", " value = value.item()\n", " for k, v in value.items():\n", " md[f\"diffusion_{k}\"] = v\n", " else:\n", " md[key] = value\n", " return md\n", "\n", " # Extract steps from filenames to identify positive/negative\n", " def get_step(file_name):\n", " parts = file_name.split(\"_steps_\")\n", " if len(parts) > 1:\n", " step_part = parts[1].split(\"__\")[0] # Safely handles __metadata\n", " return int(step_part)\n", " return 0\n", "\n", " steps_and_files = [(get_step(f), f) for f in metadata_files]\n", " steps_and_files.sort(key=lambda x: x[0]) # ascending: [negative, positive]\n", " \n", " # Get metadata for each file\n", " neg_step, neg_file = steps_and_files[0]\n", " pos_step, pos_file = steps_and_files[1]\n", " \n", " neg_path = os.path.join(dir_path, neg_file)\n", " pos_path = os.path.join(dir_path, pos_file)\n", "\n", " neg_meta = extract_metadata(neg_path)\n", " pos_meta = extract_metadata(pos_path)\n", "\n", "\n", " pair_meta = {\n", " \"id\": local_dir,\n", " \"directory\": dir_path,\n", " \"pos_npz_filename\": pos_meta[\"filename\"].item().replace(\".mp3\", \"_upsampled_vae.npz\"),\n", " \"neg_npz_filename\": neg_meta[\"filename\"].item().replace(\".mp3\", \"_upsampled_vae.npz\"),\n", " \"pos_steps\": pos_step,\n", " \"neg_steps\": neg_step,\n", " \"text\": pos_meta[\"text\"].item(),\n", " \"tags\": pos_meta[\"tags\"].item(),\n", " }\n", " # confirm that the text and tags are the same\n", " # if they are not the same continue\n", " if neg_meta[\"text\"] != pos_meta[\"text\"] or neg_meta[\"tags\"] != pos_meta[\"tags\"]:\n", " continue\n", "\n", " metas.append(pair_meta)\n", "\n", "print(f\"Collected {len(metas)} metadata entries.\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "7841a33e", "metadata": {}, "outputs": [], "source": [ "# save metas to jsonl\n", "from suno_utils.utils.text import write_jsonl\n", "\n", "out_dir = \"/home/christian/code/christian/metadata/reward_model\"\n", "\n", "# split metas into train and val\n", "\n", "# split metas into train and val\n", "train_metas = metas[: int(len(metas) * 0.95)]\n", "val_metas = metas[int(len(metas) * 0.95) :]\n", "\n", "write_jsonl(train_metas, f\"{out_dir}/tr.jsonl\")\n", "write_jsonl(val_metas, f\"{out_dir}/val.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "9fd130ad", "metadata": {}, "outputs": [], "source": [ "metas[0]" ] }, { "cell_type": "code", "execution_count": null, "id": "e05ce8e4", "metadata": {}, "outputs": [], "source": [ "# we start with a metas file that contains pairs of clips\n", "# we have path to positive and negative npz files containing the vae latents\n", "# we will load the latents, sample a random chunk, and then return the positive and negative latent chunks\n", "\n", "import torch\n", "import numpy as np\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "\n", "class RewardModelDataset(torch.utils.data.Dataset):\n", " def __init__(self, metas_path, chunk_size=750):\n", " self.metas_path = metas_path\n", " self.metas = read_jsonl(metas_path)\n", " print(f\"Loaded {len(self.metas)} metas\")\n", " self.chunk_size = chunk_size\n", "\n", " def __len__(self):\n", " return len(self.metas)\n", "\n", " def __getitem__(self, idx):\n", " meta = self.metas[idx]\n", " positive_latents = np.load(os.path.join(meta[\"directory\"], meta[\"pos_npz_filename\"]))[\"vae_latents\"]\n", " negative_latents = np.load(os.path.join(meta[\"directory\"], meta[\"neg_npz_filename\"]))[\"vae_latents\"]\n", "\n", " # sample a random chunk\n", " if positive_latents.shape[0] > self.chunk_size:\n", " start_idx = np.random.randint(0, positive_latents.shape[0] - self.chunk_size)\n", " end_idx = start_idx + self.chunk_size\n", " positive_latents = positive_latents[start_idx:end_idx]\n", " negative_latents = negative_latents[start_idx:end_idx]\n", " else:\n", " print(positive_latents.shape, negative_latents.shape)\n", " # zero pad the latents\n", " positive_latents = np.pad(positive_latents, ((self.chunk_size - positive_latents.shape[0], 0), (0, 0)), mode=\"wrap\")\n", " negative_latents = np.pad(negative_latents, ((self.chunk_size - negative_latents.shape[0], 0), (0, 0)), mode=\"wrap\")\n", "\n", " return positive_latents, negative_latents" ] }, { "cell_type": "code", "execution_count": null, "id": "cc89d7d5", "metadata": {}, "outputs": [], "source": [ "dataset = RewardModelDataset(metas_path=\"/home/christian/code/christian/metadata/reward_model/metas_v3_bootstrap_t2.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "559a68b2", "metadata": {}, "outputs": [], "source": [ "dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b54452cd", "metadata": {}, "outputs": [], "source": [ "import math\n", "import torch\n", "import torch.nn as nn\n", "\n", "\n", "class SinusoidalPositionalEncoding(nn.Module):\n", " def __init__(self, dim, max_len=2048):\n", " super().__init__()\n", " pe = torch.zeros(max_len, dim)\n", " position = torch.arange(0, max_len).unsqueeze(1)\n", " div_term = torch.exp(torch.arange(0, dim, 2) * -(math.log(10000.0) / dim))\n", " pe[:, 0::2] = torch.sin(position * div_term)\n", " pe[:, 1::2] = torch.cos(position * div_term)\n", " self.register_buffer(\"pe\", pe)\n", "\n", " def forward(self, x):\n", " # x: (B, T, D)\n", " seq_len = x.size(1)\n", " return x + self.pe[:seq_len].unsqueeze(0).to(x.dtype) # (1, T, D)\n", "\n", "\n", "class SimpleTransformerEncoder(nn.Module):\n", " def __init__(self, vae_dim, embed_dim=768, num_layers=6, num_heads=12, ff_dim=2048, dropout=0.1, max_len=2048):\n", " super().__init__()\n", " self.input_proj = nn.Linear(vae_dim, embed_dim)\n", " self.pos_encoding = SinusoidalPositionalEncoding(embed_dim, max_len)\n", "\n", " encoder_layer = nn.TransformerEncoderLayer(\n", " d_model=embed_dim,\n", " nhead=num_heads,\n", " dim_feedforward=ff_dim,\n", " dropout=dropout,\n", " batch_first=True\n", " )\n", " self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n", " self.norm = nn.LayerNorm(embed_dim)\n", " self.output_head = nn.Linear(embed_dim, 1)\n", "\n", " def forward(self, x):\n", " # x: (B, T, vae_dim)\n", " x = self.input_proj(x) # (B, T, embed_dim)\n", " x = self.pos_encoding(x) # add positional encodings\n", " x = self.transformer(x) # (B, T, embed_dim)\n", " x = self.norm(x)\n", " return self.output_head(x) # (B, T, 1)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "393a1cf3", "metadata": {}, "outputs": [], "source": [ "model = SimpleTransformerEncoder(\n", " vae_dim=128,\n", " embed_dim=768,\n", " num_layers=6,\n", " num_heads=12,\n", " ff_dim=3072,\n", " dropout=0.1,\n", " max_len=1024\n", ").cuda()\n", "\n", "x = torch.randn(2, 512, 128).cuda()\n", "out = model(x)\n", "print(out.shape) # ✅ (2, 512, 1)\n", "\n", "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)" ] }, { "cell_type": "code", "execution_count": null, "id": "b1781907", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from torch.nn import functional as F\n", "\n", "\n", "def bradley_terry_loss(\n", " r_i: torch.Tensor, r_j: torch.Tensor, labels: torch.Tensor\n", ") -> torch.Tensor:\n", " \"\"\"\n", " Compute Bradley-Terry loss for paired comparisons.\n", "\n", " Args:\n", " r_i: Logits/scores for first options in pairs, shape (batch_size,)\n", " r_j: Logits/scores for second options in pairs, shape (batch_size,)\n", " labels: Binary tensor indicating whether first option (0) or second option (1)\n", " was preferred, shape (batch_size,)\n", "\n", " Returns:\n", " Mean loss value as a torch.Tensor\n", " \"\"\"\n", " # Compute negative log likelihood using logsigmoid for numerical stability\n", " loss = -(\n", " (labels) * F.logsigmoid(r_j - r_i) + (1 - labels) * F.logsigmoid(r_i - r_j)\n", " )\n", " return loss.mean()\n" ] }, { "cell_type": "code", "execution_count": null, "id": "26b878cb", "metadata": {}, "outputs": [], "source": [ "# very simple training loop\n", "from tqdm import tqdm\n", "\n", "pbar = tqdm(dataloader)\n", "for batch_idx, batch in enumerate(pbar):\n", " optimizer.zero_grad()\n", " pos_vae, neg_vae = batch\n", " pos_vae = pos_vae.cuda()\n", " neg_vae = neg_vae.cuda()\n", "\n", " pos_logits = model(pos_vae).mean(dim=1).mean(dim=1)\n", " neg_logits = model(neg_vae).mean(dim=1).mean(dim=1)\n", "\n", " # compute the accuracy \n", " preds = (pos_logits > neg_logits).float()\n", " acc = (preds == torch.ones_like(preds)).float().mean()\n", "\n", " loss = bradley_terry_loss(pos_logits, neg_logits, torch.zeros_like(pos_logits))\n", " loss.backward()\n", " pbar.set_description(f\"Loss: {loss.item():.4f} acc: {acc.item():.4f}\")\n", " optimizer.step()" ] }, { "cell_type": "code", "execution_count": null, "id": "c7da7fa5", "metadata": {}, "outputs": [], "source": [ "# test reward model\n", "import glob\n", "import os\n", "import numpy as np\n", "import torch\n", "\n", "\n", "base_dir = \"/home/christian/code/christian/notebooks/outputs/seed_test/\"\n", "\n", "# find all npz files in the base_dir\n", "npz_files = glob.glob(os.path.join(base_dir, \"*.npz\"))\n", "\n", "# score with the model\n", "for npz_file in npz_files:\n", " data = np.load(npz_file)\n", " vae_latents = torch.from_numpy(data[\"vae_latents\"]).cuda()[:750, :].unsqueeze(0)\n", " with torch.no_grad():\n", " logits = model(vae_latents).mean(dim=1).mean(dim=1)\n", "\n", " print(logits)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "da46610c", "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 }