{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "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": [
      "  0%|          | 0/20 [00:16<?, ?it/s]\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[5], line 15\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m n \u001b[38;5;129;01min\u001b[39;00m tqdm\u001b[38;5;241m.\u001b[39mtqdm(\u001b[38;5;28mrange\u001b[39m(tot_steps)):\n\u001b[1;32m      9\u001b[0m     filepaths \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m     10\u001b[0m         fp\n\u001b[1;32m     11\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m m \u001b[38;5;129;01min\u001b[39;00m metas[\u001b[38;5;241m-\u001b[39m(n \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m) \u001b[38;5;241m*\u001b[39m chunksize : \u001b[38;5;28mlen\u001b[39m(metas) \u001b[38;5;241m-\u001b[39m n \u001b[38;5;241m*\u001b[39m chunksize]\n\u001b[1;32m     12\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m (fp \u001b[38;5;241m:=\u001b[39m m\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms3_filepath\u001b[39m\u001b[38;5;124m\"\u001b[39m, m\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maudio_filepath\u001b[39m\u001b[38;5;124m\"\u001b[39m, m\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfilepath\u001b[39m\u001b[38;5;124m\"\u001b[39m))))\n\u001b[1;32m     13\u001b[0m         \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     14\u001b[0m     ]\n\u001b[0;32m---> 15\u001b[0m     download_duration_s, load_duration_s, audio_arrays \u001b[38;5;241m=\u001b[39m \u001b[43m_load_audio\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfilepaths\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m        \u001b[49m\u001b[43mSAMPLE_RATE\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmax_duration_per_file_s\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m60\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m     t0 \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[1;32m     21\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m  embedding audio...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[4], line 12\u001b[0m, in \u001b[0;36m_load_audio\u001b[0;34m(filepaths, sample_rate, max_duration_per_file_s)\u001b[0m\n\u001b[1;32m      7\u001b[0m tmp_out_filepaths \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m      8\u001b[0m     os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(tmp_dir, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124maudio_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mn\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mget_file_ext(filepath)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      9\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m n, filepath \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(filepaths)\n\u001b[1;32m     10\u001b[0m ]\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m  downloading audio...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 12\u001b[0m confirmed_downloads \u001b[38;5;241m=\u001b[39m \u001b[43mdownload_s3_files\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfilepaths\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtmp_out_filepaths\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m    \u001b[49m\u001b[43mchunksize\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m100\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m    \u001b[49m\u001b[43mn_cores\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m16\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m    \u001b[49m\u001b[43mjoblib_backend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mthreads\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m    \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m time\u001b[38;5;241m.\u001b[39msleep(\u001b[38;5;241m5\u001b[39m) \u001b[38;5;66;03m# make sure things close\u001b[39;00m\n\u001b[1;32m     21\u001b[0m local_filepath \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m     22\u001b[0m     filepath \u001b[38;5;28;01mif\u001b[39;00m b_confirmed \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     23\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m b_confirmed, filepath \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(confirmed_downloads, tmp_out_filepaths)\n\u001b[1;32m     24\u001b[0m ]\n",
      "File \u001b[0;32m~/glockenspiel/suno_utils/suno_utils/utils/s3.py:163\u001b[0m, in \u001b[0;36mdownload_s3_files\u001b[0;34m(from_s3_filepaths, to_local_filepaths, chunksize, n_cores, joblib_backend, silent)\u001b[0m\n\u001b[1;32m    160\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m confirmed_download\n\u001b[1;32m    162\u001b[0m work_items \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mzip\u001b[39m(from_s3_filepaths, to_local_filepaths))\n\u001b[0;32m--> 163\u001b[0m confirmed_downloads \u001b[38;5;241m=\u001b[39m \u001b[43m_apply_mp\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    164\u001b[0m \u001b[43m    \u001b[49m\u001b[43mprocess_f\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    165\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwork_items\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[43m    \u001b[49m\u001b[43mchunksize\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mchunksize\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[43m    \u001b[49m\u001b[43mn_cores\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_cores\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    168\u001b[0m \u001b[43m    \u001b[49m\u001b[43mjoblib_backend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mjoblib_backend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    169\u001b[0m \u001b[43m    \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    170\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    171\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m flat_input:\n\u001b[1;32m    172\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m confirmed_downloads[\u001b[38;5;241m0\u001b[39m]\n",
      "File \u001b[0;32m~/glockenspiel/suno_utils/suno_utils/utils/s3.py:127\u001b[0m, in \u001b[0;36m_apply_mp\u001b[0;34m(f, work_items, chunksize, n_cores, joblib_backend, silent)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    126\u001b[0m     n_cores \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mmin([\u001b[38;5;28mlen\u001b[39m(work_items_chunk), n_cores])\n\u001b[0;32m--> 127\u001b[0m     tmp_out \u001b[38;5;241m=\u001b[39m \u001b[43mParallel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mn_jobs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_cores\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprefer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mjoblib_backend\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    128\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdelayed\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mwork_item\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mwork_item\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mwork_items_chunk\u001b[49m\n\u001b[1;32m    129\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m     get_reusable_executor()\u001b[38;5;241m.\u001b[39mshutdown(wait\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    131\u001b[0m     out\u001b[38;5;241m.\u001b[39mextend(tmp_out)\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/joblib/parallel.py:1952\u001b[0m, in \u001b[0;36mParallel.__call__\u001b[0;34m(self, iterable)\u001b[0m\n\u001b[1;32m   1946\u001b[0m \u001b[38;5;66;03m# The first item from the output is blank, but it makes the interpreter\u001b[39;00m\n\u001b[1;32m   1947\u001b[0m \u001b[38;5;66;03m# progress until it enters the Try/Except block of the generator and\u001b[39;00m\n\u001b[1;32m   1948\u001b[0m \u001b[38;5;66;03m# reach the first `yield` statement. This starts the aynchronous\u001b[39;00m\n\u001b[1;32m   1949\u001b[0m \u001b[38;5;66;03m# dispatch of the tasks to the workers.\u001b[39;00m\n\u001b[1;32m   1950\u001b[0m \u001b[38;5;28mnext\u001b[39m(output)\n\u001b[0;32m-> 1952\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreturn_generator \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;43mlist\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/joblib/parallel.py:1595\u001b[0m, in \u001b[0;36mParallel._get_outputs\u001b[0;34m(self, iterator, pre_dispatch)\u001b[0m\n\u001b[1;32m   1592\u001b[0m     \u001b[38;5;28;01myield\u001b[39;00m\n\u001b[1;32m   1594\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backend\u001b[38;5;241m.\u001b[39mretrieval_context():\n\u001b[0;32m-> 1595\u001b[0m         \u001b[38;5;28;01myield from\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_retrieve()\n\u001b[1;32m   1597\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mGeneratorExit\u001b[39;00m:\n\u001b[1;32m   1598\u001b[0m     \u001b[38;5;66;03m# The generator has been garbage collected before being fully\u001b[39;00m\n\u001b[1;32m   1599\u001b[0m     \u001b[38;5;66;03m# consumed. This aborts the remaining tasks if possible and warn\u001b[39;00m\n\u001b[1;32m   1600\u001b[0m     \u001b[38;5;66;03m# the user if necessary.\u001b[39;00m\n\u001b[1;32m   1601\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env/lib/python3.10/site-packages/joblib/parallel.py:1707\u001b[0m, in \u001b[0;36mParallel._retrieve\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1702\u001b[0m \u001b[38;5;66;03m# If the next job is not ready for retrieval yet, we just wait for\u001b[39;00m\n\u001b[1;32m   1703\u001b[0m \u001b[38;5;66;03m# async callbacks to progress.\u001b[39;00m\n\u001b[1;32m   1704\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m ((\u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_jobs) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m) \u001b[38;5;129;01mor\u001b[39;00m\n\u001b[1;32m   1705\u001b[0m     (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_jobs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mget_status(\n\u001b[1;32m   1706\u001b[0m         timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtimeout) \u001b[38;5;241m==\u001b[39m TASK_PENDING)):\n\u001b[0;32m-> 1707\u001b[0m     \u001b[43mtime\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msleep\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0.01\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1708\u001b[0m     \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[1;32m   1710\u001b[0m \u001b[38;5;66;03m# We need to be careful: the job list can be filling up as\u001b[39;00m\n\u001b[1;32m   1711\u001b[0m \u001b[38;5;66;03m# we empty it and Python list are not thread-safe by\u001b[39;00m\n\u001b[1;32m   1712\u001b[0m \u001b[38;5;66;03m# default hence the use of the lock\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "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_val\", np.concatenate(val_data, axis=0).astype(np.float32))\n",
    "np.save(\"/home/minz/musicfm_concat_tr\", np.concatenate(train_data, axis=0).astype(np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1540795, 1024)\n",
      "(16635, 1024)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\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_tr.npy\")\n",
    "y_arr = np.load(\"/home/minz/musicfm_concat_val.npy\")[::5]\n",
    "\n",
    "print(X_arr.shape)\n",
    "print(y_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_arr = (X_arr - x_mean) / x_std\n",
    "y_arr = (y_arr - y_mean) / y_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "calculating codebook 0...\n",
      "train codebook 0...\n",
      "Clustering 1540795 points in 1024D to 8192 clusters, redo 1 times, 25 iterations\n",
      "  Preprocessing in 1.06 s\n",
      "  Iteration 24 (28283.60 s, search 28281.17 s): objective=1.20541e+09 imbalance=1.145 nsplit=0       \n",
      "finish train codebook 0...\n",
      " score: 28.046\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": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"/home/minz/musicfm_concat_centroids_norm_8192\", codebooked_centroids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "centroids = np.load(\"/home/minz/musicfm_concat_centroids_norm_8192.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 8192, 1024)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "centroids.shape"
   ]
  },
  {
   "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
}