{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class AudioPatcher:\n",
    "    def __init__(self, patch_size):\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "    def to_patches(self, audio):\n",
    "        # audio: (bs, 2, seq_len)\n",
    "        bs, channels, seq_len = audio.shape\n",
    "        assert seq_len % self.patch_size == 0, \"Sequence length must be divisible by patch size\"\n",
    "        num_patches = seq_len // self.patch_size\n",
    "        patches = audio.view(bs, channels, num_patches, self.patch_size)  # (bs, 2, num_patches, patch_size)\n",
    "        return patches\n",
    "\n",
    "    def flatten_patches(self, patches):\n",
    "        # patches: (bs, 2, num_patches, patch_size)\n",
    "        bs, channels, num_patches, patch_size = patches.shape\n",
    "        flattened = patches.permute(0, 2, 1, 3).reshape(bs, num_patches * channels, patch_size)\n",
    "        # (bs, num_patches * 2, patch_size)\n",
    "        return flattened\n",
    "\n",
    "    def unflatten_patches(self, flattened, original_channels=2):\n",
    "        # flattened: (bs, num_patches * channels, patch_size)\n",
    "        bs, total_patches, patch_size = flattened.shape\n",
    "        num_patches = total_patches // original_channels\n",
    "        patches = flattened.view(bs, num_patches, original_channels, patch_size).permute(0, 2, 1, 3)\n",
    "        # (bs, 2, num_patches, patch_size)\n",
    "        return patches\n",
    "\n",
    "    def reconstruct_audio(self, patches):\n",
    "        # patches: (bs, 2, num_patches, patch_size)\n",
    "        bs, channels, num_patches, patch_size = patches.shape\n",
    "        audio = patches.reshape(bs, channels, num_patches * patch_size)\n",
    "        return audio\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "import torchaudio\n",
    "\n",
    "audio, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n",
    "audio = audio[:, :131072*2].unsqueeze(0)\n",
    "\n",
    "\n",
    "patcher = AudioPatcher(patch_size=1024)\n",
    "\n",
    "patches = patcher.to_patches(audio)\n",
    "print(patches.shape)\n",
    "flattened = patcher.flatten_patches(patches)\n",
    "print(flattened.shape)\n",
    "reconstructed = patcher.unflatten_patches(flattened)\n",
    "print(reconstructed.shape)\n",
    "reconstructed_audio = patcher.reconstruct_audio(reconstructed)\n",
    "print(reconstructed_audio.shape)\n",
    "\n",
    "ipd.Audio(reconstructed_audio.squeeze().numpy(), rate=sr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "class SinusoidalPositionalEncoding(torch.nn.Module):\n",
    "    def __init__(self, hidden_dim):\n",
    "        super().__init__()\n",
    "        position = torch.arange(10000).unsqueeze(1)\n",
    "        div_term = torch.exp(\n",
    "            torch.arange(0, hidden_dim, 2) * -(math.log(10000.0) / hidden_dim)\n",
    "        )\n",
    "        pe = torch.zeros(10000, hidden_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: (batch, num_patches, hidden_dim)\n",
    "        return x + self.pe[: x.size(1)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Refiner(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_dim=1024,\n",
    "        num_heads=8,\n",
    "        num_transformer_layers=12,\n",
    "        dropout=0.1,\n",
    "        patch_size=1024,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        self.patcher = AudioPatcher(patch_size)\n",
    "\n",
    "        # Position embedding\n",
    "        self.pos_embed = SinusoidalPositionalEncoding(hidden_dim)\n",
    "\n",
    "        # input layer\n",
    "        self.input_layer = torch.nn.Linear(patch_size, hidden_dim)\n",
    "\n",
    "        # Transformer encoder\n",
    "        encoder_layer = torch.nn.TransformerEncoderLayer(\n",
    "            d_model=hidden_dim,\n",
    "            nhead=num_heads,\n",
    "            dim_feedforward=hidden_dim * 4,\n",
    "            dropout=dropout,\n",
    "        )\n",
    "        self.transformer = torch.nn.TransformerEncoder(\n",
    "            encoder_layer, num_transformer_layers\n",
    "        )\n",
    "\n",
    "        self.output_layer = torch.nn.Linear(hidden_dim, patch_size)\n",
    "\n",
    "    def forward(self, input_audio):\n",
    "        # input_audio shape: (batch_size, 2, seq_len)\n",
    "        # first split into patches, which becomes (batch_size, 2, num_patches, patch_size)\n",
    "        patches = self.patcher.to_patches(input_audio)\n",
    "        flattened = self.patcher.flatten_patches(patches)\n",
    "        x = self.input_layer(flattened)  # (batch_size, num_patches, hidden_dim)\n",
    "        x = self.pos_embed(x)  # (batch_size, num_patches, hidden_dim)\n",
    "        x = self.transformer(x)  # (batch_size, num_patches, hidden_dim)\n",
    "        x = self.output_layer(x)  # (batch_size, num_patches, 1)\n",
    "        x = self.patcher.unflatten_patches(x)  # (batch_size, 2, seq_len)\n",
    "\n",
    "        # fold the patches back together\n",
    "        output_audio = self.patcher.reconstruct_audio(x)\n",
    "\n",
    "        return input_audio + output_audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "audio, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n",
    "audio = audio[:, :131072*2].unsqueeze(0)\n",
    "print(audio.shape)\n",
    "\n",
    "refiner = Refiner(hidden_dim=1024, num_heads=8, num_transformer_layers=12, dropout=0.1, patch_size=1024)\n",
    "# count model parameters\n",
    "total_params = sum(p.numel() for p in refiner.parameters())\n",
    "print(f\"Total parameters: {total_params/1e6:.2f}M\")\n",
    "\n",
    "output_audio = refiner(audio)\n",
    "\n",
    "ipd.Audio(output_audio.squeeze().detach().cpu().numpy(), rate=sr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "filepath = \"/home/christian/code/christian/metadata/dedup/combined_title_clusters_jac-0.80.jsonl\"\n",
    "\n",
    "clusters = read_jsonl(filepath)\n",
    "\n",
    "# get the first cluster\n",
    "cluster = clusters[0]\n",
    "\n",
    "# get the first item in the cluster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "rand_idx = random.randint(0, len(clusters) - 1)\n",
    "\n",
    "# get the first item in the cluster\n",
    "item = clusters[rand_idx]\n",
    "\n",
    "# get the title\n",
    "\n",
    "print(item[\"size\"], item[\"items\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "cluster_sizes = [len(c[\"items\"]) for c in clusters]\n",
    "\n",
    "# plot the cluster sizes\n",
    "plt.hist(cluster_sizes, bins=100)\n",
    "plt.xlabel(\"Cluster Size\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.yscale(\"log\")\n",
    "plt.title(\"Cluster Size Distribution\")\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a json file where we map id to the number of items in the cluster\n",
    "id_to_weight = {}\n",
    "for cluster in clusters:\n",
    "    for item in cluster[\"items\"]:\n",
    "        size = cluster[\"size\"]\n",
    "        id_to_weight[item[\"id\"]] = 1 / size\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(id_to_weight))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(id_to_weight[list(id_to_weight.keys())[100]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "output_filepath = \"/home/christian/code/christian/metadata/dedup/id_to_weight_jac-0.80.json\"\n",
    "with open(output_filepath, \"w\") as f:\n",
    "    json.dump(id_to_weight, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = \"/home/vibert/data/suno_data_monitor/outputs/v9_full_20251017_151623/train/analysis_results.json\"\n",
    "\n",
    "with open(filepath, \"r\") as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "for key, value in data.items():\n",
    "    print(key, value)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}