{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import os\n",
    "import tempfile\n",
    "import time\n",
    "import random\n",
    "\n",
    "import funcy\n",
    "import numpy as np\n",
    "import tqdm\n",
    "import torch\n",
    "\n",
    "from suno_utils.tasks.data_loader import load_audio_mp\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    SAMPLE_RATE,\n",
    "    EMBEDDING_RATE,\n",
    "    encode,\n",
    "    encode_files,\n",
    "    preload_models,\n",
    ")\n",
    "# _ = preload_models(checkpoint_filepath=\"/home/tony/Data/MERT/mert_test_8x_400k.pt\")\n",
    "\n",
    "from suno_utils.utils.s3 import read_from_s3, download_s3_files, upload_s3_files, check_s3_file_exists\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl, get_file_ext\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"/home/minz/glockenspiel/musicfm-training/\")\n",
    "from musicfm.models.musicfm_mertlong import MusicFM_MERTLong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_read_jsonl = funcy.partial(read_jsonl, allowed_keys=[\"s3_filepath\"])\n",
    "metas = read_from_s3(\"s3://suno-data/datasets/bundles/v2/music_sample/metas.jsonl\", read_f=p_read_jsonl)\n",
    "random.seed(6006)\n",
    "random.shuffle(metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load first minute of each file\n",
    "def _load_audio(filepaths, sample_rate, max_duration_per_file_s=None):\n",
    "    with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "        t0 = time.time()\n",
    "        if filepaths[0][:2] == \"s3\":\n",
    "            # if filepaths are on s3 then load them to a temp dir first\n",
    "            tmp_out_filepaths = [\n",
    "                os.path.join(tmp_dir, f\"audio_{n}.{get_file_ext(filepath)}\")\n",
    "                for n, filepath in enumerate(filepaths)\n",
    "            ]\n",
    "            print(\"  downloading audio...\")\n",
    "            confirmed_downloads = download_s3_files(\n",
    "                filepaths,\n",
    "                tmp_out_filepaths,\n",
    "                chunksize=100,\n",
    "                n_cores=16,\n",
    "                joblib_backend=\"threads\",\n",
    "                silent=True,\n",
    "            )\n",
    "            time.sleep(5) # make sure things close\n",
    "            local_filepath = [\n",
    "                filepath if b_confirmed else None\n",
    "                for b_confirmed, filepath in zip(confirmed_downloads, tmp_out_filepaths)\n",
    "            ]\n",
    "        else:\n",
    "            local_filepath = filepaths\n",
    "        download_duration_s = round(time.time() - t0, 1)\n",
    "        # remove Nones\n",
    "        safe_orig_idx, safe_filepaths = zip(*[\n",
    "            (idx, fp) for idx, fp in enumerate(local_filepath) if fp is not None\n",
    "        ])\n",
    "        print(\"  loading audio...\")\n",
    "        t0 = time.time()\n",
    "        audio_arrays = load_audio_mp(\n",
    "            safe_filepaths,\n",
    "            target_sample_rate=sample_rate,\n",
    "            max_duration_s=max_duration_per_file_s,\n",
    "            num_workers=32,\n",
    "            force_threads=True,\n",
    "        )\n",
    "        load_duration_s = round(time.time() - t0, 1)\n",
    "        # merge back into None list\n",
    "        out_audio_arrays = [None]*len(filepaths)\n",
    "        for idx, arr in zip(safe_orig_idx, audio_arrays):\n",
    "            out_audio_arrays[idx] = arr\n",
    "        assert(len(filepaths) == len(out_audio_arrays))\n",
    "        if len(filepaths) >= 10:\n",
    "            assert(np.mean([arr is not None for arr in out_audio_arrays]) >= 0.5)\n",
    "    return download_duration_s, load_duration_s, out_audio_arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/minz/anaconda3/envs/suno_env/lib/python3.10/site-packages/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations\n",
      "  warnings.warn(\n",
      "  0%|          | 0/20 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2893/2893 [03:50<00:00, 12.55it/s]\n",
      "  5%|▌         | 1/20 [04:27<1:24:50, 267.90s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1/20: 28.3s downloading, 4.9s loading, 233.5s encoding -- 25.709999999999997h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2940/2940 [03:57<00:00, 12.37it/s]\n",
      " 10%|█         | 2/20 [09:02<1:21:32, 271.81s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 2/20: 27.6s downloading, 4.8s loading, 240.8s encoding -- 26.12777777777778h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3090/3090 [04:10<00:00, 12.36it/s]\n",
      " 15%|█▌        | 3/20 [13:50<1:19:08, 279.34s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 3/20: 28.3s downloading, 5.5s loading, 253.1s encoding -- 27.462222222222223h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3033/3033 [04:06<00:00, 12.28it/s]\n",
      " 20%|██        | 4/20 [18:39<1:15:29, 283.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 4/20: 31.8s downloading, 5.2s loading, 249.9s encoding -- 26.958333333333332h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3123/3123 [04:11<00:00, 12.40it/s]\n",
      " 25%|██▌       | 5/20 [24:29<1:16:47, 307.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 5/20: 87.9s downloading, 5.1s loading, 255.3s encoding -- 27.75722222222222h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2928/2928 [03:55<00:00, 12.42it/s]\n",
      " 30%|███       | 6/20 [29:04<1:09:09, 296.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 6/20: 30.4s downloading, 4.9s loading, 238.7s encoding -- 26.02611111111111h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3079/3079 [04:08<00:00, 12.39it/s]\n",
      " 35%|███▌      | 7/20 [33:53<1:03:41, 293.93s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 7/20: 30.4s downloading, 5.0s loading, 252.0s encoding -- 27.363333333333333h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2952/2952 [04:00<00:00, 12.30it/s]\n",
      " 40%|████      | 8/20 [38:31<57:45, 288.76s/it]  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 8/20: 28.3s downloading, 5.0s loading, 243.0s encoding -- 26.238333333333333h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2851/2851 [03:51<00:00, 12.34it/s]\n",
      " 45%|████▌     | 9/20 [43:01<51:52, 282.99s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 9/20: 30.2s downloading, 4.6s loading, 234.1s encoding -- 25.341666666666665h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2962/2962 [03:59<00:00, 12.38it/s]\n",
      " 50%|█████     | 10/20 [47:41<47:00, 282.07s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 10/20: 30.9s downloading, 5.0s loading, 242.4s encoding -- 26.328333333333333h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2886/2886 [03:54<00:00, 12.32it/s]\n",
      " 55%|█████▌    | 11/20 [52:16<41:59, 279.90s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 11/20: 31.4s downloading, 4.7s loading, 237.5s encoding -- 25.652222222222225h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3054/3054 [04:07<00:00, 12.36it/s]\n",
      " 60%|██████    | 12/20 [57:00<37:28, 281.10s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 12/20: 27.2s downloading, 5.1s loading, 250.1s encoding -- 27.138333333333332h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2939/2939 [03:58<00:00, 12.33it/s]\n",
      " 65%|██████▌   | 13/20 [1:01:37<32:39, 279.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 13/20: 28.7s downloading, 4.9s loading, 241.6s encoding -- 26.118333333333332h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2957/2957 [03:57<00:00, 12.44it/s]\n",
      " 70%|███████   | 14/20 [1:06:15<27:56, 279.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 14/20: 28.0s downloading, 8.2s loading, 240.7s encoding -- 26.284444444444443h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3111/3111 [04:14<00:00, 12.24it/s]\n",
      " 75%|███████▌  | 15/20 [1:11:10<23:39, 283.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 15/20: 29.1s downloading, 6.7s loading, 257.3s encoding -- 27.647222222222222h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2990/2990 [04:01<00:00, 12.39it/s]\n",
      " 80%|████████  | 16/20 [1:15:51<18:52, 283.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 16/20: 30.1s downloading, 5.3s loading, 244.7s encoding -- 26.571666666666665h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2979/2979 [04:02<00:00, 12.29it/s]\n",
      " 85%|████████▌ | 17/20 [1:20:33<14:07, 282.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 17/20: 29.5s downloading, 5.0s loading, 245.3s encoding -- 26.47388888888889h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3127/3127 [04:11<00:00, 12.42it/s]\n",
      " 90%|█████████ | 18/20 [1:25:23<09:30, 285.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 18/20: 29.1s downloading, 5.1s loading, 255.0s encoding -- 27.79111111111111h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2926/2926 [03:54<00:00, 12.50it/s]\n",
      " 95%|█████████▌| 19/20 [1:29:58<04:41, 281.96s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 19/20: 30.9s downloading, 5.1s loading, 237.2s encoding -- 26.006666666666668h processed\n",
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n",
      "  embedding audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3053/3053 [04:07<00:00, 12.31it/s]\n",
      "100%|██████████| 20/20 [1:34:47<00:00, 284.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 20/20: 30.5s downloading, 5.0s loading, 251.0s encoding -- 27.134999999999998h processed\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "chunksize = 250\n",
    "tot_steps = 20\n",
    "train_data = []\n",
    "val_data = []\n",
    "musicfm = MusicFM_MERTLong(model_path=\"/app/suno/minz/models/musicfm_concat_epoch=51.pt\")\n",
    "musicfm = musicfm.cuda()\n",
    "musicfm = musicfm.eval()\n",
    "for n in tqdm.tqdm(range(tot_steps)):\n",
    "    filepaths = [\n",
    "        fp\n",
    "        for m in metas[-(n + 1) * chunksize : len(metas) - n * chunksize]\n",
    "        if (fp := m.get(\"s3_filepath\", m.get(\"audio_filepath\", m.get(\"filepath\"))))\n",
    "        is not None\n",
    "    ]\n",
    "    download_duration_s, load_duration_s, audio_arrays = _load_audio(\n",
    "        filepaths,\n",
    "        SAMPLE_RATE,\n",
    "        max_duration_per_file_s=8 * 60,\n",
    "    )\n",
    "    t0 = time.time()\n",
    "    print(\"  embedding audio...\")\n",
    "    stacked_arr = musicfm.encode_arrays(audio_arrays)\n",
    "    timing_encode_s = round(time.time() - t0, 1)\n",
    "    n_hours_processed = len(stacked_arr) / EMBEDDING_RATE / 60 / 60\n",
    "    # subsample audio array for better diversity\n",
    "    idx_list = list(range(stacked_arr.shape[0]))\n",
    "    random.shuffle(idx_list)\n",
    "    keep_idx = np.array(idx_list[: int(stacked_arr.shape[0] / tot_steps)])\n",
    "    stacked_arr = stacked_arr[keep_idx, :]\n",
    "    if n == tot_steps - 1:\n",
    "        val_data.append(stacked_arr)\n",
    "    else:\n",
    "        train_data.append(stacked_arr)\n",
    "    print(\n",
    "        f\" {n+1}/{tot_steps}: {download_duration_s}s downloading, {load_duration_s}s loading,\"\n",
    "        f\" {timing_encode_s}s encoding -- {n_hours_processed}h processed\"\n",
    "    )\n",
    "#  1/20: 25.0s downloading, 4.5s loading, 61.6s encoding -- 13.0h processed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"/home/minz/musicfm_concat_2s_val\", np.concatenate(val_data, axis=0).astype(np.float32))\n",
    "np.save(\"/home/minz/musicfm_concat_2s_tr\", np.concatenate(train_data, axis=0).astype(np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2272482, 1024)\n",
      "(24422, 1024)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "\n",
    "import numpy as np\n",
    "import faiss\n",
    "from sklearn.metrics.pairwise import paired_distances\n",
    "\n",
    "X_arr = np.load(\"/home/minz/musicfm_concat_2s_tr.npy\")\n",
    "y_arr = np.load(\"/home/minz/musicfm_concat_2s_val.npy\")[::5]\n",
    "\n",
    "print(X_arr.shape)\n",
    "print(y_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_mean = X_arr.mean(axis=0)\n",
    "y_mean = y_arr.mean(axis=0)\n",
    "x_std = X_arr.std(axis=0)\n",
    "y_std = y_arr.std(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calculating codebook 0...\n",
      "train codebook 0...\n",
      "Sampling a subset of 2097152 / 2272482 for training\n",
      "Clustering 2097152 points in 1024D to 8192 clusters, redo 1 times, 25 iterations\n",
      "  Preprocessing in 4.70 s\n",
      "  Iteration 24 (37750.63 s, search 37747.64 s): objective=1.42267e+09 imbalance=1.042 nsplit=0       \n",
      "finish train codebook 0...\n",
      " score: 26.231\n",
      "----------\n"
     ]
    }
   ],
   "source": [
    "n_codebooks = 1\n",
    "n_clusters = 8_192\n",
    "\n",
    "X_arr_resid = X_arr.copy()\n",
    "y_arr_resid = y_arr.copy()\n",
    "y_preds_prev = np.zeros(y_arr.shape)\n",
    "centroids_list = []\n",
    "models_list = []\n",
    "for n_codebook in range(n_codebooks):\n",
    "    print(f\"calculating codebook {n_codebook}...\")\n",
    "    faiss_model = faiss.Kmeans(\n",
    "        d=X_arr_resid.shape[1], \n",
    "        k=n_clusters, \n",
    "        niter=25, \n",
    "        nredo=1, \n",
    "        seed=n_codebook,\n",
    "        verbose=True, \n",
    "        gpu=True\n",
    "    )\n",
    "    print(f\"train codebook {n_codebook}...\")\n",
    "    faiss_model.train(X_arr_resid)\n",
    "    print(f\"finish train codebook {n_codebook}...\")\n",
    "    # score preds\n",
    "    y_cluster_preds = faiss_model.index.search(y_arr_resid, 1)[1].squeeze()\n",
    "    y_preds = faiss_model.centroids[y_cluster_preds]\n",
    "    y_arr_resid -= y_preds\n",
    "    y_preds_prev += y_preds\n",
    "    print(\" score:\", round(np.mean(paired_distances(y_arr, y_preds_prev)), 3))\n",
    "    # start stuff for next round\n",
    "    X_cluster_preds = faiss_model.index.search(X_arr_resid, 1)[1].squeeze()\n",
    "    X_preds = faiss_model.centroids[X_cluster_preds]\n",
    "    X_arr_resid -= X_preds\n",
    "    centroids_list.append(faiss_model.centroids)\n",
    "    models_list.append(faiss_model)\n",
    "    print(\"-\"*10)\n",
    "codebooked_centroids = np.stack(centroids_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"/home/minz/musicfm_concat_centroids_2s_8192\", codebooked_centroids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "centroids = np.load(\"/home/minz/musicfm_concat_centroids_2s_8192.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 8192, 1024)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "centroids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ClusterModel()"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# token comparison\n",
    "\n",
    "chunksize = 250\n",
    "tot_steps = 20\n",
    "train_data = []\n",
    "val_data = []\n",
    "musicfm = MusicFM_MERTLong(model_path=\"/app/suno/minz/models/musicfm_concat_epoch=51.pt\")\n",
    "musicfm = musicfm.cuda()\n",
    "musicfm = musicfm.eval()\n",
    "\n",
    "from suno_utils.tasks.musicfm_v2 import ClusterModel\n",
    "centroids = np.load(\"/home/minz/musicfm_concat_centroids_8192.npy\")\n",
    "cluster_model = ClusterModel(centroids)\n",
    "cluster_model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  downloading audio...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "audio needs conversion, will be slow without using joblib\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  loading audio...\n"
     ]
    }
   ],
   "source": [
    "n = 10\n",
    "filepaths = [\n",
    "    fp\n",
    "    for m in metas[-(n + 1) * chunksize : len(metas) - n * chunksize]\n",
    "    if (fp := m.get(\"s3_filepath\", m.get(\"audio_filepath\", m.get(\"filepath\"))))\n",
    "    is not None\n",
    "]\n",
    "download_duration_s, load_duration_s, audio_arrays = _load_audio(\n",
    "    filepaths,\n",
    "    SAMPLE_RATE,\n",
    "    max_duration_per_file_s=8 * 60,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_30s = audio_arrays[0][:, :24000*30]\n",
    "emb_30s = musicfm.get_latent(audio_30s.cuda()).cpu().detach()\n",
    "out_states_30s = cluster_model.encode(emb_30s)[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_states_5s = []\n",
    "for i in range(6):\n",
    "    audio_5s = audio_arrays[0][:, 24000*i*5:24000*(i+1)*5]\n",
    "    emb_5s = musicfm.get_latent(audio_5s.cuda()).cpu().detach()\n",
    "    out_states_5s.append(cluster_model.encode(emb_5s)[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.4080)\n"
     ]
    }
   ],
   "source": [
    "n = 25 * 5\n",
    "print((out_states_30s[:n] == out_states_5s[0][:n]).sum() / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.4320)\n"
     ]
    }
   ],
   "source": [
    "print((out_states_30s[n:2*n] == out_states_5s[1][:n]).sum() / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5200)\n"
     ]
    }
   ],
   "source": [
    "print((out_states_30s[2*n:3*n] == out_states_5s[2][:n]).sum() / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.4880)\n"
     ]
    }
   ],
   "source": [
    "print((out_states_30s[3*n:4*n] == out_states_5s[3][:n]).sum() / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.3920)\n"
     ]
    }
   ],
   "source": [
    "print((out_states_30s[4*n:5*n] == out_states_5s[4][:n]).sum() / 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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}