{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import write_jsonl\n",
    "import gc\n",
    "import os\n",
    "import sys\n",
    "import shutil\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "import numpy as np\n",
    "import os\n",
    "base_dir = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t4/\"\n",
    "\n",
    "# get all directories in base_dir\n",
    "dirs = os.listdir(base_dir)\n",
    "print(len(dirs))\n",
    "# get all files in each directory\n",
    "\n",
    "model_name = \"16n_25hz_v45_infill_shared_flow_resume_1_75m\"\n",
    "\n",
    "def process_dir(dirpath):\n",
    "    a_metadata_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a__metadata.npz\")\n",
    "    a_upsampled_vae_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a_upsampled_vae.npz\")\n",
    "    b_metadata_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b__metadata.npz\")\n",
    "    b_upsampled_vae_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b_upsampled_vae.npz\")\n",
    "\n",
    "    a_metadata = np.load(a_metadata_filepath, allow_pickle=True)\n",
    "    b_metadata = np.load(b_metadata_filepath, allow_pickle=True)\n",
    "\n",
    "    # put metadata into dict\n",
    "    a_metadata_dict = {}\n",
    "    b_metadata_dict = {}\n",
    "    for key in a_metadata.keys():\n",
    "        a_metadata_dict[key] = a_metadata[key].tolist()\n",
    "    for key in b_metadata.keys():\n",
    "        b_metadata_dict[key] = b_metadata[key].tolist()\n",
    "\n",
    "    a_upsampled_vae = np.load(a_upsampled_vae_filepath)\n",
    "    b_upsampled_vae = np.load(b_upsampled_vae_filepath)\n",
    "\n",
    "    # load the mp3s \n",
    "    #a_mp3_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a.mp3\")\n",
    "    #b_mp3_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b.mp3\")\n",
    "    #print(f\"a: ear {a_metadata_dict['ear_score']:0.1f}, shimmer {a_metadata_dict['shimmer_score']:0.1f}, stereo {a_metadata_dict['stereo_width']:0.1f}, hoot_cer {a_metadata_dict['hoot_cer']}\")\n",
    "    #a_mp3 = Audio.from_file(a_mp3_filepath, n_channels=2).play()\n",
    "\n",
    "    #print(f\"b: ear {b_metadata_dict['ear_score']:0.1f}, shimmer {b_metadata_dict['shimmer_score']:0.1f}, stereo {b_metadata_dict['stereo_width']:0.1f}, hoot_cer {b_metadata_dict['hoot_cer']}\")\n",
    "    #b_mp3 = Audio.from_file(b_mp3_filepath, n_channels=2).play()\n",
    "\n",
    "    return a_metadata_dict, a_upsampled_vae, b_metadata_dict, b_upsampled_vae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create metas in parallel with joblib\n",
    "from tqdm import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "def process_dir_to_meta(dirpath):\n",
    "    try:    \n",
    "        a_metadata, a_upsampled_vae, b_metadata, b_upsampled_vae = process_dir(dirpath)\n",
    "    except Exception as e:\n",
    "        #print(f\"Error processing {dirpath}: {e}\")\n",
    "        return None\n",
    "\n",
    "    tags = a_metadata[\"tags\"]\n",
    "    text = a_metadata[\"text\"]\n",
    "    semantic_codes_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_semantic.npz\")\n",
    "\n",
    "    # figure out which is positive (better) and which is negative (worse) based on ear score\n",
    "    if a_metadata[\"ear_score\"] > b_metadata[\"ear_score\"]:\n",
    "        pos_metadata = a_metadata\n",
    "        neg_metadata = b_metadata\n",
    "        pos_vae_latents_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a_upsampled_vae.npz\")\n",
    "        neg_vae_latents_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b_upsampled_vae.npz\")\n",
    "    else:\n",
    "        pos_metadata = b_metadata\n",
    "        neg_metadata = a_metadata\n",
    "        pos_vae_latents_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_b_upsampled_vae.npz\")\n",
    "        neg_vae_latents_filepath = os.path.join(base_dir, dirpath, f\"{dirpath}_{model_name}_a_upsampled_vae.npz\")\n",
    "\n",
    "    meta = {\n",
    "        \"id\": dirpath,\n",
    "        \"tags\": tags,\n",
    "        \"text\": str(text),\n",
    "        \"pos_metadata\": pos_metadata,\n",
    "        \"neg_metadata\": neg_metadata,\n",
    "        \"pos_vae_latents_filepath\": pos_vae_latents_filepath,\n",
    "        \"neg_vae_latents_filepath\": neg_vae_latents_filepath,\n",
    "        \"semantic_codes_filepath\": semantic_codes_filepath,\n",
    "    }\n",
    "    return meta\n",
    "\n",
    "# Use joblib to parallelize\n",
    "metas = []\n",
    "results = Parallel(n_jobs=-1)(\n",
    "    delayed(process_dir_to_meta)(dirpath) for dirpath in tqdm(dirs)\n",
    ")\n",
    "# Filter out any None results (from failed processing)\n",
    "metas = [meta for meta in results if meta is not None]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can we look at the distribution of ear scores for the metas?\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "pos_ear_scores = [meta[\"pos_metadata\"][\"ear_score\"] for meta in metas]  \n",
    "neg_ear_scores = [meta[\"neg_metadata\"][\"ear_score\"] for meta in metas]\n",
    "\n",
    "# 95 percentile\n",
    "pos_ear_scores_95 = np.percentile(pos_ear_scores, 95)\n",
    "pos_ear_scores_05 = np.percentile(pos_ear_scores, 5)\n",
    "print(f\"pos_ear_scores_05: {pos_ear_scores_05}, pos_ear_scores_95: {pos_ear_scores_95}\")\n",
    "\n",
    "plt.hist(pos_ear_scores, bins=100, alpha=0.5)\n",
    "plt.hist(neg_ear_scores, bins=100, alpha=0.5)\n",
    "plt.axvline(pos_ear_scores_95, color='red', linestyle='--')\n",
    "plt.axvline(pos_ear_scores_05, color='red', linestyle='--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter metas based on pos_ear_score\n",
    "print(len(metas))\n",
    "filtered_metas = []\n",
    "for meta in metas:\n",
    "    if meta[\"pos_metadata\"][\"ear_score\"] < pos_ear_scores_05 or meta[\"pos_metadata\"][\"ear_score\"] > pos_ear_scores_95:\n",
    "        continue\n",
    "    filtered_metas.append(meta)\n",
    "\n",
    "print(f\"len(filtered_metas): {len(filtered_metas)}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split metas into train and val \n",
    "\n",
    "metas_tr = filtered_metas[:int(len(filtered_metas) * 0.98)]\n",
    "metas_val = filtered_metas[int(len(filtered_metas) * 0.98):]\n",
    "\n",
    "print(f\"len(metas_tr): {len(metas_tr)}\")\n",
    "print(f\"len(metas_val): {len(metas_val)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to create a memmap for sft we will select the higheset scoring upsample_id for each base_s3_id\n",
    "# we also need to grab the correct vae latents and semantic codes and text prompt\n",
    "\n",
    "\n",
    "SEMANTIC_RATE_HZ = 25\n",
    "CHUNK_SIZE_S = 30\n",
    "CHUNK_SIZE = int(CHUNK_SIZE_S * SEMANTIC_RATE_HZ)\n",
    "OUT_DATA_DIR = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t4/memmaps/syn_sft_t2\"\n",
    "LOCAL_DATA_DIR = \"/app2/suno/data/christian/outputs/v3-bootstrap-data-t4/\"\n",
    "\n",
    "if not os.path.exists(OUT_DATA_DIR):\n",
    "    os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "else:\n",
    "    #shutil.rmtree(OUT_DATA_DIR)\n",
    "    os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "\n",
    "for dset_type in [\"val\", \"tr\"]:\n",
    "    new_metas = []\n",
    "\n",
    "    if dset_type == \"tr\":\n",
    "        dset_metas = metas_tr\n",
    "    else:\n",
    "        dset_metas = metas_val\n",
    "\n",
    "    out_mm_vae_filepath = os.path.join(OUT_DATA_DIR, f\"data_vae_{dset_type}.bin\")\n",
    "    out_mm_semantic_filepath = os.path.join(OUT_DATA_DIR, f\"data_semantic_{dset_type}.bin\")\n",
    "    out_metas_filepath = os.path.join(OUT_DATA_DIR, f\"metas_{dset_type}.jsonl\")\n",
    "\n",
    "    n_offs_v = 0\n",
    "    n_offs_s = 0\n",
    "    to_write_len_v = 0\n",
    "    to_write_len_s = 0\n",
    "    total_hours = 0  # Counter for total hours of audio\n",
    "\n",
    "    out_mm_vae = np.memmap(\n",
    "        out_mm_vae_filepath, dtype=np.float16, mode=\"w+\", shape=(1,)\n",
    "    )\n",
    "    out_mm_semantic = np.memmap(\n",
    "        out_mm_semantic_filepath, dtype=np.uint16, mode=\"w+\", shape=(1,)\n",
    "    )\n",
    "\n",
    "    # clear the metas file\n",
    "    with open(out_metas_filepath, \"w\") as f:\n",
    "        f.write(\"\")\n",
    "\n",
    "    # Create a tqdm progress bar with hours counter\n",
    "    pbar = tqdm(dset_metas)\n",
    "    pbar.set_description(\"Hours: 0.00\")\n",
    "\n",
    "    for idx, (meta) in enumerate(pbar):\n",
    "\n",
    "        semantic_codes_filepath = meta[\"semantic_codes_filepath\"]\n",
    "        pos_vae_latents_filepath = meta[\"pos_vae_latents_filepath\"]\n",
    "\n",
    "        try:\n",
    "            semantic_data = np.load(semantic_codes_filepath)[\"semantic_codes\"]\n",
    "            vae_data = np.load(pos_vae_latents_filepath)[\"vae_latents\"].astype(np.float16)\n",
    "        except Exception as e:\n",
    "            raise ValueError(f\"Error loading {semantic_codes_filepath}: {e}\")\n",
    "\n",
    "        # check that vae and semantic data are the same length\n",
    "        assert vae_data.shape[0] == semantic_data.shape[0]\n",
    "\n",
    "        # convert vae_data to float16\n",
    "        vae_data = vae_data.astype(np.float16)\n",
    "\n",
    "        num_chunks = vae_data.shape[0] // CHUNK_SIZE\n",
    "\n",
    "        to_write_len_v = vae_data[:750,:].size * num_chunks\n",
    "        to_write_len_s = semantic_data[:750].size * num_chunks\n",
    "        \n",
    "        if to_write_len_s == 0 or to_write_len_v == 0:\n",
    "            continue\n",
    "\n",
    "        out_mm_vae = np.memmap(\n",
    "            out_mm_vae_filepath,\n",
    "            dtype=np.float16,\n",
    "            mode=\"r+\",\n",
    "            shape=(n_offs_v + to_write_len_v,),\n",
    "        )\n",
    "        \n",
    "        out_mm_semantic = np.memmap(\n",
    "            out_mm_semantic_filepath,\n",
    "            dtype=np.uint16,\n",
    "            mode=\"r+\",\n",
    "            shape=(n_offs_s + to_write_len_s,),\n",
    "        )\n",
    "\n",
    "        # Add to total hours counter\n",
    "        audio_duration_hours = (num_chunks * CHUNK_SIZE_S) / 3600\n",
    "        total_hours += audio_duration_hours\n",
    "        \n",
    "        # Update progress bar description with current total hours\n",
    "        pbar.set_description(f\"Hours: {total_hours:.2f}\")\n",
    "\n",
    "        for i in range(num_chunks):\n",
    "            # create a new meta\n",
    "            new_meta = {\n",
    "                \"upsample_id\": meta[\"id\"],\n",
    "                \"text\" : meta[\"text\"],\n",
    "                \"tags\" : meta[\"tags\"],\n",
    "                \"start_s\": i*CHUNK_SIZE_S,\n",
    "                \"end_s\": (i+1)*CHUNK_SIZE_S,\n",
    "                \"original_duration_s\": vae_data.shape[0] / SEMANTIC_RATE_HZ,\n",
    "                \"n_vae_tokens\": CHUNK_SIZE,\n",
    "                \"n_semantic_tokens\": CHUNK_SIZE,\n",
    "            }\n",
    "            new_metas.append(new_meta)\n",
    "\n",
    "            vae_chunk = vae_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE, :]\n",
    "            semantic_chunk = semantic_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE]\n",
    "\n",
    "            # convert vae_chunk to float16\n",
    "\n",
    "            out_mm_vae[n_offs_v : n_offs_v + vae_chunk.size] = vae_chunk.reshape(\n",
    "                -1,\n",
    "            )\n",
    "            out_mm_semantic[n_offs_s : n_offs_s + semantic_chunk.size] = semantic_chunk.reshape(\n",
    "                -1,\n",
    "            )\n",
    "\n",
    "            n_offs_s += semantic_chunk.size\n",
    "            n_offs_v += vae_chunk.size\n",
    "\n",
    "    print(f\"Total hours of audio added: {total_hours:.2f} for {dset_type} set\")\n",
    "  \n",
    "    write_jsonl(\n",
    "        new_metas,\n",
    "        os.path.join(out_metas_filepath),\n",
    "        do_append=True\n",
    "    )\n",
    "\n",
    "    out_mm_semantic.flush()\n",
    "    out_mm_vae.flush()\n",
    "    del out_mm_semantic, out_mm_vae, f\n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}