{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f90714bd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:17:56.563424Z",
     "iopub.status.busy": "2024-06-26T19:17:56.563275Z",
     "iopub.status.idle": "2024-06-26T19:17:56.567481Z",
     "shell.execute_reply": "2024-06-26T19:17:56.567082Z",
     "shell.execute_reply.started": "2024-06-26T19:17:56.563405Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e8b43cc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:17:56.568155Z",
     "iopub.status.busy": "2024-06-26T19:17:56.568021Z",
     "iopub.status.idle": "2024-06-26T19:17:58.050491Z",
     "shell.execute_reply": "2024-06-26T19:17:58.049878Z",
     "shell.execute_reply.started": "2024-06-26T19:17:56.568141Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/tony/Work/neon/sunoGPT/\")\n",
    "from data_utils import get_batch, tokenize_batch, _load_tokenizer, get_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3bc93e35",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:17:58.051619Z",
     "iopub.status.busy": "2024-06-26T19:17:58.051406Z",
     "iopub.status.idle": "2024-06-26T19:18:24.506552Z",
     "shell.execute_reply": "2024-06-26T19:18:24.505946Z",
     "shell.execute_reply.started": "2024-06-26T19:17:58.051600Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-26_19:17:58]: Failed to import xformers.\n",
      "[2024-06-26_19:17:58]: Failed to import flash_attn RMSNorm. Falling back to torch RMSNorm.\n",
      "[2024-06-26_19:18:24]: number of parameters: 2615M\n"
     ]
    }
   ],
   "source": [
    "from modules.gpt import GPT, GPTConfig, GPTTrainConfig\n",
    "model_cfg = GPTConfig(\n",
    "    n_layer=24,\n",
    "    n_head=24,\n",
    "    d_head=128,\n",
    "    n_kv_head=4,\n",
    "    block_size=8832,\n",
    "    t_memmap=6016,\n",
    "    t_audio=6272,\n",
    "    t_text=2560,\n",
    ")\n",
    "train_cfg = GPTTrainConfig()\n",
    "model = GPT(model_cfg, train_cfg)\n",
    "cfg = model.config\n",
    "device = \"cpu\"\n",
    "tokenizer = _load_tokenizer(tokenizer_fp=\"/app/suno/data/chirp_v4/base/tokenizer_60k.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7560d3f4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:18:24.507697Z",
     "iopub.status.busy": "2024-06-26T19:18:24.507452Z",
     "iopub.status.idle": "2024-06-26T19:18:31.219254Z",
     "shell.execute_reply": "2024-06-26T19:18:31.218064Z",
     "shell.execute_reply.started": "2024-06-26T19:18:24.507679Z"
    }
   },
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models\n",
    "from suno_utils.tasks.dac_2c_12cb import (\n",
    "    encode as codec_encode,\n",
    "    decode as codec_decode,\n",
    "    EMBEDDING_RATE as CODEC_EMBEDDING_RATE,\n",
    ")\n",
    "_ = preload_codec_models(\"/app/suno/models/chirp_v2/dac_2c_25x12.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f0eef05",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:18:31.221142Z",
     "iopub.status.busy": "2024-06-26T19:18:31.220404Z",
     "iopub.status.idle": "2024-06-26T19:18:31.230915Z",
     "shell.execute_reply": "2024-06-26T19:18:31.230470Z",
     "shell.execute_reply.started": "2024-06-26T19:18:31.221103Z"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "import json\n",
    "from data_utils import read_jsonl\n",
    "from suno_utils.utils.text import read_json\n",
    "\n",
    "artist_condition = True\n",
    "cover_condition = True\n",
    "pack = True\n",
    "batch_size = 2\n",
    "batch_size_tokens = cfg.block_size * batch_size\n",
    "local_data_shard_dir = None\n",
    "\n",
    "t_memmap = cfg.t_memmap\n",
    "t_text = cfg.t_text\n",
    "semantic_n_codebooks = cfg.semantic_n_codebooks\n",
    "coarse_n_codebooks = cfg.coarse_n_codebooks\n",
    "semantic_vocab_size = cfg.semantic_vocab_size\n",
    "coarse_vocab_size = cfg.coarse_vocab_size\n",
    "\n",
    "print_with_time_master = print\n",
    "\n",
    "def load_dataset(\n",
    "    data_dir: str,\n",
    "    filename: str,\n",
    "    info_filename: str,\n",
    "    metas_filename: str,\n",
    "    weights_multiplier_map: dict,\n",
    "    is_finetune: bool,\n",
    ") -> list:\n",
    "    dataset_names = []\n",
    "    data_idx_lists = []  # used to randomly sample from the dataset\n",
    "    data_weights = []\n",
    "    data = np.memmap(os.path.join(data_dir, filename), dtype=np.uint16, mode=\"r\")\n",
    "    data = data.reshape(-1, t_memmap, semantic_n_codebooks + coarse_n_codebooks)\n",
    "    assert data[:100, :, :semantic_n_codebooks].max() <= semantic_vocab_size\n",
    "    assert data[:100, :, semantic_n_codebooks:].max() <= coarse_vocab_size\n",
    "    with open(os.path.join(data_dir, info_filename)) as f:\n",
    "        infos = json.load(f)\n",
    "    metas = read_jsonl(os.path.join(data_dir, metas_filename))\n",
    "    assert len(data) == len(metas), (len(data), len(metas))\n",
    "\n",
    "    artist_to_songs = defaultdict(list)\n",
    "    for i, m in enumerate(metas):\n",
    "        if \"artist\" in m:\n",
    "            artist_to_songs[f\"{m['dataset']}__{m['artist']}\"].append(i)\n",
    "    artist_to_songs = {k: v for k, v in artist_to_songs.items() if len(v) > 1}\n",
    "    if artist_condition:\n",
    "        assert len(artist_to_songs) > 0, \"no artist data found\"\n",
    "        print_with_time_master(f\"found {len(artist_to_songs):,} samples with artists on main process\")\n",
    "\n",
    "    idx_set = set()  # for checking that we don't have any duplicates\n",
    "    has_cover = False\n",
    "    # make sure we turn into int since keys in json get auto turned into strings\n",
    "    for dset_name in infos.keys():\n",
    "        if \"idx_map\" in infos[dset_name]:\n",
    "            infos[dset_name][\"idx_map\"] = {int(k): v for k, v in infos[dset_name][\"idx_map\"].items()}\n",
    "    for dset_name, info in infos.items():\n",
    "        dataset_names.append(dset_name)\n",
    "\n",
    "        if info.get(\"task\", \"default\") == \"default\":\n",
    "            assert \"idx_list\" in info\n",
    "            idx_list = info[\"idx_list\"][:]\n",
    "            idx_set |= set(idx_list)\n",
    "        elif info[\"task\"] == \"covers\":\n",
    "            assert pack, \"for now pack needs to be active to do covers\"\n",
    "            assert batch_size_tokens >= t_memmap * 2 + t_text, \"for covers we need double the blocksize\"\n",
    "            has_cover = True\n",
    "            # dict from original idx to list of covers idx\n",
    "            # use original idxs as the idx_list\n",
    "            idx_list = []\n",
    "            n_covers = 0\n",
    "            for idx, child_idx_l in info[\"idx_map\"].items():\n",
    "                idx_list.append(int(idx))\n",
    "                idx_set.add(int(idx))\n",
    "                idx_set |= set(child_idx_l)\n",
    "                n_covers += len(child_idx_l)\n",
    "            print_with_time_master(\n",
    "                f\"found {len(idx_list):,} samples with {n_covers:,} total covers on main process\"\n",
    "            )\n",
    "        else:\n",
    "            raise ValueError(f\"unknown task for {dset_name} in info file\")\n",
    "        random.shuffle(idx_list)\n",
    "\n",
    "        data_idx_lists.append(idx_list)\n",
    "        data_weights.append(len(idx_list) * weights_multiplier_map.get(dset_name, 1.0))\n",
    "    if cover_condition:\n",
    "        assert has_cover, \"no cover data found\"\n",
    "    weights_norm = np.sum(data_weights)\n",
    "    data_weights = [v / weights_norm for v in data_weights]\n",
    "\n",
    "    if not is_finetune:\n",
    "        print_with_time_master(f\"indexed {len(idx_set)/len(data)*100:.1f}% of data\")\n",
    "    for k in weights_multiplier_map.keys():\n",
    "        assert k in dataset_names\n",
    "\n",
    "    del idx_set\n",
    "    shard_info = \"\" if local_data_shard_dir is None else \" (sharded)\"\n",
    "    print_with_time_master(f\"{len(data):,} lines of {filename} loaded.{shard_info}\")\n",
    "    assert len(data) == len(metas)\n",
    "    assert len(infos) == len(dataset_names) == len(data_weights) == len(data_idx_lists)\n",
    "    return (\n",
    "        dataset_names,\n",
    "        data_idx_lists,\n",
    "        data_weights,\n",
    "        data,\n",
    "        metas,\n",
    "        infos,\n",
    "        artist_to_songs,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "be14a4ef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:18:31.232403Z",
     "iopub.status.busy": "2024-06-26T19:18:31.232242Z",
     "iopub.status.idle": "2024-06-26T19:18:31.324705Z",
     "shell.execute_reply": "2024-06-26T19:18:31.324237Z",
     "shell.execute_reply.started": "2024-06-26T19:18:31.232386Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found 231 samples with artists on main process\n",
      "found 76 samples with 357 total covers on main process\n",
      "indexed 99.0% of data\n",
      "5,879 lines of data_val.bin loaded.\n"
     ]
    }
   ],
   "source": [
    "data_dir = \"/app/suno/data/chirp_v4/multi\"\n",
    "\n",
    "val_filename = \"data_val.bin\"\n",
    "val_info_filename = \"info_val.json\"\n",
    "val_metas_filename = \"metas_val.jsonl\"\n",
    "\n",
    "weights_multiplier_map = {}\n",
    "is_finetune = False\n",
    "\n",
    "(\n",
    "    val_dataset_names,\n",
    "    val_data_idx_lists,\n",
    "    val_data_weights,\n",
    "    val_data,\n",
    "    val_metas,\n",
    "    val_info,\n",
    "    val_artist_to_songs,\n",
    ") = load_dataset(\n",
    "    data_dir,\n",
    "    val_filename,\n",
    "    val_info_filename,\n",
    "    val_metas_filename,\n",
    "    weights_multiplier_map,\n",
    "    is_finetune,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "63fb75a7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:18:31.326684Z",
     "iopub.status.busy": "2024-06-26T19:18:31.326539Z",
     "iopub.status.idle": "2024-06-26T19:18:31.345037Z",
     "shell.execute_reply": "2024-06-26T19:18:31.344637Z",
     "shell.execute_reply.started": "2024-06-26T19:18:31.326669Z"
    }
   },
   "outputs": [],
   "source": [
    "data_sampling_info = {\n",
    "    \"cfg\": cfg,\n",
    "    \"train_cfg\": train_cfg,\n",
    "    \"batch_size\": batch_size,\n",
    "    \"batch_size_tokens\": batch_size_tokens,\n",
    "    \"tokenizer_fp\": \"/app/suno/data/chirp_v4/base/tokenizer_60k.json\",\n",
    "    \"device\": device,\n",
    "    \"device_type\": \"cuda\" if \"cuda\" in str(device) else \"cpu\",\n",
    "    \"val\": {\n",
    "        \"data\": val_data,\n",
    "        \"metas\": val_metas,\n",
    "        \"infos\": val_info,\n",
    "        \"artist_to_songs\": val_artist_to_songs,\n",
    "        \"names\": val_dataset_names,\n",
    "        \"weights\": val_data_weights,\n",
    "        \"idx_lists\": val_data_idx_lists,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c300c398",
   "metadata": {},
   "source": [
    "### get_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e7e7f262",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-26T19:18:31.345715Z",
     "iopub.status.busy": "2024-06-26T19:18:31.345579Z",
     "iopub.status.idle": "2024-06-26T19:18:31.389169Z",
     "shell.execute_reply": "2024-06-26T19:18:31.388735Z",
     "shell.execute_reply.started": "2024-06-26T19:18:31.345700Z"
    }
   },
   "outputs": [],
   "source": [
    "suffix_first = True\n",
    "row_idx, x_arr = get_sample(\n",
    "    data_sampling_info,\n",
    "    \"val\",\n",
    "    dataset_idx=None,\n",
    "    row_idx=None,  # absolute, overrides dataset_idx\n",
    "    use_private=False,\n",
    "    inference=False,\n",
    "    suppress_text=False,\n",
    "    dummy_data=False,\n",
    "    min_text_offs=None,\n",
    "    suffix_first=suffix_first,\n",
    "    dropout_semantic=False,\n",
    "    allow_artist_condition=artist_condition,\n",
    "    allow_cover=cover_condition,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f12d8eb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba2d91a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4005ec3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7feffdf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a3bc38",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d7c4346",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}