{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import IPython\n",
    "import torchaudio\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3, list_s3_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/splice_all_samples_data_cleaned.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #how many metas are less than or equal to 30 duration_s\n",
    "print(len(metas))\n",
    "short_metas = [meta for meta in metas if meta[\"duration\"] <= 30_000]\n",
    "print(len(short_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asset_types = set([meta[\"asset_type_slug\"] for meta in metas])\n",
    "print(asset_types)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = metas[777000]\n",
    "print(meta[\"name\"])\n",
    "print(meta[\"duration\"]) # this is this duration in milliseconds\n",
    "for key, val in meta.items():\n",
    "    print(key, val)\n",
    "\n",
    "s3_filepath = f\"s3://suno-data/datasets/harvest/splice/audio/{meta['uuid']}.mp3\"\n",
    "audio, sr = read_from_s3(s3_filepath, read_f=torchaudio.load)\n",
    "print(audio.shape)\n",
    "duration_s = audio.shape[1] / sr\n",
    "print(duration_s)\n",
    "\n",
    "IPython.display.Audio(audio.numpy(), rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bundle_metas = []\n",
    "total_duration_s = 0\n",
    "\n",
    "pbar = tqdm(metas)\n",
    "# create a bundle metas\n",
    "for meta in pbar:\n",
    "    if meta[\"duration\"] >= 30_000 or meta[\"duration\"] < 1_000:\n",
    "        continue\n",
    "\n",
    "    s3_filepath = f\"s3://suno-data/datasets/harvest/splice/audio/{meta['uuid']}.mp3\"\n",
    "    duration_s = meta[\"duration\"] / 1000\n",
    "\n",
    "    tags = meta.get(\"tags\", None)\n",
    "    tags_list = []\n",
    "    if tags is not None:\n",
    "        for tag in tags:\n",
    "            tags_list.append(tag[\"label\"])\n",
    "\n",
    "    key = meta.get(\"key\", None)\n",
    "    bpm = meta.get(\"bpm\", None)\n",
    "\n",
    "    if key is not None:\n",
    "        tags_list.append(f\"key: {key}\")\n",
    "\n",
    "    if bpm is not None:\n",
    "        tags_list.append(f\"{bpm} bpm\")\n",
    "\n",
    "    # add the name as a tag\n",
    "    tags_list.append(meta[\"name\"])\n",
    "\n",
    "    # split the name on slash\n",
    "    name_parts = meta[\"name\"].split(\"/\")\n",
    "    if len(name_parts) > 1:\n",
    "        for part in name_parts:\n",
    "            tags_list.append(part)\n",
    "\n",
    "    # add the duration as a tag\n",
    "    tags_list.append(f\"duration_s: {duration_s}\")\n",
    "\n",
    "    new_meta = {\n",
    "        \"name\": meta[\"name\"],\n",
    "        \"id\": meta[\"uuid\"],\n",
    "        \"duration_s\": meta[\"duration\"] / 1000,\n",
    "        \"uuid\": meta[\"uuid\"],\n",
    "        \"s3_filepath\": s3_filepath,\n",
    "        \"tags\": tags_list,\n",
    "        \"key\": key,\n",
    "        \"bpm\": bpm,\n",
    "    }\n",
    "    bundle_metas.append(new_meta)\n",
    "\n",
    "    total_duration_s += meta[\"duration\"] / 1000\n",
    "    #pbar.set_postfix({\"total_duration_s\": total_duration_s})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_duration_h = total_duration_s / 3600\n",
    "print(f\"Total duration: {total_duration_h:.2f} hours\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print(len(bundle_metas))\n",
    "bundle_metas[6000]\n",
    "\n",
    "# save metas locally to jsonl\n",
    "output_filepath = \"/home/christian/code/christian/metadata/bundles/splice_samples_30s/metas.jsonl\"\n",
    "write_jsonl(bundle_metas, output_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = \"s3://suno-data/datasets/bundles/v4/splice_wetdry/dac_vae_tuned_25hz/part_0.npz\"\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "data = read_from_s3(filepath, read_f=np.load)\n",
    "print(data.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memamp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the metas from the bundle\n",
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/bundles/splice_samples_30s/metas.jsonl\")\n",
    "print(metas[0])\n",
    "\n",
    "# create metas map\n",
    "metas_map = {meta[\"id\"]: meta for meta in metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import write_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3, list_s3_dir\n",
    "import gc\n",
    "import sys\n",
    "import shutil\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 = \"/app/suno/data/splice_samples_30s/dac_vae_tuned_25hz\"\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first get a list of all the npz parts in the s3 bucket\n",
    "s3_bucket = \"s3://suno-data/datasets/bundles/v4/splice_samples_30s/dac_vae_tuned_25hz\"\n",
    "\n",
    "part_filepaths = list_s3_dir(s3_bucket)\n",
    "for part_filepath in part_filepaths:\n",
    "    print(part_filepath)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first get a list of all the npz parts in the s3 bucket\n",
    "s3_bucket = \"s3://suno-data/datasets/bundles/v4/splice_samples_30s/dac_vae_tuned_25hz\"\n",
    "\n",
    "part_filepaths = list_s3_dir(s3_bucket)\n",
    "# get all the part names\n",
    "part_names = [part_filepath[0].split(\"/\")[-1].split(\".\")[0] for part_filepath in part_filepaths]\n",
    "unique_part_names = list(set(part_names))\n",
    "print(len(unique_part_names))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import json\n",
    "\n",
    "idx = 1\n",
    "npz_filepath = f\"s3://suno-data/datasets/bundles/v4/splice_samples_30s/dac_vae_tuned_25hz/part_{idx}.npz\"\n",
    "\n",
    "vae_data = read_from_s3(npz_filepath, read_f=np.load)\n",
    "meta_ids = vae_data.keys()\n",
    "for meta_id in meta_ids:\n",
    "    print(meta_id)\n",
    "    meta_vae_data = vae_data[meta_id]\n",
    "    metadata = metas_map[meta_id]\n",
    "    print(metadata)\n",
    "    print(meta_vae_data.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load silence vae\n",
    "silence_vae = np.load(\"/home/christian/code/christian/metadata/dac_vae_tuned_25hz_30s_silence.npz\")[\"vae_data\"]\n",
    "print(silence_vae.shape)\n",
    "# swap the first two dimensions\n",
    "silence_vae = silence_vae.transpose(1, 0)\n",
    "print(silence_vae.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_parts = (0, 100)\n",
    "tr_parts = (101, 4148)\n",
    "\n",
    "use_local = True\n",
    "\n",
    "for dset_type in [\"tr\"]: #, \"tr\"\n",
    "\n",
    "    dset_filepaths = []\n",
    "    if dset_type == \"val\":\n",
    "        part_indices = range(val_parts[0], val_parts[1])\n",
    "    else:\n",
    "        part_indices = range(tr_parts[0], tr_parts[1])\n",
    "\n",
    "    metas = []\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",
    "\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(part_indices)\n",
    "    pbar.set_description(\"Hours: 0.00\")\n",
    "\n",
    "    for idx, part_idx in enumerate(pbar):\n",
    "        if use_local:\n",
    "            npz_filepath = f\"/mnt/localdisk/tmp/bundles/v4/splice_samples_30s/dac_vae_tuned_25hz/part_{idx}.npz\"\n",
    "            try:\n",
    "                vae_data = np.load(npz_filepath)\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "        else:\n",
    "            npz_filepath = f\"/home/christian/code/christian/data/splice_samples_30s/dac_vae_tuned_25hz/part_{idx}.npz\"\n",
    "            try:\n",
    "                vae_data = read_from_s3(npz_filepath, read_f=np.load)\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "        \n",
    "        # first collect the stuff we will write to disk\n",
    "        # here we have to iterate over each meta_id in vae_data\n",
    "        to_write_len_v = 0\n",
    "        chunks = []\n",
    "        for meta_id in vae_data.keys():\n",
    "            meta_vae_data = vae_data[meta_id].astype(np.float16)\n",
    "            metadata = metas_map[meta_id]\n",
    "            vae_chunk = meta_vae_data[:750, :]\n",
    "            n_vae_tokens = vae_chunk.shape[0]\n",
    "\n",
    "            # lets append silence vae data to the end of the chunk\n",
    "            pad_len = 750 - n_vae_tokens\n",
    "            cropped_silence_vae_data = silence_vae[:pad_len, :]\n",
    "            vae_chunk = np.concatenate([vae_chunk, cropped_silence_vae_data], axis=0)\n",
    "\n",
    "            # crop to 30s\n",
    "            to_write_len_v += vae_chunk.size\n",
    "\n",
    "            # create a new meta\n",
    "            new_meta = {\n",
    "                \"id\": id,\n",
    "                \"original_duration_s\": meta_vae_data.shape[0] / SEMANTIC_RATE_HZ,\n",
    "                \"n_vae_tokens\": n_vae_tokens,\n",
    "            }\n",
    "            for key, value in metadata.items():\n",
    "                new_meta[key] = value\n",
    "            metas.append(new_meta)\n",
    "\n",
    "            chunks.append(vae_chunk.reshape(-1))\n",
    "\n",
    "            # Add to total hours counter\n",
    "            audio_duration_hours = (meta_vae_data.shape[0] / SEMANTIC_RATE_HZ) / 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",
    "        # now write to disk\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",
    "        # convert chunks to single array\n",
    "        vae_chunks = np.concatenate(chunks, axis=0)\n",
    "\n",
    "        # convert vae_chunk to float16\n",
    "        out_mm_vae[n_offs_v : n_offs_v + vae_chunks.size] = vae_chunks\n",
    "\n",
    "        n_offs_v += vae_chunks.size\n",
    "\n",
    "    print(f\"Total hours of audio added: {total_hours:.2f} for {dset_type} set\")\n",
    "\n",
    "    print(len(metas))\n",
    "    write_jsonl(\n",
    "        metas,\n",
    "        os.path.join(out_metas_filepath),\n",
    "        do_append=True\n",
    "    )\n",
    "\n",
    "    out_mm_vae.flush()\n",
    "    del out_mm_vae, f\n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_gpt45",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}