{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torchaudio\n",
    "from torch import nn\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from einops import rearrange\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "class MSDConfig:\n",
    "    tag_names = (\n",
    "        'rock',\n",
    "        'pop',\n",
    "        'indie',\n",
    "        'alternative',\n",
    "        'electronic',\n",
    "        'hip-hop',\n",
    "        'metal',\n",
    "        'jazz',\n",
    "        'punk',\n",
    "        'folk',\n",
    "        'alternative rock',\n",
    "        'indie rock',\n",
    "        'dance',\n",
    "        'hard rock',\n",
    "        '00s',\n",
    "        'soul',\n",
    "        'hardcore',\n",
    "        '80s',\n",
    "        'country',\n",
    "        'classic rock',\n",
    "        'punk rock',\n",
    "        'blues',\n",
    "        'chillout',\n",
    "        'experimental',\n",
    "        'heavy metal',\n",
    "        'death metal',\n",
    "        '90s',\n",
    "        'reggae',\n",
    "        'progressive rock',\n",
    "        'ambient',\n",
    "        'acoustic',\n",
    "        'beautiful',\n",
    "        'british',\n",
    "        'rnb',\n",
    "        'funk',\n",
    "        'metalcore',\n",
    "        'mellow',\n",
    "        'world',\n",
    "        'guitar',\n",
    "        'trance',\n",
    "        'indie pop',\n",
    "        'christian',\n",
    "        'house',\n",
    "    )\n",
    "    n_tags = len(tag_names)\n",
    "\n",
    "class Res2DMaxPoolModule(nn.Module):\n",
    "    def __init__(self, input_channels, output_channels, pooling=2):\n",
    "        super(Res2DMaxPoolModule, self).__init__()\n",
    "        self.conv_1 = nn.Conv2d(input_channels, output_channels, 3, padding=1)\n",
    "        self.bn_1 = nn.BatchNorm2d(output_channels)\n",
    "        self.conv_2 = nn.Conv2d(output_channels, output_channels, 3, padding=1)\n",
    "        self.bn_2 = nn.BatchNorm2d(output_channels)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.mp = nn.MaxPool2d(pooling)\n",
    "\n",
    "        # residual\n",
    "        self.diff = False\n",
    "        if input_channels != output_channels:\n",
    "            self.conv_3 = nn.Conv2d(input_channels, output_channels, 3, padding=1)\n",
    "            self.bn_3 = nn.BatchNorm2d(output_channels)\n",
    "            self.diff = True\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.bn_2(self.conv_2(self.relu(self.bn_1(self.conv_1(x)))))\n",
    "        if self.diff:\n",
    "            x = self.bn_3(self.conv_3(x))\n",
    "        out = x + out\n",
    "        out = self.mp(self.relu(out))\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResFrontEnd(nn.Module):\n",
    "    \"\"\"\n",
    "    Evaluation of CNN based Music Tagging.\n",
    "    Won et al., 2020\n",
    "    \n",
    "    Note that, different from the original work, we only stack 3 convolutional layers instead of 7.\n",
    "    After the convolution layers, we flatten the time-frequency representation to be a vector.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, conv_ndim, attention_ndim, nfreq, nharmonics=1):\n",
    "        super(ResFrontEnd, self).__init__()\n",
    "        self.input_bn = nn.BatchNorm2d(nharmonics)\n",
    "        self.layer1 = Res2DMaxPoolModule(nharmonics, conv_ndim, pooling=(2, 2))\n",
    "        self.layer2 = Res2DMaxPoolModule(conv_ndim, conv_ndim, pooling=(2, 2))\n",
    "        self.layer3 = Res2DMaxPoolModule(conv_ndim, conv_ndim, pooling=(2, 1))\n",
    "        fc_ndim = nfreq // 2 // 2 // 2 * conv_ndim\n",
    "        self.fc = nn.Linear(fc_ndim, attention_ndim)\n",
    "\n",
    "    def forward(self, hcqt):\n",
    "        # batch normalization\n",
    "        out = self.input_bn(hcqt)\n",
    "\n",
    "        # CNN\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "\n",
    "        # permute and channel control\n",
    "        b, c, f, t = out.shape\n",
    "        out = out.permute(0, 3, 1, 2)  # batch, time, conv_ndim, freq\n",
    "        out = out.contiguous().view(b, t, -1)  # batch, time, fc_ndim\n",
    "        out = self.fc(out)  # batch, time, attention_ndim\n",
    "        return out\n",
    "\n",
    "\n",
    "# Transformer modules\n",
    "\"\"\"\n",
    "    Referenced PyTorch implementation of Vision Transformer by Lucidrains.\n",
    "    https://github.com/lucidrains/vit-pytorch.git\n",
    "\"\"\"\n",
    "class Residual(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(x, **kwargs) + x\n",
    "\n",
    "\n",
    "class PreNorm(nn.Module):\n",
    "    def __init__(self, dim, fn):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(self.norm(x), **kwargs)\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dim, hidden_dim, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(hidden_dim, dim),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head * heads\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\n",
    "\n",
    "        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)\n",
    "        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        b, n, _, h = *x.shape, self.heads\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=-1)\n",
    "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)\n",
    "\n",
    "        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale\n",
    "        mask_value = -torch.finfo(dots.dtype).max\n",
    "\n",
    "        if mask is not None:\n",
    "            mask = F.pad(mask.flatten(1), (1, 0), value=True)\n",
    "            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'\n",
    "            mask = mask[:, None, :] * mask[:, :, None]\n",
    "            dots.masked_fill_(~mask, mask_value)\n",
    "            del mask\n",
    "\n",
    "        attn = dots.softmax(dim=-1)\n",
    "\n",
    "        out = torch.einsum('bhij,bhjd->bhid', attn, v)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        out = self.to_out(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([])\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(\n",
    "                nn.ModuleList(\n",
    "                    [\n",
    "                        Residual(\n",
    "                            PreNorm(\n",
    "                                dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)\n",
    "                            )\n",
    "                        ),\n",
    "                        Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))),\n",
    "                    ]\n",
    "                )\n",
    "            )\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x, mask=mask)\n",
    "            x = ff(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class MusicTaggingTransformer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        conv_ndim=16,\n",
    "        n_mels=128,\n",
    "        sample_rate=22050,\n",
    "        n_fft=1024,\n",
    "        f_min=0,\n",
    "        f_max=11025,\n",
    "        attention_ndim=256,\n",
    "        attention_nheads=8,\n",
    "        attention_nlayers=4,\n",
    "        attention_max_len=512,\n",
    "        dropout=0.1,\n",
    "        n_seq_cls=1,\n",
    "        n_token_cls=1,\n",
    "    ):\n",
    "        super(MusicTaggingTransformer, self).__init__()\n",
    "        # Input preprocessing\n",
    "        self.spec = self.spec = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate,\n",
    "                                                                     n_fft=n_fft,\n",
    "                                                                     f_min=f_min,\n",
    "                                                                     f_max=f_max,\n",
    "                                                                     n_mels=n_mels,\n",
    "                                                                     power=2)\n",
    "        self.amplitude_to_db = torchaudio.transforms.AmplitudeToDB()\n",
    "\n",
    "        # Input embedding\n",
    "        self.frontend = ResFrontEnd(conv_ndim, attention_ndim, n_mels)\n",
    "\n",
    "        # Positional embedding\n",
    "        self.pos_embedding = nn.Parameter(torch.randn(1, attention_max_len + 1, attention_ndim))\n",
    "        self.cls_token = nn.Parameter(torch.randn(attention_ndim))\n",
    "\n",
    "        # transformer\n",
    "        self.transformer = Transformer(\n",
    "            attention_ndim,\n",
    "            attention_nlayers,\n",
    "            attention_nheads,\n",
    "            attention_ndim // attention_nheads,\n",
    "            attention_ndim * 4,\n",
    "            dropout,\n",
    "        )\n",
    "        self.to_latent = nn.Identity()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        # projection for sequence classification\n",
    "        self.mlp_head = nn.Sequential(\n",
    "            nn.LayerNorm(attention_ndim), nn.Linear(attention_ndim, n_seq_cls)\n",
    "        )\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): (batch, time)\n",
    "\n",
    "        Returns:\n",
    "            x (torch.Tensor): (batch, n_seq_cls)\n",
    "\n",
    "        \"\"\"\n",
    "        # Input preprocessing\n",
    "        x = self.spec(x)\n",
    "        x = self.amplitude_to_db(x)\n",
    "        x = x.unsqueeze(1)\n",
    "\n",
    "        # Input embedding\n",
    "        x = self.frontend(x)\n",
    "\n",
    "        # Positional embedding with a [CLS] token\n",
    "        cls_token = self.cls_token.repeat(x.shape[0], 1, 1)\n",
    "        x = torch.cat((cls_token, x), dim=1)\n",
    "        x += self.pos_embedding[:, : x.size(1)]\n",
    "        x = self.dropout(x)\n",
    "\n",
    "        # transformer\n",
    "        x = self.transformer(x)\n",
    "\n",
    "        # projection for sequence classification\n",
    "        x = self.to_latent(x[:, 0])\n",
    "        x = self.mlp_head(x)\n",
    "        x = self.sigmoid(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MusicTaggingTransformer(\n",
       "  (spec): MelSpectrogram(\n",
       "    (spectrogram): Spectrogram()\n",
       "    (mel_scale): MelScale()\n",
       "  )\n",
       "  (amplitude_to_db): AmplitudeToDB()\n",
       "  (frontend): ResFrontEnd(\n",
       "    (input_bn): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (layer1): Res2DMaxPoolModule(\n",
       "      (conv_1): Conv2d(1, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv_2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU()\n",
       "      (mp): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
       "      (conv_3): Conv2d(1, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (layer2): Res2DMaxPoolModule(\n",
       "      (conv_1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv_2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU()\n",
       "      (mp): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (layer3): Res2DMaxPoolModule(\n",
       "      (conv_1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv_2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (bn_2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU()\n",
       "      (mp): MaxPool2d(kernel_size=(2, 1), stride=(2, 1), padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (fc): Linear(in_features=2048, out_features=256, bias=True)\n",
       "  )\n",
       "  (transformer): Transformer(\n",
       "    (layers): ModuleList(\n",
       "      (0-3): 4 x ModuleList(\n",
       "        (0): Residual(\n",
       "          (fn): PreNorm(\n",
       "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "            (fn): Attention(\n",
       "              (to_qkv): Linear(in_features=256, out_features=768, bias=False)\n",
       "              (to_out): Sequential(\n",
       "                (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "                (1): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (1): Residual(\n",
       "          (fn): PreNorm(\n",
       "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "            (fn): FeedForward(\n",
       "              (net): Sequential(\n",
       "                (0): Linear(in_features=256, out_features=1024, bias=True)\n",
       "                (1): GELU(approximate='none')\n",
       "                (2): Dropout(p=0.1, inplace=False)\n",
       "                (3): Linear(in_features=1024, out_features=256, bias=True)\n",
       "                (4): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (to_latent): Identity()\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (mlp_head): Sequential(\n",
       "    (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "    (1): Linear(in_features=256, out_features=50, bias=True)\n",
       "  )\n",
       "  (sigmoid): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load model# Teacher model\n",
    "checkpoint_path = \"/home/christian/code/christian/checkpoints/teacher.ckpt\"\n",
    "model = MusicTaggingTransformer(conv_ndim=128, attention_ndim=256, n_seq_cls=50)\n",
    "S = torch.load(checkpoint_path)\n",
    "S = {k[8:]: v for k, v in S.items()}\n",
    "model.load_state_dict(S)\n",
    "\n",
    "\n",
    "def get_labels(logits):\n",
    "    # Get top tags from logits\n",
    "    sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
    "    \n",
    "    # Map indices to tag names using MSDConfig\n",
    "    tags = []\n",
    "    for i, value in zip(sorted_indices.squeeze().tolist(), sorted_logits.squeeze().tolist()):\n",
    "        # Only include tags with logits above threshold\n",
    "        if i < len(MSDConfig.tag_names):\n",
    "            tags.append((MSDConfig.tag_names[i], value))\n",
    "    return tags\n",
    "\n",
    "model.eval()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model loaded!\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Ditto(\n",
       "  (music_encoder): MusicEncoder(\n",
       "    (music_encoder): MusicFM_MERTLong(\n",
       "      (preprocessor_melspec_2048): MelSTFT(\n",
       "        (mel_stft): MelSpectrogram(\n",
       "          (spectrogram): Spectrogram()\n",
       "          (mel_scale): MelScale()\n",
       "        )\n",
       "        (amplitude_to_db): AmplitudeToDB()\n",
       "      )\n",
       "      (quantizer_melspec_2048_0): RandomProjectionQuantizer()\n",
       "      (conv): Conv2dSubsampling(\n",
       "        (conv): Sequential(\n",
       "          (0): Res2dModule(\n",
       "            (conv1): Conv2d(1, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "            (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (relu): ReLU()\n",
       "            (conv3): Conv2d(1, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "            (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "          (1): Res2dModule(\n",
       "            (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "            (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (relu): ReLU()\n",
       "            (conv3): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "            (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (linear): Linear(in_features=16384, out_features=1024, bias=True)\n",
       "      )\n",
       "      (conformer): Wav2Vec2ConformerEncoder(\n",
       "        (embed_positions): Wav2Vec2ConformerRotaryPositionalEmbedding()\n",
       "        (pos_conv_embed): Wav2Vec2ConformerPositionalConvEmbedding(\n",
       "          (conv): ParametrizedConv1d(\n",
       "            1024, 1024, kernel_size=(128,), stride=(1,), padding=(64,), groups=16\n",
       "            (parametrizations): ModuleDict(\n",
       "              (weight): ParametrizationList(\n",
       "                (0): _WeightNorm()\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (padding): Wav2Vec2ConformerSamePadLayer()\n",
       "          (activation): GELUActivation()\n",
       "        )\n",
       "        (layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "        (layers): ModuleList(\n",
       "          (0-11): 12 x Wav2Vec2ConformerEncoderLayer(\n",
       "            (ffn1_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn1): Wav2Vec2ConformerFeedForward(\n",
       "              (intermediate_dropout): Dropout(p=0.1, inplace=False)\n",
       "              (intermediate_dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "              (intermediate_act_fn): SiLU()\n",
       "              (output_dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              (output_dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (self_attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "            (self_attn): Wav2Vec2ConformerSelfAttention(\n",
       "              (linear_q): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (linear_k): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (linear_v): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (linear_out): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (conv_module): Wav2Vec2ConformerConvolutionModule(\n",
       "              (layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "              (pointwise_conv1): Conv1d(1024, 2048, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              (glu): GLU(dim=1)\n",
       "              (depthwise_conv): Conv1d(1024, 1024, kernel_size=(31,), stride=(1,), padding=(15,), groups=1024, bias=False)\n",
       "              (batch_norm): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "              (activation): SiLU()\n",
       "              (pointwise_conv2): Conv1d(1024, 1024, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (ffn2_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn2): Wav2Vec2ConformerFeedForward(\n",
       "              (intermediate_dropout): Dropout(p=0.1, inplace=False)\n",
       "              (intermediate_dense): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "              (intermediate_act_fn): SiLU()\n",
       "              (output_dense): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "              (output_dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (linear): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "      (loss): CrossEntropyLoss()\n",
       "    )\n",
       "  )\n",
       "  (text_encoder): TextEncoder(\n",
       "    (text_encoder): XLMRobertaModel(\n",
       "      (embeddings): XLMRobertaEmbeddings(\n",
       "        (word_embeddings): Embedding(250002, 768, padding_idx=1)\n",
       "        (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "        (token_type_embeddings): Embedding(1, 768)\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (encoder): XLMRobertaEncoder(\n",
       "        (layer): ModuleList(\n",
       "          (0-11): 12 x XLMRobertaLayer(\n",
       "            (attention): XLMRobertaAttention(\n",
       "              (self): XLMRobertaSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): XLMRobertaSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): XLMRobertaIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              (intermediate_act_fn): GELUActivation()\n",
       "            )\n",
       "            (output): XLMRobertaOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (pooler): XLMRobertaPooler(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (activation): Tanh()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (music_projection): Projection(\n",
       "    (linear_1): Linear(in_features=1024, out_features=128, bias=False)\n",
       "    (linear_2): Linear(in_features=128, out_features=128, bias=False)\n",
       "    (layer_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       "  (text_projection): Projection(\n",
       "    (linear_1): Linear(in_features=768, out_features=128, bias=False)\n",
       "    (linear_2): Linear(in_features=128, out_features=128, bias=False)\n",
       "    (layer_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load ditto v2\n",
    "from suno_utils.models.ditto_v2.ditto_v2 import Ditto\n",
    "\n",
    "model_filepath = \"s3://suno-data/minz/models/ditto_v2_epoch_57.pt\"\n",
    "\n",
    "ditto_model = Ditto(\n",
    "        latent_dim=128,\n",
    "        model_path=model_filepath,\n",
    "        is_flash=False,\n",
    "        is_serving=False,\n",
    "    )\n",
    "\n",
    "ditto_model.eval()\n",
    "ditto_model.cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SquimObjective(\n",
       "  (encoder): Encoder(\n",
       "    (conv1d): Conv1d(1, 256, kernel_size=(64,), stride=(32,), bias=False)\n",
       "  )\n",
       "  (dprnn): DPRNN(\n",
       "    (row_rnn): ModuleList(\n",
       "      (0-1): 2 x SingleRNN(\n",
       "        (rnn): LSTM(256, 256, batch_first=True, bidirectional=True)\n",
       "        (proj): Linear(in_features=512, out_features=256, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (col_rnn): ModuleList(\n",
       "      (0-1): 2 x SingleRNN(\n",
       "        (rnn): LSTM(256, 256, batch_first=True, bidirectional=True)\n",
       "        (proj): Linear(in_features=512, out_features=256, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (row_norm): ModuleList(\n",
       "      (0-1): 2 x GroupNorm(1, 256, eps=1e-08, affine=True)\n",
       "    )\n",
       "    (col_norm): ModuleList(\n",
       "      (0-1): 2 x GroupNorm(1, 256, eps=1e-08, affine=True)\n",
       "    )\n",
       "    (conv): Sequential(\n",
       "      (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (1): PReLU(num_parameters=1)\n",
       "    )\n",
       "  )\n",
       "  (branches): ModuleList(\n",
       "    (0-1): 2 x Sequential(\n",
       "      (0): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (linear2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.0, inplace=False)\n",
       "        (dropout2): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (1): AutoPool(\n",
       "        (softmax): Softmax(dim=1)\n",
       "      )\n",
       "      (2): Sequential(\n",
       "        (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (1): PReLU(num_parameters=1)\n",
       "        (2): Linear(in_features=256, out_features=1, bias=True)\n",
       "        (3): RangeSigmoid(\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=256, out_features=1024, bias=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (linear2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.0, inplace=False)\n",
       "        (dropout2): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (1): AutoPool(\n",
       "        (softmax): Softmax(dim=1)\n",
       "      )\n",
       "      (2): Sequential(\n",
       "        (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (1): PReLU(num_parameters=1)\n",
       "        (2): Linear(in_features=256, out_features=1, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load torchsquim\n",
    "from torchaudio.pipelines import SQUIM_OBJECTIVE, SQUIM_SUBJECTIVE\n",
    "\n",
    "objective_model = SQUIM_OBJECTIVE.get_model()\n",
    "objective_model.eval()\n",
    "objective_model.cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load clap model\n",
    "# pip install laion-clap\n",
    "# pip install transformers==4.30.0\n",
    "import laion_clap\n",
    "\n",
    "clap_model = laion_clap.CLAP_Module(enable_fusion=False)\n",
    "clap_model.load_ckpt() # download the default pretrained checkpoint.\n",
    "clap_model.eval()\n",
    "clap_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load demucs for stem split\n",
    "from suno_utils.tasks.demucs import (\n",
    "    preload_models as preload_stem_models,\n",
    "    load_model as load_stem_model,\n",
    "    split_vocals,\n",
    "    AUDIO_STEMS,\n",
    ")\n",
    "\n",
    "stem_model = load_stem_model(device=\"cuda\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of prompts: 170\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/170 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 19%|█▉        | 32/170 [00:01<00:04, 30.76it/s]/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_pytree.py:185: FutureWarning: optree is installed but the version is too old to support PyTorch Dynamo in C++ pytree. C++ pytree support is disabled. Please consider upgrading optree using `python3 -m pip install --upgrade 'optree>=0.13.0'`.\n",
      "  warnings.warn(\n",
      "100%|██████████| 170/170 [00:34<00:00,  4.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of prompts with audio: 170/170\n"
     ]
    }
   ],
   "source": [
    "import joblib\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "# use prompt for evaluation \n",
    "prompts_df = pd.read_csv(\"/home/christian/code/christian/sunoBench/prompts/suno_bench_prompts_v1.csv\")\n",
    "print(f\"Number of prompts: {len(prompts_df)}\")\n",
    "\n",
    "# pull audio from s3 directory based on prompt id\n",
    "#ckpt_name = \"2025-02-04_01-33-55-step_95000_infer\"\n",
    "#ckpt_name = \"2025-01-31_03-14-13-last_ckpt_infer\"\n",
    "\n",
    "#6b\n",
    "ckpt_name = \"2025-02-04_21-04-31-last_ckpt_infer\" # 6b base no cfg\n",
    "#ckpt_name = \"2025-02-10_16-52-41-step_9000_infer\" # 6b finetune\n",
    "base_audio_dir = f\"s3://suno-data/christian/outputs/{ckpt_name}\"\n",
    "\n",
    "local_audio_dir = f\"/home/christian/code/christian/sunoBench/outputs/{ckpt_name}\"\n",
    "os.makedirs(local_audio_dir, exist_ok=True)\n",
    "prompts = []\n",
    "\n",
    "def process_row(row):\n",
    "    s3_filepath = f\"{base_audio_dir}/{row.id}.mp3\"\n",
    "    audio_filepath = f\"{local_audio_dir}/{row.id}.mp3\"\n",
    "    if not os.path.exists(audio_filepath):\n",
    "        os.system(f\"aws s3 cp {s3_filepath} {audio_filepath} > /dev/null 2>&1\")\n",
    "\n",
    "    # load the audio\n",
    "    try:\n",
    "        x, sr = torchaudio.load(audio_filepath)\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading audio: {e}\")\n",
    "        return None\n",
    "\n",
    "    prompt = {\n",
    "        \"prompt_id\": row.id,\n",
    "        \"audio_filepath\": audio_filepath,\n",
    "        \"sr\": sr,\n",
    "        \"audio\": x,\n",
    "        \"tag\": row.tags,\n",
    "        \"text\": row.text\n",
    "    }\n",
    "\n",
    "    return prompt\n",
    "\n",
    "\n",
    "# multiprocess the prompts\n",
    "prompts = joblib.Parallel(n_jobs=32)(joblib.delayed(process_row)(row) for row in tqdm(prompts_df.itertuples(), total=len(prompts_df)))\n",
    "\n",
    "# remove None prompts\n",
    "prompts = [p for p in prompts if p is not None]\n",
    "\n",
    "print(f\"Number of prompts with audio: {len(prompts)}/{len(prompts_df)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/6798c1e3-7c67-418e-b90d-2df13963a33a.mp3\n",
      "female vocalist\n",
      "\n",
      "[Verse 1]\n",
      "Boots on trail at break of dawn\n",
      "Mist and shadows carry on\n",
      "Each step higher than before\n",
      "Nature's secrets to explore\n",
      "\n",
      "[Chorus]\n",
      "Up the mountain we will climb\n",
      "Leave the world of clocks and time\n",
      "Rocky paths beneath our feet\n",
      "Make this journey so complete\n",
      "\n",
      "[Verse 2]\n",
      "Alpine meadows bloom with pride\n",
      "Mountain streams flow by our side\n",
      "Thin air makes us catch our breath\n",
      "Views that make us forget the rest\n",
      "\n",
      "[Bridge]\n",
      "Summit calling from above\n",
      "Testing limits, showing love\n",
      "For these peaks that touch the sky\n",
      "Where the eagles dare to fly\n",
      "\n",
      "[Outro]\n",
      "As we make our way back down\n",
      "carrying memories we found\n",
      "These mountains now part of me\n",
      "Till next time we climb them free\n",
      "\n"
     ]
    }
   ],
   "source": [
    "idx = 167\n",
    "\n",
    "print(prompts[idx][\"audio_filepath\"])\n",
    "print(prompts[idx][\"tag\"])\n",
    "print(prompts[idx][\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autotagging (Reciprocal rank)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 47/170 [00:13<00:34,  3.61it/s, reciprocal_rank=0.45] \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 18\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# resample to 22050\u001b[39;00m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m audio_dict[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msr\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m22050\u001b[39m:\n\u001b[0;32m---> 18\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[43mtorchaudio\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunctional\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maudio_dict\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43msr\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m22050\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;66;03m# convert to mono\u001b[39;00m\n\u001b[1;32m     21\u001b[0m x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mmean(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, keepdim\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torchaudio/functional/functional.py:1519\u001b[0m, in \u001b[0;36mresample\u001b[0;34m(waveform, orig_freq, new_freq, lowpass_filter_width, rolloff, resampling_method, beta)\u001b[0m\n\u001b[1;32m   1515\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m waveform\n\u001b[1;32m   1517\u001b[0m gcd \u001b[38;5;241m=\u001b[39m math\u001b[38;5;241m.\u001b[39mgcd(\u001b[38;5;28mint\u001b[39m(orig_freq), \u001b[38;5;28mint\u001b[39m(new_freq))\n\u001b[0;32m-> 1519\u001b[0m kernel, width \u001b[38;5;241m=\u001b[39m \u001b[43m_get_sinc_resample_kernel\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1520\u001b[0m \u001b[43m    \u001b[49m\u001b[43morig_freq\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1521\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnew_freq\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgcd\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1523\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlowpass_filter_width\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1524\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrolloff\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1525\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresampling_method\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1526\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbeta\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1527\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwaveform\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1528\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwaveform\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1530\u001b[0m resampled \u001b[38;5;241m=\u001b[39m _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width)\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m resampled\n",
      "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torchaudio/functional/functional.py:1426\u001b[0m, in \u001b[0;36m_get_sinc_resample_kernel\u001b[0;34m(orig_freq, new_freq, gcd, lowpass_filter_width, rolloff, resampling_method, beta, device, dtype)\u001b[0m\n\u001b[1;32m   1423\u001b[0m \u001b[38;5;66;03m# we do not use built in torch windows here as we need to evaluate the window\u001b[39;00m\n\u001b[1;32m   1424\u001b[0m \u001b[38;5;66;03m# at specific positions, not over a regular grid.\u001b[39;00m\n\u001b[1;32m   1425\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m resampling_method \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msinc_interp_hann\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m-> 1426\u001b[0m     window \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcos\u001b[49m\u001b[43m(\u001b[49m\u001b[43mt\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mmath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpi\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mlowpass_filter_width\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\n\u001b[1;32m   1427\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1428\u001b[0m     \u001b[38;5;66;03m# sinc_interp_kaiser\u001b[39;00m\n\u001b[1;32m   1429\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m beta \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/_tensor.py:39\u001b[0m, in \u001b[0;36m_handle_torch_function_and_wrap_type_error_to_not_implemented.<locals>.wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m has_torch_function(args):\n\u001b[1;32m     38\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(wrapped, args, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 39\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# score a model \n",
    "reciprocal_ranks = []\n",
    "target_tags = {}\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "for audio_dict in pbar:\n",
    "    # parse the tag from the filename\n",
    "    tag = audio_dict[\"tag\"]\n",
    "    x = audio_dict[\"audio\"]\n",
    "\n",
    "    # trim to 30 seconds\n",
    "    start_idx = audio_dict[\"sr\"]\n",
    "    end_idx = start_idx + audio_dict[\"sr\"]*30\n",
    "    x = x[:, start_idx:end_idx]\n",
    "\n",
    "    # resample to 22050\n",
    "    if audio_dict[\"sr\"] != 22050:\n",
    "        x = torchaudio.functional.resample(x, audio_dict[\"sr\"], 22050)\n",
    "\n",
    "    # convert to mono\n",
    "    x = x.mean(dim=0, keepdim=True)\n",
    "\n",
    "    x = x / x.abs().max().clamp(min=1e-5)\n",
    "    x = x.cuda()\n",
    "    # run model\n",
    "    with torch.no_grad():\n",
    "        y = model(x)\n",
    "        y = y.cpu()\n",
    "\n",
    "    # convert to tags using MSDConfig\n",
    "    tags = get_labels(y)\n",
    "    tags = [t[0] for t in tags]\n",
    "\n",
    "    # look for tags in the prompt tags\n",
    "    prompt_tags = str(audio_dict[\"tag\"])\n",
    "\n",
    "    if prompt_tags not in target_tags:\n",
    "        target_tags[prompt_tags] = []\n",
    "\n",
    "    prompt_tag_ranks = []\n",
    "    for tag in tags:\n",
    "        if tag in prompt_tags:\n",
    "            reciprocal_rank = 1 / (tags.index(tag) + 1)\n",
    "            prompt_tag_ranks.append(reciprocal_rank)\n",
    "            target_tags[prompt_tags].append(reciprocal_rank)\n",
    "\n",
    "    # create poppiness score\n",
    "    # each track has a poppiness score\n",
    "    pop_reciprocal_rank = 1 / (tags.index(\"pop\") + 1)\n",
    "    \n",
    "    if len(prompt_tag_ranks) > 0:\n",
    "        mean_prompt_tag_rank = np.mean(prompt_tag_ranks)\n",
    "    else:\n",
    "        mean_prompt_tag_rank = 0\n",
    "    \n",
    "    reciprocal_ranks.append(mean_prompt_tag_rank)\n",
    "    \n",
    "    # add the reciprocal rank to the prompts dict\n",
    "    audio_dict[\"reciprocal_rank\"] = mean_prompt_tag_rank\n",
    "    audio_dict[\"pop_reciprocal_rank\"] = pop_reciprocal_rank\n",
    "    pbar.set_postfix(reciprocal_rank=np.mean(reciprocal_ranks))\n",
    "\n",
    "print(f\"mean reciprocal rank: {np.mean(reciprocal_ranks):.2f} +/- {np.std(reciprocal_ranks):.2f}\")\n",
    "print(f\"mean poppiness reciprocal rank: {np.mean([p['pop_reciprocal_rank'] for p in prompts]):.2f} +/- {np.std([p['pop_reciprocal_rank'] for p in prompts]):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort the target tags by mean reciprocal rank\n",
    "target_tags_sorted = sorted(target_tags.items(), key=lambda x: np.mean(x[1]), reverse=True)\n",
    "\n",
    "for tag_name, reciprocal_ranks in target_tags_sorted:\n",
    "    # print the tag name and the mean reciprocal rank\n",
    "    # make the print with padding \n",
    "    print(f\"{tag_name: <20} {np.mean(reciprocal_ranks):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CLAP audio<->prompt similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get audio embeddings from audio data\n",
    "# sample rate should be 48000\n",
    "target_tags = {}\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "similarity_scores = []\n",
    "for audio_dict in pbar:\n",
    "    x = audio_dict[\"audio\"]\n",
    "    sr = audio_dict[\"sr\"]\n",
    "    if sr != 48000:\n",
    "        x = torchaudio.functional.resample(x, sr, 48000)\n",
    "\n",
    "    # crop to 30 seconds\n",
    "    #x = x[:, :48000*30]\n",
    "\n",
    "    with torch.no_grad():\n",
    "        audio_data = x.mean(dim=0, keepdim=True)\n",
    "        # peak normalize\n",
    "        audio_data = audio_data / audio_data.abs().max().clamp(min=1e-5)\n",
    "        audio_embed = clap_model.get_audio_embedding_from_data(x = audio_data, use_tensor=True)\n",
    "        audio_embed = audio_embed[0].cpu()\n",
    "\n",
    "        # Get text embedings from texts:\n",
    "        text_data = f\"a song in the style of {audio_dict['tag']}\"\n",
    "        # for some reason we have to pass a list with more than one string\n",
    "        text_embed = clap_model.get_text_embedding([text_data, text_data], use_tensor=True)\n",
    "        text_embed = text_embed[0].cpu()\n",
    "\n",
    "    if text_data not in target_tags:    \n",
    "        target_tags[text_data] = []\n",
    "\n",
    "    # measure cosine similarity\n",
    "    similarity = torch.nn.functional.cosine_similarity(audio_embed, text_embed, dim=0)\n",
    "    pbar.set_postfix(similarity=np.mean(similarity_scores))\n",
    "    similarity_scores.append(similarity.item())\n",
    "    target_tags[text_data].append(similarity.item())\n",
    "\n",
    "    # add the similarity to the prompts dict\n",
    "    audio_dict[\"clap_similarity\"] = similarity.item()\n",
    "\n",
    "print(f\"mean similarity: {np.mean(similarity_scores):.2f} +/- {np.std(similarity_scores):.2f}\")\n",
    "print(f\"max similarity: {np.max(similarity_scores):.2f} min similarity: {np.min(similarity_scores):.2f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort the target tags by mean similarity\n",
    "target_tags_sorted = sorted(target_tags.items(), key=lambda x: np.mean(x[1]), reverse=True)\n",
    "\n",
    "for tag_name, similarities in target_tags_sorted:\n",
    "    print(f\"{tag_name: <20} {np.mean(similarities):.2f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ditto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first compute the text embeddings for the things we want to score\n",
    "classes = [\"male vocalist\", \"female vocalist\", \"instrumental\", \"pop\"]\n",
    "text_embeds = {class_name: None for class_name in classes}\n",
    "for class_name in classes:\n",
    "    text_embed = ditto_model.text_to_latent(class_name)\n",
    "    text_embed = text_embed[0].cpu()\n",
    "    text_embeds[class_name] = text_embed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170/170 [02:11<00:00,  1.30it/s]\n"
     ]
    }
   ],
   "source": [
    "# test ditto v2\n",
    "\n",
    "# Get audio embeddings from audio data\n",
    "target_tags = {}\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "similarity_scores = []\n",
    "for audio_dict in pbar:\n",
    "    x = audio_dict[\"audio\"]\n",
    "    sr = audio_dict[\"sr\"]\n",
    "    if sr != 24_000:\n",
    "        x = torchaudio.functional.resample(x, sr, 24_000)\n",
    "\n",
    "    # crop to 30 seconds\n",
    "    #x = x[:, :48000*30]\n",
    "\n",
    "    with torch.no_grad():\n",
    "        audio_data = x.mean(dim=0, keepdim=True)\n",
    "        # peak normalize\n",
    "        audio_data = audio_data / audio_data.abs().max().clamp(min=1e-5)\n",
    "        # Get text embedings from texts:\n",
    "        text_data = f\"{audio_dict['tag']}\"\n",
    "\n",
    "        text_emb = ditto_model.text_to_latent(text_data)\n",
    "        audio_emb = ditto_model.music_to_latent(audio_data.cuda(), task=\"self_sim\")\n",
    "\n",
    "        # for some reason we have to pass a list with more than one string\n",
    "        audio_emb = audio_emb[0].cpu()\n",
    "        text_emb = text_emb[0].cpu()\n",
    "\n",
    "    # measure similarity between audio and text\n",
    "    similarity = torch.nn.functional.cosine_similarity(audio_emb, text_emb, dim=0)\n",
    "    audio_dict[\"ditto_v2_similarity\"] = similarity.item()\n",
    "\n",
    "    # measure the similarity between audio embed and each class text embed\n",
    "    for class_name in classes:\n",
    "        similarity = torch.nn.functional.cosine_similarity(audio_emb, text_embeds[class_name], dim=0)\n",
    "        similarity_scores.append(similarity.item())\n",
    "        clean_class_name = class_name.replace(\" \", \"_\")\n",
    "        audio_dict[f\"ditto_v2_{clean_class_name}_similarity\"] = similarity.item()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clip duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "170\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAikAAAGdCAYAAADXIOPgAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy80BEi2AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAfKklEQVR4nO3df2zU9eHH8dfV0gPW3tVS2mtHCxUQRCzbUOtNZSiVUg0DrQkiiegIBleMUDelm4p1W8o0UXTBusQJmlg7NYLzBzAstsytoFS6Up0dJbjW0R+Kaa9UOZC+v38Y7utJdR7c9d5cn4/kk/Q+n08/9z7eNH3mc5/71GGMMQIAALBMXLQHAAAAMBAiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICV4qM9gK/r7+/XwYMHlZSUJIfDEe3hAACA78AYo97eXmVmZiouLjznQKyLlIMHDyorKyvawwAAAKegra1NY8aMCcuxrIuUpKQkSV++SJfLFeXRAACA78Ln8ykrKyvwezwcrIuUE2/xuFwuIgUAgDNMOC/VCOlNo4qKCuXm5gYCwuv1avPmzYHtM2fOlMPhCFqWLVsWtsECAIChI6QzKWPGjNGaNWs0ceJEGWP09NNPa968edqzZ4/OP/98SdLSpUv1wAMPBL5n5MiR4R0xAAAYEkKKlLlz5wY9/t3vfqeKigrt3LkzECkjR46Ux+MJ3wgBAMCQdMqfETp+/LiqqqrU19cnr9cbWP/ss88qNTVVU6dOVWlpqT777LNvPY7f75fP5wtaAAAAQr5wdu/evfJ6vTpy5IgSExO1ceNGTZkyRZJ04403auzYscrMzFRjY6PuvvtuNTc366WXXvrG45WXl6usrOzUXwEAAIhJDmOMCeUbjh49qtbWVvX09OjFF1/Uk08+qdra2kCofNX27ds1a9YstbS0aPz48QMez+/3y+/3Bx6f+AhTT08Pn+4BAOAM4fP55Ha7w/r7O+RI+br8/HyNHz9ef/zjH0/a1tfXp8TERG3ZskUFBQXf6XiReJEAACCyIvH7+7TvW9vf3x90JuSrGhoaJEkZGRmn+zQAAGCICemalNLSUhUWFio7O1u9vb2qrKxUTU2Ntm7dqv3796uyslJXX321Ro0apcbGRq1cuVIzZsxQbm5upMYPAABiVEiR0tXVpZtuuknt7e1yu93Kzc3V1q1bddVVV6mtrU1vvPGG1q5dq76+PmVlZamoqEj33HNPpMYOAABi2GlfkxJuXJMCAMCZx8prUgAAACKBSAEAAFYiUgAAgJVCvuMsAHuMW/VatIcQsg/XXBPtIQA4Q3AmBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWCmkSKmoqFBubq5cLpdcLpe8Xq82b94c2H7kyBEVFxdr1KhRSkxMVFFRkTo7O8M+aAAAEPtCipQxY8ZozZo1qq+v1+7du3XllVdq3rx5eu+99yRJK1eu1CuvvKIXXnhBtbW1OnjwoK677rqIDBwAAMQ2hzHGnM4BUlJS9NBDD+n666/X6NGjVVlZqeuvv16S9MEHH+i8885TXV2dLrnkku90PJ/PJ7fbrZ6eHrlcrtMZGhDzxq16LdpDCNmHa66J9hAAREAkfn+f8jUpx48fV1VVlfr6+uT1elVfX69jx44pPz8/sM/kyZOVnZ2turq6bzyO3++Xz+cLWgAAAEKOlL179yoxMVFOp1PLli3Txo0bNWXKFHV0dCghIUHJyclB+6enp6ujo+Mbj1deXi632x1YsrKyQn4RAAAg9oQcKZMmTVJDQ4N27dql2267TYsXL9b7779/ygMoLS1VT09PYGlrazvlYwEAgNgRH+o3JCQkaMKECZKk6dOn65133tGjjz6qBQsW6OjRo+ru7g46m9LZ2SmPx/ONx3M6nXI6naGPHAAAxLTTvk9Kf3+//H6/pk+frmHDhqm6ujqwrbm5Wa2trfJ6vaf7NAAAYIgJ6UxKaWmpCgsLlZ2drd7eXlVWVqqmpkZbt26V2+3WkiVLVFJSopSUFLlcLt1+++3yer3f+ZM9AAAAJ4QUKV1dXbrpppvU3t4ut9ut3Nxcbd26VVdddZUk6ZFHHlFcXJyKiork9/tVUFCgxx9/PCIDBwAAse2075MSbtwnBfjuuE8KAFtYdZ8UAACASCJSAACAlYgUAABgpZDvkwIAAP4f14ZFDmdSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJW44yyAQcXdOQF8V5xJAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlUKKlPLycl100UVKSkpSWlqa5s+fr+bm5qB9Zs6cKYfDEbQsW7YsrIMGAACxL6RIqa2tVXFxsXbu3Klt27bp2LFjmj17tvr6+oL2W7p0qdrb2wPLgw8+GNZBAwCA2Bcfys5btmwJerxhwwalpaWpvr5eM2bMCKwfOXKkPB5PeEYIAACGpNO6JqWnp0eSlJKSErT+2WefVWpqqqZOnarS0lJ99tln33gMv98vn88XtAAAAIR0JuWr+vv7tWLFCl166aWaOnVqYP2NN96osWPHKjMzU42Njbr77rvV3Nysl156acDjlJeXq6ys7FSHAQAAYtQpR0pxcbGampr01ltvBa2/9dZbA19fcMEFysjI0KxZs7R//36NHz/+pOOUlpaqpKQk8Njn8ykrK+tUhwUAAGLEKUXK8uXL9eqrr2rHjh0aM2bMt+6bl5cnSWppaRkwUpxOp5xO56kMAwAAxLCQIsUYo9tvv10bN25UTU2NcnJy/uf3NDQ0SJIyMjJOaYAAAGBoCilSiouLVVlZqZdffllJSUnq6OiQJLndbo0YMUL79+9XZWWlrr76ao0aNUqNjY1auXKlZsyYodzc3Ii8AAAAEJtCipSKigpJX96w7avWr1+vm2++WQkJCXrjjTe0du1a9fX1KSsrS0VFRbrnnnvCNmAAADA0hPx2z7fJyspSbW3taQ0IAABA4m/3AAAASxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArhRQp5eXluuiii5SUlKS0tDTNnz9fzc3NQfscOXJExcXFGjVqlBITE1VUVKTOzs6wDhoAAMS+kCKltrZWxcXF2rlzp7Zt26Zjx45p9uzZ6uvrC+yzcuVKvfLKK3rhhRdUW1urgwcP6rrrrgv7wAEAQGyLD2XnLVu2BD3esGGD0tLSVF9frxkzZqinp0d/+tOfVFlZqSuvvFKStH79ep133nnauXOnLrnkkvCNHAAAxLTTuialp6dHkpSSkiJJqq+v17Fjx5Sfnx/YZ/LkycrOzlZdXd2Ax/D7/fL5fEELAABASGdSvqq/v18rVqzQpZdeqqlTp0qSOjo6lJCQoOTk5KB909PT1dHRMeBxysvLVVZWdqrDAMJm3KrXoj0EAMBXnPKZlOLiYjU1Namqquq0BlBaWqqenp7A0tbWdlrHAwAAseGUzqQsX75cr776qnbs2KExY8YE1ns8Hh09elTd3d1BZ1M6Ozvl8XgGPJbT6ZTT6TyVYQAAgBgW0pkUY4yWL1+ujRs3avv27crJyQnaPn36dA0bNkzV1dWBdc3NzWptbZXX6w3PiAEAwJAQ0pmU4uJiVVZW6uWXX1ZSUlLgOhO3260RI0bI7XZryZIlKikpUUpKilwul26//XZ5vV4+2QMAAEISUqRUVFRIkmbOnBm0fv369br55pslSY888oji4uJUVFQkv9+vgoICPf7442EZLAAAGDpCihRjzP/cZ/jw4Vq3bp3WrVt3yoMCAADgb/cAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsFB/tAQCA7catei3aQzglH665JtpDAE4LZ1IAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYKOVJ27NihuXPnKjMzUw6HQ5s2bQrafvPNN8vhcAQtc+bMCdd4AQDAEBFypPT19WnatGlat27dN+4zZ84ctbe3B5bnnnvutAYJAACGnvhQv6GwsFCFhYXfuo/T6ZTH4znlQQEAAETkmpSamhqlpaVp0qRJuu2223To0KFv3Nfv98vn8wUtAAAAYY+UOXPm6JlnnlF1dbV+//vfq7a2VoWFhTp+/PiA+5eXl8vtdgeWrKyscA8JAACcgUJ+u+d/ueGGGwJfX3DBBcrNzdX48eNVU1OjWbNmnbR/aWmpSkpKAo99Ph+hAgAAIv8R5HPOOUepqalqaWkZcLvT6ZTL5QpaAAAAIh4pH330kQ4dOqSMjIxIPxUAAIghIb/dc/jw4aCzIgcOHFBDQ4NSUlKUkpKisrIyFRUVyePxaP/+/brrrrs0YcIEFRQUhHXgAAAgtoUcKbt379YVV1wReHziepLFixeroqJCjY2Nevrpp9Xd3a3MzEzNnj1bv/nNb+R0OsM3agAAEPNCjpSZM2fKGPON27du3XpaAwIAAJD42z0AAMBSRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwEpECAACsRKQAAAArESkAAMBKRAoAALASkQIAAKxEpAAAACsRKQAAwEpECgAAsBKRAgAArESkAAAAKxEpAADASkQKAACwUny0B4DYNG7Va9EeAgDgDMeZFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYKeRI2bFjh+bOnavMzEw5HA5t2rQpaLsxRvfdd58yMjI0YsQI5efna9++feEaLwAAGCJCjpS+vj5NmzZN69atG3D7gw8+qMcee0xPPPGEdu3ape9973sqKCjQkSNHTnuwAABg6Aj5Zm6FhYUqLCwccJsxRmvXrtU999yjefPmSZKeeeYZpaena9OmTbrhhhtOb7QAAGDICOs1KQcOHFBHR4fy8/MD69xut/Ly8lRXVzfg9/j9fvl8vqAFAAAgrLfF7+jokCSlp6cHrU9PTw9s+7ry8nKVlZWFcxgxh1vMAwCGoqh/uqe0tFQ9PT2Bpa2tLdpDAgAAFghrpHg8HklSZ2dn0PrOzs7Atq9zOp1yuVxBCwAAQFgjJScnRx6PR9XV1YF1Pp9Pu3btktfrDedTAQCAGBfyNSmHDx9WS0tL4PGBAwfU0NCglJQUZWdna8WKFfrtb3+riRMnKicnR/fee68yMzM1f/78cI4bAADEuJAjZffu3briiisCj0tKSiRJixcv1oYNG3TXXXepr69Pt956q7q7u3XZZZdpy5YtGj58ePhGDQAAYl7IkTJz5kwZY75xu8Ph0AMPPKAHHnjgtAYGAACGtqh/ugcAAGAgRAoAALASkQIAAKwU1jvOAgDscSberfrDNddEewiwCGdSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJW44ywAwBpn4l1yETmcSQEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJXCHin333+/HA5H0DJ58uRwPw0AAIhx8ZE46Pnnn6833njj/58kPiJPAwAAYlhE6iE+Pl4ejycShwYAAENERK5J2bdvnzIzM3XOOedo0aJFam1tjcTTAACAGBb2Myl5eXnasGGDJk2apPb2dpWVlenyyy9XU1OTkpKSTtrf7/fL7/cHHvt8vnAPCQAAnIHCHimFhYWBr3Nzc5WXl6exY8fq+eef15IlS07av7y8XGVlZeEeBgAAOMNF/CPIycnJOvfcc9XS0jLg9tLSUvX09ASWtra2SA8JAACcASIeKYcPH9b+/fuVkZEx4Han0ymXyxW0AAAAhD1SfvGLX6i2tlYffvih/vGPf+jaa6/VWWedpYULF4b7qQAAQAwL+zUpH330kRYuXKhDhw5p9OjRuuyyy7Rz506NHj063E8FAABiWNgjpaqqKtyHBAAAQxB/uwcAAFiJSAEAAFYiUgAAgJWIFAAAYKUh9+eJx616LdpDAAAA3wFnUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFYiUgAAgJWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFiJSAEAAFaKWKSsW7dO48aN0/Dhw5WXl6e33347Uk8FAABiUEQi5c9//rNKSkq0evVqvfvuu5o2bZoKCgrU1dUViacDAAAxKCKR8vDDD2vp0qW65ZZbNGXKFD3xxBMaOXKknnrqqUg8HQAAiEHx4T7g0aNHVV9fr9LS0sC6uLg45efnq66u7qT9/X6//H5/4HFPT48kyefzhXtokqR+/2cROS4AAGeKSPyOPXFMY0zYjhn2SPnkk090/PhxpaenB61PT0/XBx98cNL+5eXlKisrO2l9VlZWuIcGAAAkuddG7ti9vb1yu91hOVbYIyVUpaWlKikpCTzu7+/Xp59+qlGjRqm3t1dZWVlqa2uTy+WK4iiHNp/PxzxYgHmwA/NgB+bBDl+dh6SkJPX29iozMzNsxw97pKSmpuqss85SZ2dn0PrOzk55PJ6T9nc6nXI6nUHrkpOTJUkOh0OS5HK5+E9oAebBDsyDHZgHOzAPdjgxD+E6g3JC2C+cTUhI0PTp01VdXR1Y19/fr+rqanm93nA/HQAAiFERebunpKREixcv1oUXXqiLL75Ya9euVV9fn2655ZZIPB0AAIhBEYmUBQsW6OOPP9Z9992njo4O/eAHP9CWLVtOupj2f3E6nVq9evVJbwdhcDEPdmAe7MA82IF5sEOk58FhwvlZIQAAgDDhb/cAAAArESkAAMBKRAoAALASkQIAAKxkbaSsW7dO48aN0/Dhw5WXl6e333472kOKKTt27NDcuXOVmZkph8OhTZs2BW03xui+++5TRkaGRowYofz8fO3bty9on08//VSLFi2Sy+VScnKylixZosOHDw/iqzjzlZeX66KLLlJSUpLS0tI0f/58NTc3B+1z5MgRFRcXa9SoUUpMTFRRUdFJN0tsbW3VNddco5EjRyotLU2//OUv9cUXXwzmSzmjVVRUKDc3N3BDKq/Xq82bNwe2MwfRsWbNGjkcDq1YsSKwjrmIvPvvv18OhyNomTx5cmD7oM6BsVBVVZVJSEgwTz31lHnvvffM0qVLTXJysuns7Iz20GLG66+/bn7961+bl156yUgyGzduDNq+Zs0a43a7zaZNm8w///lP89Of/tTk5OSYzz//PLDPnDlzzLRp08zOnTvN3/72NzNhwgSzcOHCQX4lZ7aCggKzfv1609TUZBoaGszVV19tsrOzzeHDhwP7LFu2zGRlZZnq6mqze/duc8kll5gf//jHge1ffPGFmTp1qsnPzzd79uwxr7/+uklNTTWlpaXReElnpL/85S/mtddeM//+979Nc3Oz+dWvfmWGDRtmmpqajDHMQTS8/fbbZty4cSY3N9fccccdgfXMReStXr3anH/++aa9vT2wfPzxx4HtgzkHVkbKxRdfbIqLiwOPjx8/bjIzM015eXkURxW7vh4p/f39xuPxmIceeiiwrru72zidTvPcc88ZY4x5//33jSTzzjvvBPbZvHmzcTgc5r///e+gjT3WdHV1GUmmtrbWGPPlv/uwYcPMCy+8ENjnX//6l5Fk6urqjDFfBmdcXJzp6OgI7FNRUWFcLpfx+/2D+wJiyNlnn22efPJJ5iAKent7zcSJE822bdvMT37yk0CkMBeDY/Xq1WbatGkDbhvsObDu7Z6jR4+qvr5e+fn5gXVxcXHKz89XXV1dFEc2dBw4cEAdHR1Bc+B2u5WXlxeYg7q6OiUnJ+vCCy8M7JOfn6+4uDjt2rVr0MccK3p6eiRJKSkpkqT6+nodO3YsaC4mT56s7OzsoLm44IILgm6WWFBQIJ/Pp/fee28QRx8bjh8/rqqqKvX19cnr9TIHUVBcXKxrrrkm6N9c4udhMO3bt0+ZmZk655xztGjRIrW2tkoa/DmI+l9B/rpPPvlEx48fP+nutOnp6frggw+iNKqhpaOjQ5IGnIMT2zo6OpSWlha0PT4+XikpKYF9EJr+/n6tWLFCl156qaZOnSrpy3/nhISEwB/dPOHrczHQXJ3Yhu9m79698nq9OnLkiBITE7Vx40ZNmTJFDQ0NzMEgqqqq0rvvvqt33nnnpG38PAyOvLw8bdiwQZMmTVJ7e7vKysp0+eWXq6mpadDnwLpIAYaq4uJiNTU16a233or2UIakSZMmqaGhQT09PXrxxRe1ePFi1dbWRntYQ0pbW5vuuOMObdu2TcOHD4/2cIaswsLCwNe5ubnKy8vT2LFj9fzzz2vEiBGDOhbr3u5JTU3VWWedddKVwp2dnfJ4PFEa1dBy4t/52+bA4/Goq6sraPsXX3yhTz/9lHk6BcuXL9err76qN998U2PGjAms93g8Onr0qLq7u4P2//pcDDRXJ7bhu0lISNCECRM0ffp0lZeXa9q0aXr00UeZg0FUX1+vrq4u/ehHP1J8fLzi4+NVW1urxx57TPHx8UpPT2cuoiA5OVnnnnuuWlpaBv3nwbpISUhI0PTp01VdXR1Y19/fr+rqanm93iiObOjIycmRx+MJmgOfz6ddu3YF5sDr9aq7u1v19fWBfbZv367+/n7l5eUN+pjPVMYYLV++XBs3btT27duVk5MTtH369OkaNmxY0Fw0NzertbU1aC727t0bFI3btm2Ty+XSlClTBueFxKD+/n75/X7mYBDNmjVLe/fuVUNDQ2C58MILtWjRosDXzMXgO3z4sPbv36+MjIzB/3kI+bLfQVBVVWWcTqfZsGGDef/9982tt95qkpOTg64Uxunp7e01e/bsMXv27DGSzMMPP2z27Nlj/vOf/xhjvvwIcnJysnn55ZdNY2OjmTdv3oAfQf7hD39odu3aZd566y0zceJEPoIcottuu8243W5TU1MT9HG/zz77LLDPsmXLTHZ2ttm+fbvZvXu38Xq9xuv1Braf+Ljf7NmzTUNDg9myZYsZPXo0H7kMwapVq0xtba05cOCAaWxsNKtWrTIOh8P89a9/NcYwB9H01U/3GMNcDIY777zT1NTUmAMHDpi///3vJj8/36Smppquri5jzODOgZWRYowxf/jDH0x2drZJSEgwF198sdm5c2e0hxRT3nzzTSPppGXx4sXGmC8/hnzvvfea9PR043Q6zaxZs0xzc3PQMQ4dOmQWLlxoEhMTjcvlMrfccovp7e2Nwqs5cw00B5LM+vXrA/t8/vnn5uc//7k5++yzzciRI821115r2tvbg47z4YcfmsLCQjNixAiTmppq7rzzTnPs2LFBfjVnrp/97Gdm7NixJiEhwYwePdrMmjUrECjGMAfR9PVIYS4ib8GCBSYjI8MkJCSY73//+2bBggWmpaUlsH0w58BhjDGnfA4IAAAgQqy7JgUAAEAiUgAAgKWIFAAAYCUiBQAAWIlIAQAAViJSAACAlYgUAABgJSIFAABYiUgBAABWIlIAAICViBQAAGAlIgUAAFjp/wAIEdyaGacQOgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean duration: 279.02 +/- 119.13\n",
      "max duration: 479.96 s duration: 20.36 s\n",
      "3 / 170 clips are under 60 seconds\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "durations = []\n",
    "under_60s = 0\n",
    "print(len(prompts))\n",
    "for idx, audio_dict in enumerate(prompts):\n",
    "    duration_s = audio_dict[\"audio\"].shape[1] / audio_dict[\"sr\"]\n",
    "    durations.append(duration_s)\n",
    "    if duration_s < 60:\n",
    "        under_60s += 1\n",
    "\n",
    "    # add the duration to the prompts dict\n",
    "    audio_dict[\"duration_s\"] = duration_s\n",
    "        \n",
    "plt.hist(durations, bins=10)\n",
    "plt.show()\n",
    "print(f\"mean duration: {np.mean(durations):.2f} +/- {np.std(durations):.2f}\")\n",
    "print(f\"max duration: {np.max(durations):.2f} s duration: {np.min(durations):.2f} s\")\n",
    "print(f\"{under_60s} / {len(prompts)} clips are under 60 seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vocal detection\n",
    "\n",
    "Use basic stem split model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170/170 [11:35<00:00,  4.09s/it, percent_vocal=54.71%, pesq=1.86, si_sdr=5.97, stoi=0.848]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "54.71% of clips have vocals\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from suno_utils.audio import Audio\n",
    "import tempfile\n",
    "\n",
    "def compute_rms_energy(audio_tensor: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"Compute RMS energy of audio tensor with overlapping frames.\n",
    "    \n",
    "    Args:\n",
    "        audio_tensor (torch.Tensor): Input audio tensor of shape (channels, samples) or (samples,)\n",
    "        \n",
    "    Returns:\n",
    "        torch.Tensor: RMS energy values for each frame\n",
    "    \"\"\"\n",
    "    # Convert to mono if stereo\n",
    "    if len(audio_tensor.shape) > 1 and audio_tensor.shape[0] == 2:\n",
    "        audio_tensor = torch.mean(audio_tensor, dim=0)\n",
    "    elif len(audio_tensor.shape) == 1:\n",
    "        audio_tensor = audio_tensor\n",
    "    else:\n",
    "        raise ValueError(f\"Unexpected audio tensor shape: {audio_tensor.shape}\")\n",
    "        \n",
    "    # Parameters for frame analysis\n",
    "    frame_length = 2048  # ~46ms at 44.1kHz\n",
    "    hop_length = 512    # 75% overlap\n",
    "    \n",
    "    # Pad audio to handle partial frames at end\n",
    "    pad_length = (frame_length - (len(audio_tensor) % frame_length)) % frame_length\n",
    "    audio_tensor = torch.nn.functional.pad(audio_tensor, (0, pad_length))\n",
    "    \n",
    "    # Create overlapping frames\n",
    "    frames = audio_tensor.unfold(0, frame_length, hop_length)\n",
    "    \n",
    "    # Compute RMS energy for each frame\n",
    "    rms = torch.sqrt(torch.mean(frames**2, dim=1))\n",
    "    \n",
    "    return rms\n",
    "\n",
    "vocal_count = []\n",
    "vocal_stoi = []\n",
    "vocal_pesq = []\n",
    "vocal_si_sdr = []\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "for audio_dict in pbar:\n",
    "    with torch.no_grad():\n",
    "        audio_tensor = audio_dict[\"audio\"]\n",
    "        sr = audio_dict[\"sr\"]\n",
    "\n",
    "        # resample to 441000\n",
    "        if sr != 44100:\n",
    "            audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 44100)\n",
    "\n",
    "        audio_tensor = audio_tensor.unsqueeze(0).cuda()\n",
    "        output = stem_model(audio_tensor).squeeze(0)\n",
    "\n",
    "        stems = {}\n",
    "        for stem_idx, stem_name in enumerate(AUDIO_STEMS):\n",
    "            stems[stem_name] = output[stem_idx].cpu()\n",
    "\n",
    "        # compute rms energy for vocals\n",
    "        vocals = stems[\"vocals\"]\n",
    "        rms_energy = compute_rms_energy(vocals)\n",
    "        mean_rms_energy = rms_energy.mean().item()\n",
    "\n",
    "        if mean_rms_energy > 0.01:\n",
    "            vocal_count.append(1)\n",
    "            # if we have vocals, measure the objective metrics\n",
    "            # while we are here, measure the vocal quality\n",
    "            vocals_16k = torchaudio.functional.resample(vocals.mean(dim=0).unsqueeze(0), 44100, 16000)\n",
    "\n",
    "            # write vocal to temp mp3 file\n",
    "            #with tempfile.NamedTemporaryFile(suffix=\".mp3\") as temp_file:\n",
    "            #torchaudio.save(\"audio.mp3\", vocals_16k, 16000)\n",
    "            #result = seg(\"audio.mp3\")\n",
    "            #male_duration_s = 0\n",
    "            #female_duration_s = 0\n",
    "            #for segment_type, start, end in result:\n",
    "            #    if segment_type == \"male\":\n",
    "            #        male_duration_s += end - start\n",
    "            #    elif segment_type == \"female\":\n",
    "            #        female_duration_s += end - start\n",
    "\n",
    "            #audio_dict[\"male_duration_s\"] = male_duration_s\n",
    "            #audio_dict[\"female_duration_s\"] = female_duration_s\n",
    "            #audio_dict[\"predicted_gender\"] = \"male\" if male_duration_s > female_duration_s else \"female\"\n",
    "\n",
    "            # now measure the objective metrics\n",
    "            stoi_hyp, pesq_hyp, si_sdr_hyp = objective_model(vocals_16k.cuda())\n",
    "            vocal_stoi.append(stoi_hyp.cpu().item())\n",
    "            vocal_pesq.append(pesq_hyp.cpu().item())\n",
    "            vocal_si_sdr.append(si_sdr_hyp.cpu().item())\n",
    "\n",
    "            # add the metrics to the prompts dict\n",
    "            audio_dict[\"vocal_stoi\"] = stoi_hyp.cpu().item()\n",
    "            audio_dict[\"vocal_pesq\"] = pesq_hyp.cpu().item()\n",
    "            audio_dict[\"vocal_si_sdr\"] = si_sdr_hyp.cpu().item()\n",
    "            audio_dict[\"has_vocals\"] = True\n",
    "\n",
    "        else:\n",
    "            vocal_count.append(0)\n",
    "            audio_dict[\"has_vocals\"] = False\n",
    "            \n",
    "        percent_vocal = np.mean(vocal_count) * 100\n",
    "        pbar.set_postfix(percent_vocal=f\"{percent_vocal:.2f}%\", stoi=np.mean(vocal_stoi), pesq=np.mean(vocal_pesq), si_sdr=np.mean(vocal_si_sdr))\n",
    "\n",
    "\n",
    "print(f\"{np.mean(vocal_count)*100:.2f}% of clips have vocals\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Audio quality\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170/170 [00:17<00:00,  9.54it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAh8AAAGfCAYAAAD/BbCUAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy80BEi2AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAd/klEQVR4nO3dfXBV9Z348U9ISOJoHgQlgRrAhypsldrFNcaH1rLpMspYHdNpbbuWdti6nUanknW2sNWl1q5h1CmsLdiui7DdkWZLp3bX4qJtutKphWqjzCAKK4oDXZo4dpcEcLggOb8/9me2qajcm+QbLr5eM2emOffccz9fL5B3b869KcmyLAsAgETGjPYAAMC7i/gAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSKsvn4K9+9atx++23D9p3zjnnxNatWyMi4sCBA/FXf/VX0dHREblcLmbPnh3Lly+Purq6o36M/v7+2L17d1RVVUVJSUk+4wEAoyTLsti7d29MmjQpxox5+9c28oqPiIj3ve998dOf/vT/TlD2f6eYP39+rF27NtasWRM1NTVx4403xrXXXhtPPPHEUZ9/9+7d0dDQkO9YAMAxYNeuXXHaaae97TF5x0dZWVnU19e/aX9vb2+sWLEiVq9eHbNmzYqIiJUrV8b06dNj48aNcdFFFx3V+auqqiLif4evrq7OdzwAYBT09fVFQ0PDwPfxt5N3fLzwwgsxadKkqKysjKampmhvb4/JkydHV1dXHDp0KJqbmweOnTZtWkyePDk2bNjwlvGRy+Uil8sNfL13796IiKiurhYfAFBkjuaSibwuOG1sbIxVq1bFunXr4r777osdO3bEZZddFnv37o3u7u4oLy+P2traQfepq6uL7u7utzxne3t71NTUDGx+5AIAx7e8Xvm44oorBv73jBkzorGxMaZMmRLf//7344QTTihogIULF0ZbW9vA12+8bAMAHJ+G9Fbb2traOPvss2P79u1RX18fBw8ejD179gw6pqen54jXiLyhoqJi4EcsftQCAMe/IcXHvn374sUXX4yJEyfGzJkzY+zYsdHZ2Tlw+7Zt22Lnzp3R1NQ05EEBgONDXj92ueWWW+Kqq66KKVOmxO7du2PRokVRWloan/zkJ6OmpibmzZsXbW1tMW7cuKiuro6bbropmpqajvqdLgDA8S+v+PjNb34Tn/zkJ+N3v/tdnHrqqXHppZfGxo0b49RTT42IiCVLlsSYMWOipaVl0IeMAQC8oSTLsmy0h/h9fX19UVNTE729va7/AIAikc/3b7/bBQBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAksrrQ8YYHVMXrB3tEfL28uI5oz0CAMcor3wAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBSPueDEeGzSQB4K175AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJDWk+Fi8eHGUlJTEzTffPLDvwIED0draGuPHj4+TTjopWlpaoqenZ6hzAgDHiYLj46mnnorvfOc7MWPGjEH758+fHw8//HCsWbMm1q9fH7t3745rr712yIMCAMeHguJj37598elPfzruv//+OPnkkwf29/b2xooVK+Ib3/hGzJo1K2bOnBkrV66MX/7yl7Fx48ZhGxoAKF4FxUdra2vMmTMnmpubB+3v6uqKQ4cODdo/bdq0mDx5cmzYsOGI58rlctHX1zdoAwCOX2X53qGjoyOefvrpeOqpp950W3d3d5SXl0dtbe2g/XV1ddHd3X3E87W3t8ftt9+e7xgAQJHK65WPXbt2xZe+9KV48MEHo7KyclgGWLhwYfT29g5su3btGpbzAgDHprzio6urK1555ZX44z/+4ygrK4uysrJYv3593HvvvVFWVhZ1dXVx8ODB2LNnz6D79fT0RH19/RHPWVFREdXV1YM2AOD4ldePXf70T/80Nm/ePGjf5z73uZg2bVp8+ctfjoaGhhg7dmx0dnZGS0tLRERs27Ytdu7cGU1NTcM3NQBQtPKKj6qqqjj33HMH7TvxxBNj/PjxA/vnzZsXbW1tMW7cuKiuro6bbropmpqa4qKLLhq+qQGAopX3BafvZMmSJTFmzJhoaWmJXC4Xs2fPjuXLlw/3wwAARaoky7JstIf4fX19fVFTUxO9vb2u//j/pi5YO9ojvCu8vHjOaI8AULTy+f7td7sAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUnnFx3333RczZsyI6urqqK6ujqampvj3f//3gdsPHDgQra2tMX78+DjppJOipaUlenp6hn1oAKB45RUfp512WixevDi6urri17/+dcyaNSuuvvrq2LJlS0REzJ8/Px5++OFYs2ZNrF+/Pnbv3h3XXnvtiAwOABSnkizLsqGcYNy4cXH33XfHxz72sTj11FNj9erV8bGPfSwiIrZu3RrTp0+PDRs2xEUXXXTE++dyucjlcgNf9/X1RUNDQ/T29kZ1dfVQRjtuTF2wdrRHeFd4efGc0R4BoGj19fVFTU3NUX3/Lviaj8OHD0dHR0fs378/mpqaoqurKw4dOhTNzc0Dx0ybNi0mT54cGzZseMvztLe3R01NzcDW0NBQ6EgAQBHIOz42b94cJ510UlRUVMQXvvCFeOihh+KP/uiPoru7O8rLy6O2tnbQ8XV1ddHd3f2W51u4cGH09vYObLt27cp7EQBA8SjL9w7nnHNObNq0KXp7e+MHP/hBzJ07N9avX1/wABUVFVFRUVHw/QGA4pJ3fJSXl8dZZ50VEREzZ86Mp556Kv7+7/8+PvGJT8TBgwdjz549g1796Onpifr6+mEbGAAobkP+nI/+/v7I5XIxc+bMGDt2bHR2dg7ctm3btti5c2c0NTUN9WEAgONEXq98LFy4MK644oqYPHly7N27N1avXh2PP/54PProo1FTUxPz5s2Ltra2GDduXFRXV8dNN90UTU1Nb/lOFwDg3Sev+HjllVfiM5/5TPz2t7+NmpqamDFjRjz66KPxkY98JCIilixZEmPGjImWlpbI5XIxe/bsWL58+YgMDgAUpyF/zsdwy+d9wu8WPucjDZ/zAVC4JJ/zAQBQCPEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASCrv3+0Cx6ti/DwVn00CFCOvfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSyis+2tvb40/+5E+iqqoqJkyYENdcc01s27Zt0DEHDhyI1tbWGD9+fJx00knR0tISPT09wzo0AFC88oqP9evXR2tra2zcuDF+8pOfxKFDh+LP/uzPYv/+/QPHzJ8/Px5++OFYs2ZNrF+/Pnbv3h3XXnvtsA8OABSnsnwOXrdu3aCvV61aFRMmTIiurq744Ac/GL29vbFixYpYvXp1zJo1KyIiVq5cGdOnT4+NGzfGRRddNHyTAwBFaUjXfPT29kZExLhx4yIioqurKw4dOhTNzc0Dx0ybNi0mT54cGzZsOOI5crlc9PX1DdoAgONXwfHR398fN998c1xyySVx7rnnRkREd3d3lJeXR21t7aBj6+rqoru7+4jnaW9vj5qamoGtoaGh0JEAgCJQcHy0trbGs88+Gx0dHUMaYOHChdHb2zuw7dq1a0jnAwCObXld8/GGG2+8MX784x/Hz3/+8zjttNMG9tfX18fBgwdjz549g1796Onpifr6+iOeq6KiIioqKgoZAwAoQnm98pFlWdx4443x0EMPxc9+9rM4/fTTB90+c+bMGDt2bHR2dg7s27ZtW+zcuTOampqGZ2IAoKjl9cpHa2trrF69Ov71X/81qqqqBq7jqKmpiRNOOCFqampi3rx50dbWFuPGjYvq6uq46aaboqmpyTtdAICIyDM+7rvvvoiIuPzyywftX7lyZXz2s5+NiIglS5bEmDFjoqWlJXK5XMyePTuWL18+LMMCAMUvr/jIsuwdj6msrIxly5bFsmXLCh4KADh+FXTBaTGbumDtaI8AAO9qfrEcAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASKpstAcACjd1wdrRHiFvLy+eM9ojAKPMKx8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkso7Pn7+85/HVVddFZMmTYqSkpL40Y9+NOj2LMvib//2b2PixIlxwgknRHNzc7zwwgvDNS8AUOTyjo/9+/fH+9///li2bNkRb7/rrrvi3nvvjW9/+9vxq1/9Kk488cSYPXt2HDhwYMjDAgDFryzfO1xxxRVxxRVXHPG2LMti6dKlceutt8bVV18dERHf/e53o66uLn70ox/FddddN7RpAYCiN6zXfOzYsSO6u7ujubl5YF9NTU00NjbGhg0bjnifXC4XfX19gzYA4PiV9ysfb6e7uzsiIurq6gbtr6urG7jtD7W3t8ftt98+nGMAx7CpC9aO9gh5e3nxnNEeAY4ro/5ul4ULF0Zvb+/AtmvXrtEeCQAYQcMaH/X19RER0dPTM2h/T0/PwG1/qKKiIqqrqwdtAMDxa1jj4/TTT4/6+vro7Owc2NfX1xe/+tWvoqmpaTgfCgAoUnlf87Fv377Yvn37wNc7duyITZs2xbhx42Ly5Mlx8803x9e//vV473vfG6effnrcdtttMWnSpLjmmmuGc24AoEjlHR+//vWv48Mf/vDA121tbRERMXfu3Fi1alX89V//dezfvz9uuOGG2LNnT1x66aWxbt26qKysHL6pAYCiVZJlWTbaQ/y+vr6+qKmpid7e3hG5/qMYr7QHRpd3u8A7y+f796i/2wUAeHcRHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBISnwAAEmJDwAgKfEBACQlPgCApMQHAJCU+AAAkhIfAEBS4gMASEp8AABJiQ8AICnxAQAkJT4AgKTEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACSEh8AQFLiAwBIqmy0BwA41k1dsHa0R3jXeHnxnNEeIW/F+OdjtP87e+UDAEhKfAAASYkPACCpEYuPZcuWxdSpU6OysjIaGxvjySefHKmHAgCKyIjEx7/8y79EW1tbLFq0KJ5++ul4//vfH7Nnz45XXnllJB4OACgiI/Jul2984xvx+c9/Pj73uc9FRMS3v/3tWLt2bTzwwAOxYMGCQcfmcrnI5XIDX/f29kZERF9f30iMFv2510bkvAAM3Uj92z+SivH7ykj8d37jnFmWvfPB2TDL5XJZaWlp9tBDDw3a/5nPfCb76Ec/+qbjFy1alEWEzWaz2Wy242DbtWvXO7bCsL/y8eqrr8bhw4ejrq5u0P66urrYunXrm45fuHBhtLW1DXzd398f//3f/x3jx4+PkpKSIc/T19cXDQ0NsWvXrqiurh7y+Y411lfcrK94Hc9ri7C+Yjca68uyLPbu3RuTJk16x2NH/UPGKioqoqKiYtC+2traYX+c6urq4/IP2Busr7hZX/E6ntcWYX3FLvX6ampqjuq4Yb/g9JRTTonS0tLo6ekZtL+npyfq6+uH++EAgCIz7PFRXl4eM2fOjM7OzoF9/f390dnZGU1NTcP9cABAkRmRH7u0tbXF3Llz44ILLogLL7wwli5dGvv37x9490tKFRUVsWjRojf9aOd4YX3FzfqK1/G8tgjrK3bH+vpKsuxo3hOTv29961tx9913R3d3d5x//vlx7733RmNj40g8FABQREYsPgAAjsTvdgEAkhIfAEBS4gMASEp8AABJHRfxsWzZspg6dWpUVlZGY2NjPPnkk297/Jo1a2LatGlRWVkZ5513XjzyyCOJJi1MPuvbsmVLtLS0xNSpU6OkpCSWLl2abtAC5bO++++/Py677LI4+eST4+STT47m5uZ3fL5HWz7r++EPfxgXXHBB1NbWxoknnhjnn39+/PM//3PCafOX79+/N3R0dERJSUlcc801IzvgEOSztlWrVkVJScmgrbKyMuG0+cv3uduzZ0+0trbGxIkTo6KiIs4+++xj+t/PfNZ3+eWXv+n5KykpiTlz5iScOD/5Pn9Lly6Nc845J0444YRoaGiI+fPnx4EDBxJN+weG4XfJjaqOjo6svLw8e+CBB7ItW7Zkn//857Pa2tqsp6fniMc/8cQTWWlpaXbXXXdlzz33XHbrrbdmY8eOzTZv3px48qOT7/qefPLJ7JZbbsm+973vZfX19dmSJUvSDpynfNf3qU99Klu2bFn2zDPPZM8//3z22c9+Nqupqcl+85vfJJ786OS7vv/4j//IfvjDH2bPPfdctn379mzp0qVZaWlptm7dusSTH5181/eGHTt2ZO95z3uyyy67LLv66qvTDJunfNe2cuXKrLq6Ovvtb387sHV3dyee+ujlu75cLpddcMEF2ZVXXpn94he/yHbs2JE9/vjj2aZNmxJPfnTyXd/vfve7Qc/ds88+m5WWlmYrV65MO/hRynd9Dz74YFZRUZE9+OCD2Y4dO7JHH300mzhxYjZ//vzEk/+voo+PCy+8MGttbR34+vDhw9mkSZOy9vb2Ix7/8Y9/PJszZ86gfY2Njdlf/uVfjuichcp3fb9vypQpx3x8DGV9WZZlr7/+elZVVZX90z/900iNOCRDXV+WZdkHPvCB7NZbbx2J8YaskPW9/vrr2cUXX5z94z/+YzZ37txjNj7yXdvKlSuzmpqaRNMNXb7ru++++7IzzjgjO3jwYKoRh2Sof/eWLFmSVVVVZfv27RupEYck3/W1trZms2bNGrSvra0tu+SSS0Z0zrdS1D92OXjwYHR1dUVzc/PAvjFjxkRzc3Ns2LDhiPfZsGHDoOMjImbPnv2Wx4+mQtZXTIZjfa+99locOnQoxo0bN1JjFmyo68uyLDo7O2Pbtm3xwQ9+cCRHLUih6/va174WEyZMiHnz5qUYsyCFrm3fvn0xZcqUaGhoiKuvvjq2bNmSYty8FbK+f/u3f4umpqZobW2Nurq6OPfcc+POO++Mw4cPpxr7qA3Hvy0rVqyI6667Lk488cSRGrNghazv4osvjq6uroEfzbz00kvxyCOPxJVXXplk5j806r/VdiheffXVOHz4cNTV1Q3aX1dXF1u3bj3ifbq7u494fHd394jNWahC1ldMhmN9X/7yl2PSpElvCspjQaHr6+3tjfe85z2Ry+WitLQ0li9fHh/5yEdGety8FbK+X/ziF7FixYrYtGlTggkLV8jazjnnnHjggQdixowZ0dvbG/fcc09cfPHFsWXLljjttNNSjH3UClnfSy+9FD/72c/i05/+dDzyyCOxffv2+OIXvxiHDh2KRYsWpRj7qA3135Ynn3wynn322VixYsVIjTgkhazvU5/6VLz66qtx6aWXRpZl8frrr8cXvvCF+Ju/+ZsUI79JUccH726LFy+Ojo6OePzxx4/5C/vyUVVVFZs2bYp9+/ZFZ2dntLW1xRlnnBGXX375aI82JHv37o3rr78+7r///jjllFNGe5xh19TUNOiXZ1588cUxffr0+M53vhN33HHHKE42PPr7+2PChAnxD//wD1FaWhozZ86M//qv/4q77777mIuPoVqxYkWcd955ceGFF472KMPm8ccfjzvvvDOWL18ejY2NsX379vjSl74Ud9xxR9x2223J5ynq+DjllFOitLQ0enp6Bu3v6emJ+vr6I96nvr4+r+NHUyHrKyZDWd8999wTixcvjp/+9KcxY8aMkRyzYIWub8yYMXHWWWdFRMT5558fzz//fLS3tx9z8ZHv+l588cV4+eWX46qrrhrY19/fHxERZWVlsW3btjjzzDNHduijNBx/98aOHRsf+MAHYvv27SMx4pAUsr6JEyfG2LFjo7S0dGDf9OnTo7u7Ow4ePBjl5eUjOnM+hvL87d+/Pzo6OuJrX/vaSI44JIWs77bbbovrr78+/uIv/iIiIs4777zYv39/3HDDDfGVr3wlxoxJexVGUV/zUV5eHjNnzozOzs6Bff39/dHZ2Tno/4H8vqampkHHR0T85Cc/ecvjR1Mh6ysmha7vrrvuijvuuCPWrVsXF1xwQYpRCzJcz19/f3/kcrmRGHFI8l3ftGnTYvPmzbFp06aB7aMf/Wh8+MMfjk2bNkVDQ0PK8d/WcDx3hw8fjs2bN8fEiRNHasyCFbK+Sy65JLZv3z4QjBER//mf/xkTJ048psIjYmjP35o1ayKXy8Wf//mfj/SYBStkfa+99tqbAuONkMxG41e8jcplrsOoo6Mjq6ioyFatWpU999xz2Q033JDV1tYOvMXt+uuvzxYsWDBw/BNPPJGVlZVl99xzT/b8889nixYtOubfapvP+nK5XPbMM89kzzzzTDZx4sTslltuyZ555pnshRdeGK0lvK1817d48eKsvLw8+8EPfjDobXF79+4drSW8rXzXd+edd2aPPfZY9uKLL2bPPfdcds8992RlZWXZ/fffP1pLeFv5ru8PHcvvdsl3bbfffnv26KOPZi+++GLW1dWVXXfddVllZWW2ZcuW0VrC28p3fTt37syqqqqyG2+8Mdu2bVv24x//OJswYUL29a9/fbSW8LYK/bN56aWXZp/4xCdSj5u3fNe3aNGirKqqKvve976XvfTSS9ljjz2WnXnmmdnHP/7xUZm/6OMjy7Lsm9/8ZjZ58uSsvLw8u/DCC7ONGzcO3PahD30omzt37qDjv//972dnn312Vl5enr3vfe/L1q5dm3ji/OSzvh07dmQR8abtQx/6UPrBj1I+65syZcoR17do0aL0gx+lfNb3la98JTvrrLOyysrK7OSTT86ampqyjo6OUZj66OX79+/3HcvxkWX5re3mm28eOLauri678sors6effnoUpj56+T53v/zlL7PGxsasoqIiO+OMM7K/+7u/y15//fXEUx+9fNe3devWLCKyxx57LPGkhclnfYcOHcq++tWvZmeeeWZWWVmZNTQ0ZF/84hez//mf/0k/eJZlJVk2Gq+3AADvVkV9zQcAUHzEBwCQlPgAAJISHwBAUuIDAEhKfAAASYkPACAp8QEAJCU+AICkxAcAkJT4AACS+n+LZNMeKpF4RQAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean stereo width: 0.21 +/- 0.12\n",
      "min stereo width: 0.00 max stereo width: 0.81\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def calculate_stereo_width(waveform):\n",
    "    # Split into left and right channels\n",
    "    left = waveform[0]\n",
    "    right = waveform[1]\n",
    "  \n",
    "    # Compute mid/side representation\n",
    "    mid = (left + right) / 2\n",
    "    side = (left - right) / 2\n",
    "    \n",
    "    # Compute RMS energy of mid and side channels\n",
    "    mid_energy = torch.sqrt(torch.mean(mid ** 2))\n",
    "    side_energy = torch.sqrt(torch.mean(side ** 2))\n",
    "    \n",
    "    # Compute stereo width based on mid/side ratio\n",
    "    # Normalize to range 0-1 using sigmoid-like function\n",
    "    width_ratio = (side_energy / (mid_energy + 1e-8)).item()\n",
    "    stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5)\n",
    "\n",
    "    return stereo_width.item()\n",
    "\n",
    "stereo_widths = []\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "for audio_dict in pbar:\n",
    "    stereo_width = calculate_stereo_width(audio_dict[\"audio\"])\n",
    "    stereo_widths.append(stereo_width)\n",
    "    audio_dict[\"stereo_width\"] = stereo_width\n",
    "\n",
    "# look at the distribution of stereo widths\n",
    "plt.hist(stereo_widths, bins=10)\n",
    "plt.show()\n",
    "\n",
    "# look at the mean and std of the stereo widths\n",
    "print(f\"mean stereo width: {np.mean(stereo_widths):.2f} +/- {np.std(stereo_widths):.2f}\")\n",
    "print(f\"min stereo width: {np.min(stereo_widths):.2f} max stereo width: {np.max(stereo_widths):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Audio start\n",
    "Detect if we have silence at the beginning of the audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170/170 [00:53<00:00,  3.19it/s]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "# we will use the rms energy to detect silence\n",
    "def find_start_of_signal_rms(waveform, sr, window_size=16384, threshold=0.005):\n",
    "    \"\"\"\n",
    "    Finds the time in seconds where the RMS energy of the waveform exceeds the given threshold.\n",
    "    \n",
    "    :param waveform: PyTorch tensor containing the audio waveform.\n",
    "    :param sr: Integer, the sample rate of the audio.\n",
    "    :param window_size: Integer, the number of samples over which to calculate the RMS.\n",
    "    :param threshold: Float, the RMS value above which the signal is considered to have started.\n",
    "    :return: float, the time in seconds where RMS exceeds the threshold, or -1 if not found.\n",
    "    \"\"\"\n",
    "    # Calculate the squared values\n",
    "    squared_waveform = waveform**2\n",
    "    \n",
    "    # Compute rolling window RMS\n",
    "    rms = torch.sqrt(squared_waveform.unfold(0, window_size, 1).mean(dim=1))\n",
    "    # Find the first index where the RMS exceeds the threshold\n",
    "    start_index = torch.nonzero(rms > threshold, as_tuple=True)[0]\n",
    "    \n",
    "    if len(start_index) > 0:\n",
    "        # Convert sample index to seconds\n",
    "        # Divide by hop size (1) instead of window_size to get correct time\n",
    "        start_time = start_index[0].item() / sr\n",
    "        return start_time\n",
    "    else:\n",
    "        return -1  # Return -1 if no RMS exceeds the threshold\n",
    "\n",
    "start_times = []\n",
    "\n",
    "pbar = tqdm(prompts)\n",
    "for audio_dict in pbar:\n",
    "    start_time = find_start_of_signal_rms(audio_dict[\"audio\"].mean(dim=0), audio_dict[\"sr\"])\n",
    "    start_times.append(start_time)\n",
    "    audio_dict[\"start_time\"] = start_time\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start time histogram\n",
    "plt.hist([prompt[\"start_time\"] for prompt in prompts], bins=25)\n",
    "plt.show()\n",
    "\n",
    "print(f\"mean start time: {np.mean(start_times):.2f} +/- {np.std(start_times):.2f}\")\n",
    "print(f\"min start time: {np.min(start_times):.2f} max start time: {np.max(start_times):.2f}\")\n",
    "\n",
    "# find the audio_filepath for the prompt with the longest start time\n",
    "# sort the prompts by start time\n",
    "prompts_by_start_time = sorted(prompts, key=lambda x: x[\"start_time\"], reverse=True)\n",
    "\n",
    "# print the audio_filepath for the prompt with the longest start time\n",
    "print(prompts_by_start_time[0][\"audio_filepath\"], prompts_by_start_time[0][\"start_time\"])\n",
    "print(prompts_by_start_time[1][\"audio_filepath\"], prompts_by_start_time[1][\"start_time\"])\n",
    "print(prompts_by_start_time[2][\"audio_filepath\"], prompts_by_start_time[2][\"start_time\"])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts_with_audio = []\n",
    "\n",
    "for prompt in prompts:\n",
    "    new_prompt = prompt.copy()\n",
    "    new_prompt.pop(\"audio\")\n",
    "    prompts_with_audio.append(new_prompt)\n",
    "\n",
    "# add the ckpt_name to the prompts_with_audio\n",
    "for prompt in prompts_with_audio:\n",
    "    prompt[\"ckpt_name\"] = ckpt_name\n",
    "\n",
    "# construct the dataframe\n",
    "df = pd.DataFrame(prompts_with_audio)\n",
    "df.head()\n",
    "\n",
    "# save dataframe\n",
    "output_filename = f\"suno_bench_prompts_v1_{ckpt_name}.csv\"\n",
    "output_filepath = os.path.join(\"/home/christian/code/christian/sunoBench/results\", output_filename)\n",
    "df.to_csv(output_filepath)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/788da196-93d2-4fdf-93a1-d0e72f36ca46.mp3 3.3449747562408447\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/22b9ddf6-a3b1-4160-bf73-686923f2ec6b.mp3 3.1060805320739746\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/94978b72-c4c5-4eb2-958d-97ce276ab6c7.mp3 3.0801260471343994\n",
      "----------------------------------------------------------------------------------------------------\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/f4d42964-27bd-4665-aa07-1941cd6536a1.mp3 1.1629421710968018\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/7206ef30-754b-4da9-b323-9a4574c15b76.mp3 1.1632678508758545\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/32f1776b-c3b9-4b46-ad93-bd97b5189924.mp3 1.1645466089248657\n"
     ]
    }
   ],
   "source": [
    "# print the top highest vocal_pesq\n",
    "# sort by vocal_pesq\n",
    "df_sorted = df.sort_values(by=\"vocal_pesq\", ascending=False)\n",
    "for idx, row in list(df_sorted.iterrows())[:3]:\n",
    "    print(row[\"audio_filepath\"], row[\"vocal_pesq\"])\n",
    "\n",
    "print(\"-\"*100)\n",
    "# print the top highest clap similarity\n",
    "df_sorted = df.sort_values(by=\"vocal_pesq\", ascending=True)\n",
    "for idx, row in list(df_sorted.iterrows())[:3]:\n",
    "    print(row[\"audio_filepath\"], row[\"vocal_pesq\"])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/46573434-52fa-4a13-ae49-d6a7531db160.mp3\n",
      "/home/christian/code/christian/sunoBench/outputs/2025-02-04_21-04-31-last_ckpt_infer/eab10a44-db3a-49dc-b92d-1e0d45c1f5b5.mp3\n",
      "Number of times [Instrumental] has vocals: 2/22 (9.09%)\n"
     ]
    }
   ],
   "source": [
    "# if the text is \"[Instrumental]\" and has_vocals is True, print the audio_filepath\n",
    "# count the number of times this happens\n",
    "unwanted_vocal_count = 0\n",
    "instrumental_count = 0\n",
    "\n",
    "for idx, row in df.iterrows():\n",
    "    if row[\"text\"] == \"[Instrumental]\":\n",
    "        instrumental_count += 1\n",
    "        if row[\"has_vocals\"]:\n",
    "            print(row[\"audio_filepath\"])\n",
    "            unwanted_vocal_count += 1\n",
    "\n",
    "unwanted_vocal_percentage = (unwanted_vocal_count / instrumental_count)\n",
    "print(f\"Number of times [Instrumental] has vocals: {unwanted_vocal_count}/{instrumental_count} ({unwanted_vocal_percentage:.2%})\")\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean has_vocals for male vocalist: 80.00%\n",
      "Mean ditto_v2_male_vocalist_similarity for male vocalist: 0.02\n",
      "Mean has_vocals for female vocalist: 80.00%\n",
      "Mean ditto_v2_female_vocalist_similarity for female vocalist: 0.08\n"
     ]
    }
   ],
   "source": [
    "# number of clips with no vocals when the tag is \"male vocalist\"\n",
    "\n",
    "# check the has_vocals column for these clips and take the mean\n",
    "mean_has_vocals = df[df[\"tag\"] == \"male vocalist\"][\"has_vocals\"].mean()\n",
    "# also check the ditto_v2_male_vocalist_similarity column for these clips\n",
    "mean_ditto_v2_male_vocalist_similarity = df[df[\"tag\"] == \"male vocalist\"][\"ditto_v2_male_vocalist_similarity\"].mean()\n",
    "print(f\"Mean has_vocals for male vocalist: {mean_has_vocals:.2%}\")\n",
    "print(f\"Mean ditto_v2_male_vocalist_similarity for male vocalist: {mean_ditto_v2_male_vocalist_similarity:.2f}\")\n",
    "\n",
    "# number of clips with no vocals when the tag is \"female vocalist\"\n",
    "# check the has_vocals column for these clips\n",
    "mean_has_vocals = df[df[\"tag\"] == \"female vocalist\"][\"has_vocals\"].mean()\n",
    "print(f\"Mean has_vocals for female vocalist: {mean_has_vocals:.2%}\")\n",
    "# also check the ditto_v2_female_vocalist_similarity column for these clips\n",
    "mean_ditto_v2_female_vocalist_similarity = df[df[\"tag\"] == \"female vocalist\"][\"ditto_v2_female_vocalist_similarity\"].mean()\n",
    "print(f\"Mean ditto_v2_female_vocalist_similarity for female vocalist: {mean_ditto_v2_female_vocalist_similarity:.2f}\")\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we want to load two results csv files and prepare a report\n",
    "df_a = pd.read_csv(\"/home/christian/code/christian/sunoBench/results/suno_bench_prompts_v1_2025-02-04_21-04-31-last_ckpt_infer.csv\")\n",
    "df_b = pd.read_csv(\"/home/christian/code/christian/sunoBench/results/suno_bench_prompts_v1_2025-02-10_16-52-41-step_9000_infer.csv\")\n",
    "\n",
    "# merge the two dataframes\n",
    "df_merged = pd.concat([df_a, df_b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find any values of start_time that are NaN\n",
    "nan_start_time = df_merged[df_b[\"start_time\"].isna()]\n",
    "print(f\"Number of clips with NaN start time: {len(nan_start_time)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean ditto_v2_male_vocalist_ratio: 0.035\n",
      "Mean ditto_v2_female_vocalist_ratio: -0.035\n"
     ]
    }
   ],
   "source": [
    "# we want to compute, for cases where the tag is \"male vocalist\" \n",
    "# for each row, check how much larger the ditto_v2_male_vocalist_similarity is compared to the ditto_v2_similarity\n",
    "df_merged[\"ditto_v2_male_vocalist_ratio\"] = df_merged[\"ditto_v2_male_vocalist_similarity\"] - df_merged[\"ditto_v2_female_vocalist_similarity\"]\n",
    "# print the mean of the ditto_v2_male_vocalist_ratio\n",
    "print(f\"Mean ditto_v2_male_vocalist_ratio: {df_merged['ditto_v2_male_vocalist_ratio'].mean():.3f}\")\n",
    "\n",
    "# do the same for \"female vocalist\"\n",
    "df_merged[\"ditto_v2_female_vocalist_ratio\"] = df_merged[\"ditto_v2_female_vocalist_similarity\"] - df_merged[\"ditto_v2_male_vocalist_similarity\"]\n",
    "# print the mean of the ditto_v2_female_vocalist_ratio\n",
    "print(f\"Mean ditto_v2_female_vocalist_ratio: {df_merged['ditto_v2_female_vocalist_ratio'].mean():.3f}\")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['Unnamed: 0', 'prompt_id', 'audio_filepath', 'sr', 'tag', 'text',\n",
       "       'reciprocal_rank', 'pop_reciprocal_rank', 'ditto_v2_similarity',\n",
       "       'ditto_v2_male_vocalist_similarity',\n",
       "       'ditto_v2_female_vocalist_similarity',\n",
       "       'ditto_v2_instrumental_similarity', 'ditto_v2_pop_similarity',\n",
       "       'duration_s', 'vocal_stoi', 'vocal_pesq', 'vocal_si_sdr', 'has_vocals',\n",
       "       'stereo_width', 'start_time', 'ckpt_name',\n",
       "       'ditto_v2_male_vocalist_ratio', 'ditto_v2_female_vocalist_ratio'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_merged.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metric              2025-02-04_21-04-31-last_ckpt_infer               2025-02-10_16-52-41-step_9000_infer               \n",
      "------------------------------------------------------------------------------------------------------------------------\n",
      "Reciprocal Rank     0.427 ± 0.425                                     0.438 ± 0.438                                     \n",
      "Ditto Similarity    0.245 ± 0.140                                     0.235 ± 0.135                                     \n",
      "Stereo Width        0.145 ± 0.100                                     0.117 ± 0.100                                     \n",
      "Start Time (s)      0.719 ± 2.445                                     0.563 ± 1.931                                     \n",
      "Vocal STOI          0.896 ± 0.081                                     0.906 ± 0.078                                     \n",
      "Vocal PESQ          2.155 ± 0.574                                     2.213 ± 0.569                                     \n",
      "Vocal SI-SDR        8.93 ± 7.20                                       9.97 ± 7.20                                       \n",
      "Duration (s)        217.4 ± 114.4                                     205.7 ± 91.3                                      \n",
      "pop_reciprocal_rank 0.158 ± 0.161                                     0.141 ± 0.098                                     \n",
      "Min Duration (s)    2.9s                                              25.3s                                              \n",
      "# Clips <30s        4                                                 1                                                 \n",
      "Max Duration (s)    480.0s                                            480.0s                                            \n",
      "# Mono Clips        27                                                 34                                                 \n",
      "Instr. Fail %       13.6% ± 0.0%                                      9.1% ± 0.0%                                      \n",
      "Vocal Fail %        24.3% ± 0.0%                                      25.0% ± 0.0%                                      \n"
     ]
    }
   ],
   "source": [
    "# Create a dictionary to store metrics for each model\n",
    "metrics = {\n",
    "    'Reciprocal Rank': ('reciprocal_rank', '{:.3f} ± {:.3f}'),\n",
    "    'Ditto Similarity': ('ditto_v2_similarity', '{:.3f} ± {:.3f}'),\n",
    "    'Stereo Width': ('stereo_width', '{:.3f} ± {:.3f}'),\n",
    "    'Start Time (s)': ('start_time', '{:.3f} ± {:.3f}'),\n",
    "    'Vocal STOI': ('vocal_stoi', '{:.3f} ± {:.3f}'),\n",
    "    'Vocal PESQ': ('vocal_pesq', '{:.3f} ± {:.3f}'),\n",
    "    'Vocal SI-SDR': ('vocal_si_sdr', '{:.2f} ± {:.2f}'),\n",
    "    'Duration (s)': ('duration_s', '{:.1f} ± {:.1f}'),\n",
    "    'pop_reciprocal_rank': ('pop_reciprocal_rank', '{:.3f} ± {:.3f}')\n",
    "}\n",
    "\n",
    "# Get model names\n",
    "model_names = df_merged['ckpt_name'].unique()\n",
    "\n",
    "# Print header\n",
    "print(f\"{'Metric':<20}\", end='')\n",
    "for model in model_names:\n",
    "    print(f\"{model:<50}\", end='')\n",
    "print(\"\\n\" + \"-\"*120)\n",
    "\n",
    "# Print each metric row\n",
    "for metric_name, (col_name, format_str) in metrics.items():\n",
    "    # Drop NaN values before computing stats\n",
    "    stats = df_merged.dropna(subset=[col_name]).groupby(\"ckpt_name\")[col_name].agg([\"mean\", \"std\"])\n",
    "    print(f\"{metric_name:<20}\", end='')\n",
    "    for model in model_names:\n",
    "        if model in stats.index:\n",
    "            formatted_stat = format_str.format(stats.loc[model]['mean'], stats.loc[model]['std'])\n",
    "        else:\n",
    "            formatted_stat = \"N/A\"\n",
    "        print(f\"{formatted_stat:<50}\", end='')\n",
    "    print()\n",
    "\n",
    "# Add shortest duration\n",
    "print(f\"{'Min Duration (s)':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    min_duration = model_df['duration_s'].min()\n",
    "    print(f\"{min_duration:.1f}s\" + \" \"*46, end='')\n",
    "print()\n",
    "\n",
    "# add number of clips under 30 seconds\n",
    "print(f\"{'# Clips <30s':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    under_30 = model_df[model_df['duration_s'] < 30].shape[0]\n",
    "    print(f\"{under_30}\" + \" \"*49, end='')\n",
    "print()\n",
    "\n",
    "# Add longest duration\n",
    "print(f\"{'Max Duration (s)':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    max_duration = model_df['duration_s'].max()\n",
    "    print(f\"{max_duration:.1f}s\" + \" \"*44, end='')\n",
    "print()\n",
    "\n",
    "# count number of mono clips, if stereo width is less than 0.05\n",
    "mono_count = df_merged[df_merged['stereo_width'] < 0.05].shape[0]\n",
    "print(f\"{'# Mono Clips':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    mono_count = model_df[model_df['stereo_width'] < 0.05].shape[0]\n",
    "    print(f\"{mono_count}\" + \" \"*49, end='')\n",
    "print()\n",
    "\n",
    "# Add instrumental failure rate\n",
    "print(f\"{'Instr. Fail %':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    instrumental_mask = model_df['text'] == '[Instrumental]'\n",
    "    total_instrumental = instrumental_mask.sum()\n",
    "    failed = model_df[instrumental_mask & model_df['has_vocals']].shape[0]\n",
    "    fail_rate = failed / total_instrumental if total_instrumental > 0 else 0\n",
    "    print(f\"{fail_rate:.1%} ± 0.0%\"+\" \"*38, end='')\n",
    "print()\n",
    "\n",
    "# Add vocal failure rate\n",
    "print(f\"{'Vocal Fail %':<20}\", end='')\n",
    "for model in model_names:\n",
    "    model_df = df_merged[df_merged['ckpt_name'] == model]\n",
    "    vocal_mask = model_df['text'] != '[Instrumental]'\n",
    "    total_vocal = vocal_mask.sum()\n",
    "    failed = model_df[vocal_mask & ~model_df['has_vocals']].shape[0]\n",
    "    fail_rate = failed / total_vocal if total_vocal > 0 else 0\n",
    "    print(f\"{fail_rate:.1%} ± 0.0%\"+\" \"*38, end='')\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test for bpm\n",
    "# print the filepath for the prompt with the lowest stereo width\n",
    "df_sorted = df_merged.sort_values(by=\"stereo_width\", ascending=True)\n",
    "print(df_sorted[\"audio_filepath\"].iloc[0], df_sorted[\"stereo_width\"].iloc[0])\n",
    "\n",
    "# print the filepath for the prompt with the highest stereo width\n",
    "df_sorted = df_merged.sort_values(by=\"stereo_width\", ascending=False)\n",
    "print(df_sorted[\"audio_filepath\"].iloc[0], df_sorted[\"stereo_width\"].iloc[0])\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
}
