{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import funcy\n",
    "import IPython\n",
    "import numpy as np\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n",
    "\n",
    "from suno_utils.utils.text import (    \n",
    "    write_jsonl,\n",
    "    read_jsonl,\n",
    "    write_json,\n",
    "    read_json,\n",
    "    normalize_whitespace,\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n",
    "\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    preload_models as preload_codec_models,\n",
    "    decode as codec_decode,\n",
    "    encode as codec_encode,\n",
    "    decode_stream_to_full_audio,\n",
    ")\n",
    "\n",
    "_ = preload_codec_models(CODEC_FILEPATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VAE_DIM = 128\n",
    "VAE_N_MEMMAP_TOKENS = 750\n",
    "\n",
    "SEMANTIC_N_CODEBOOKS = 1\n",
    "SEMANTIC_N_MEMMAP_TOKENS = 750\n",
    "\n",
    "#base_dir =  \"/app/suno/data/diff_syn_dpo/genius_t6_sampled_10k+diffv2_v1_t18\"\n",
    "#base_dir = \"/app/suno/data/dpo/diffv2_v1_t18/\"\n",
    "base_dir = \"/app/suno/data/diff_syn_dpo/interesting_clips_ahi_d3_20250504\"\n",
    "metas = read_jsonl(f\"{base_dir}/metas_tr.jsonl\", progress=True)\n",
    "vae_memmap_filepath = f\"{base_dir}/data_vae_tr.bin\"\n",
    "semantic_memmap_filepath = f\"{base_dir}/data_semantic_tr.bin\"\n",
    "\n",
    "# load memmaps\n",
    "vae_memmap = np.memmap(vae_memmap_filepath, dtype=np.float16, mode=\"r\")\n",
    "semantic_memmap = np.memmap(semantic_memmap_filepath, dtype=np.uint16, mode=\"r\")\n",
    "\n",
    "# reshape memmaps\n",
    "vae_data = vae_memmap.reshape(-1, VAE_N_MEMMAP_TOKENS, VAE_DIM)\n",
    "semantic_data = semantic_memmap.reshape(-1, SEMANTIC_N_MEMMAP_TOKENS, SEMANTIC_N_CODEBOOKS)\n",
    "\n",
    "print(vae_data.shape, semantic_data.shape, len(metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "\n",
    "#idx = 2700\n",
    "idx = 220\n",
    "\n",
    "neg_idx = idx if idx % 2 == 0 else idx + 1\n",
    "pos_idx = idx + 1\n",
    "\n",
    "meta = metas[neg_idx]\n",
    "neg_vae = vae_data[neg_idx]\n",
    "pos_vae = vae_data[pos_idx]\n",
    "\n",
    "neg_audio = codec_decode(neg_vae)\n",
    "pos_audio = codec_decode(pos_vae)\n",
    "\n",
    "neg_audio.play()\n",
    "pos_audio.play()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Merge\n",
    "\n",
    "Use this to create a new dpo memmap by merging existing memmaps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "\n",
    "#base_dir_1 = \"/app/suno/data/dpo/diffv2_v1_t18/\"\n",
    "base_dir_1 = \"/app2/suno/data/dpo/diff2_v2_d3_v1/\"\n",
    "base_dir_2 = \"/app/suno/data/diff_syn_dpo/genius_t6_sampled_10k\"\n",
    "\n",
    "out_dir = \"/app/suno/data/diff_syn_dpo/genius_t6_sampled_10k+diff2_v2_d3_v1\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "for split in [\"val\"]:\n",
    "    metas_1 = read_jsonl(f\"{base_dir_1}/metas_{split}.jsonl\", progress=True)\n",
    "    metas_2 = read_jsonl(f\"{base_dir_2}/metas_{split}.jsonl\", progress=True)\n",
    "\n",
    "    new_metas = metas_1 + metas_2\n",
    "\n",
    "    write_jsonl(new_metas, f\"{out_dir}/metas_{split}.jsonl\")\n",
    "\n",
    "    # merge vae memmaps\n",
    "    vae_memmap_1 = np.memmap(f\"{base_dir_1}/data_vae_{split}.bin\", dtype=np.float16, mode=\"r\")\n",
    "    vae_memmap_2 = np.memmap(f\"{base_dir_2}/data_vae_{split}.bin\", dtype=np.float16, mode=\"r\")\n",
    "\n",
    "    new_vae_memmap = np.concatenate([vae_memmap_1, vae_memmap_2], axis=0)\n",
    "    \n",
    "    new_vae_memmap_filepath = f\"{out_dir}/data_vae_{split}.bin\"\n",
    "    new_vae_memmap.tofile(new_vae_memmap_filepath)\n",
    "\n",
    "    # merge semantic memmaps\n",
    "    semantic_memmap_1 = np.memmap(f\"{base_dir_1}/data_semantic_{split}.bin\", dtype=np.uint16, mode=\"r\")\n",
    "    semantic_memmap_2 = np.memmap(f\"{base_dir_2}/data_semantic_{split}.bin\", dtype=np.uint16, mode=\"r\")\n",
    "\n",
    "    new_semantic_memmap = np.concatenate([semantic_memmap_1, semantic_memmap_2], axis=0)\n",
    "\n",
    "    new_semantic_memmap_filepath = f\"{out_dir}/data_semantic_{split}.bin\"\n",
    "    new_semantic_memmap.tofile(new_semantic_memmap_filepath)\n",
    "\n",
    "\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
}