{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "08218f19",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:11:42.649240Z",
     "start_time": "2023-12-12T17:11:42.647782Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "76fc0b2c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:11:43.092643Z",
     "start_time": "2023-12-12T17:11:43.091052Z"
    }
   },
   "outputs": [],
   "source": [
    "# vocab: \n",
    "#   0-60_000 text\n",
    "#   1x0-3999   semantic\n",
    "#   8x0-4095  coarse\n",
    "\n",
    "#   4000 semantic pad token\n",
    "#   4001 semantic infer token\n",
    "#   4096 coarse pad token\n",
    "#   4097 coarse infer token\n",
    "#   4098 semantic interleave token\n",
    "\n",
    "# Memmaps:\n",
    "#   Nx9x3584 for audio tokens\n",
    "# Jsons:\n",
    "#   N*Dict with meta keys \n",
    "#     \"dataset\"\n",
    "#     \"original_id\", \"original_duration_s\",\n",
    "#     \"start_s\", \"end_s\", \n",
    "#     \"text_segments\", \"private_text_segments\",\n",
    "#     \"text\", \"private_text\",\n",
    "#     \"tags\", \"private_tags\",\n",
    "#     \"views\",\n",
    "#   Dict with meta keys {\"dataset\": [\"idx_list\"]}\n",
    "\n",
    "# Bundles (mert_v2_2x1k & dac_2c_25_8):\n",
    "# s3://suno-data/datasets/bundles/\n",
    "#  v1/youtube_music\n",
    "#  v1/genius_hq\n",
    "#  v1/jamendo\n",
    "#  v1/imslp\n",
    "#  v1/freesound\n",
    "#  v2/pond5\n",
    "#  v2/deezer\n",
    "#  v2/ytm_tagged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ce43e8ad",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:11:44.109807Z",
     "start_time": "2023-12-12T17:11:43.768857Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5164dde2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:12:34.869331Z",
     "start_time": "2023-12-12T17:12:34.866847Z"
    }
   },
   "outputs": [],
   "source": [
    "# TODO: add original_filepath on s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0e33ff3d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:11:46.861155Z",
     "start_time": "2023-12-12T17:11:44.879422Z"
    }
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import tqdm\n",
    "import torch\n",
    "import funcy\n",
    "import json\n",
    "import gc\n",
    "import re\n",
    "import random\n",
    "import tempfile\n",
    "import collections\n",
    "from collections import defaultdict\n",
    "from joblib import Parallel, delayed\n",
    "from transformers import BertTokenizer\n",
    "import time\n",
    "\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl, write_json, read_json, normalize_whitespace\n",
    "from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists\n",
    "\n",
    "TEXT_CODEBOOK_SIZE = 60_001\n",
    "TEXT_PAD_TOKEN = TEXT_CODEBOOK_SIZE\n",
    "TEXT_VOCAB_SIZE = 60_032\n",
    "\n",
    "SEMANTIC_CODEBOOK_SIZE = 4000\n",
    "SEMANTIC_N_CODEBOOKS = 1\n",
    "SEMANTIC_PAD_TOKEN = SEMANTIC_CODEBOOK_SIZE\n",
    "SEMANTIC_INFER_TOKEN = SEMANTIC_CODEBOOK_SIZE + 1\n",
    "SEMANTIC_VOCAB_SIZE = 4032\n",
    "SEMANTIC_RATE_HZ = 25\n",
    "SEMANTIC_SHIFT_FACTOR = 50\n",
    "assert(SEMANTIC_VOCAB_SIZE == (np.floor(SEMANTIC_CODEBOOK_SIZE // 64) + 1) * 64)\n",
    "\n",
    "COARSE_CODEBOOK_SIZE = 4096\n",
    "COARSE_N_CODEBOOKS = 8\n",
    "COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE\n",
    "COARSE_INFER_TOKEN = COARSE_CODEBOOK_SIZE + 1\n",
    "COARSE_VOCAB_SIZE = 4160\n",
    "COARSE_RATE_HZ = 25\n",
    "COARSE_SHIFT_FACTOR = 5\n",
    "assert(COARSE_CODEBOOK_SIZE + 3 < COARSE_VOCAB_SIZE)\n",
    "assert(COARSE_VOCAB_SIZE % 64 == 0)\n",
    "\n",
    "assert(SEMANTIC_RATE_HZ == COARSE_RATE_HZ)\n",
    "\n",
    "BLOCK_SIZE = 4288\n",
    "N_TOKENS_TEXT = 1152\n",
    "N_TOKENS_AUDIO = 3008  # max 120s of audio\n",
    "N_PAD_TOKENS_AUDIO = 3008\n",
    "# make sure we have enough space for shift 10\n",
    "assert(\n",
    "    BLOCK_SIZE >= (\n",
    "        N_TOKENS_TEXT + N_TOKENS_AUDIO + \n",
    "        SEMANTIC_N_CODEBOOKS * SEMANTIC_SHIFT_FACTOR + \n",
    "        (COARSE_N_CODEBOOKS - 1) * COARSE_SHIFT_FACTOR\n",
    "    )\n",
    ")\n",
    "\n",
    "SEMANTIC_EMBED_DIR = \"mert_25_2x4k\"\n",
    "CODEC_EMBED_DIR = \"dac_2c_25_8\"\n",
    "\n",
    "METAS_DIR = \"/home/tony/Work/tony/FineTuning_chirp_v2/metadata/\"\n",
    "OUT_DATA_DIR = \"/app/suno/data/chirp_v2_finetune_v9_classical_1\"\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b30bfb77",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:12:42.815276Z",
     "start_time": "2023-12-12T17:12:41.845934Z"
    }
   },
   "outputs": [],
   "source": [
    "# load manifests of IDs and text and tags etc\n",
    "meta_info_map = {\n",
    "    # \"genius_hq\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"genius_hq_v4_extra_filter.jsonl\"))},\n",
    "#    \"genius_hq\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"genius_hq_v6.jsonl\"))},\n",
    "#     \"youtube_music\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"youtube_music.jsonl\"))},\n",
    "#     \"freesound\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"freesound.jsonl\"))},\n",
    "#     \"jamendo\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"jamendo.jsonl\"))},\n",
    "     \"imslp\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"imslp.jsonl\"))},\n",
    "#     \"pond5_music\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"pond5_music.jsonl\"))},\n",
    "#     \"deezer\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"deezer.jsonl\"))},\n",
    "#     \"ytm_tagged\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"ytm_tagged.jsonl\"))},\n",
    "#     \"musescore\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"musescore.jsonl\"))},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dfc85bb4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-12T17:13:28.055827Z",
     "start_time": "2023-12-12T17:13:28.052665Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "703b9449-217e-4626-806c-0bb755a0cdb4 {'id': '703b9449-217e-4626-806c-0bb755a0cdb4', 'tags': ['strings (Nos.1-9) soprano', 'Baroque', 'Purcell', 'continuo (No.10)', 'Abdelazer', 'Z.570', 'Henry'], 'private_tags': ['strings (Nos.1-9) soprano', 'Baroque', 'Purcell', 'Complete Performance', 'continuo (No.10)', 'Abdelazer', 'Z.570', 'Henry']}\n"
     ]
    }
   ],
   "source": [
    "for k, v in meta_info_map[\"imslp\"].items():\n",
    "    print(k, v)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9c043b02",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:45:51.171683Z",
     "start_time": "2023-12-10T23:45:50.900153Z"
    }
   },
   "outputs": [],
   "source": [
    "def _verify_stuff(dset_name, meta_info):\n",
    "    if dset_name == \"genius_hq\":\n",
    "        assert (\n",
    "            \"private_text_segments\" in meta_info\n",
    "            or \"private_text\" in meta_info\n",
    "            or \"text_segments\" in meta_info\n",
    "            or \"text\" in meta_info\n",
    "        )\n",
    "\n",
    "\n",
    "def _trim_to_common(arr_1, arr_2):\n",
    "    common_len = min(len(arr_1), len(arr_2))\n",
    "    arr_1 = arr_1[:common_len]\n",
    "    arr_2 = arr_2[:common_len]\n",
    "    return arr_1, arr_2\n",
    "\n",
    "\n",
    "def _parse_arrays(dset_name, meta_info, semantic_arr, coarse_arr):\n",
    "    _verify_stuff(dset_name, meta_info)\n",
    "    # prep segment metas (use semantic for timekeeping)\n",
    "    segments_info = []\n",
    "    # first check if we have known segments\n",
    "    if \"text_segments\" in meta_info:\n",
    "        for m in meta_info[\"text_segments\"]:\n",
    "            segments_info.append(\n",
    "                (\n",
    "                    int(round(m[\"start_s\"] * SEMANTIC_RATE_HZ)),\n",
    "                    min(len(semantic_arr), int(round(m[\"end_s\"] * SEMANTIC_RATE_HZ))),\n",
    "                    m[\"text\"],\n",
    "                    m.get(\"private_text\"),\n",
    "                    True,\n",
    "                )\n",
    "            )\n",
    "    else:\n",
    "        # randomize offset to not get only multiples if no text available\n",
    "        offs = 0\n",
    "        #         if \"text\" not in meta_info and random.random() > 0.5:\n",
    "        #             # don't randomize offsets if we have lyrics\n",
    "        #             offs = random.randint(1, N_TOKENS_AUDIO-1)\n",
    "        #             segments_info.append((0, min(len(semantic_arr), offs), None, None, False))\n",
    "        total_steps = int(np.ceil((len(semantic_arr) - offs) / N_PAD_TOKENS_AUDIO))\n",
    "        for n in range(total_steps):\n",
    "            start_idx = offs + n * N_PAD_TOKENS_AUDIO\n",
    "            end_idx = min(len(semantic_arr), offs + (n + 1) * N_PAD_TOKENS_AUDIO)\n",
    "            if end_idx - start_idx < SEMANTIC_RATE_HZ:\n",
    "                # might as well skip mini ones\n",
    "                continue\n",
    "            segments_info.append(\n",
    "                (\n",
    "                    start_idx,\n",
    "                    end_idx,\n",
    "                    meta_info.get(\"text\", \"\")\n",
    "                    + f\"[part {round(start_idx / len(semantic_arr), 2)}]\",  # add text only to first piece\n",
    "                    meta_info.get(\"private_text\")\n",
    "                    if n == 0\n",
    "                    else None,  # add text only to first piece\n",
    "                    False,\n",
    "                )\n",
    "            )\n",
    "    #             if dset_name == \"musescore\":\n",
    "    #                 # break after first piece only text-audio pairs useful here\n",
    "    #                 # TODO: lang here is a hack\n",
    "    #                 meta_info[\"lang\"] = \"en\"\n",
    "    #                 break\n",
    "    #    print(segments_info)\n",
    "    arr_list = []\n",
    "    for sem_start_idx, sem_end_idx, text, private_text, is_aligned in segments_info:\n",
    "        if (\n",
    "            sem_end_idx - sem_start_idx > N_TOKENS_AUDIO\n",
    "            or sem_end_idx - sem_start_idx < SEMANTIC_RATE_HZ  # arbitrary\n",
    "        ):\n",
    "            continue\n",
    "        coarse_start_idx = int(round(sem_start_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ))\n",
    "        coarse_end_idx = int(round(sem_end_idx * COARSE_RATE_HZ / SEMANTIC_RATE_HZ))\n",
    "        assert sem_end_idx >= 0 and coarse_start_idx >= 0\n",
    "        if sem_end_idx > len(semantic_arr) or coarse_end_idx > len(coarse_arr):\n",
    "            continue\n",
    "        # get array segments\n",
    "        arr_s = semantic_arr[sem_start_idx:sem_end_idx, :SEMANTIC_N_CODEBOOKS].copy()\n",
    "        arr_c = coarse_arr[coarse_start_idx:coarse_end_idx, :COARSE_N_CODEBOOKS].copy()\n",
    "        # fix any alignment mistakes\n",
    "        arr_s, arr_c = _trim_to_common(arr_s, arr_c)\n",
    "        assert len(arr_s) == len(arr_c)\n",
    "        # concat and stack\n",
    "        if len(arr_c) < N_TOKENS_AUDIO:\n",
    "            arr_c = np.pad(\n",
    "                arr_c,\n",
    "                ((0, N_TOKENS_AUDIO - len(arr_c)), (0, 0)),\n",
    "                constant_values=COARSE_PAD_TOKEN,\n",
    "                mode=\"constant\",\n",
    "            )\n",
    "            arr_s = np.pad(\n",
    "                arr_s,\n",
    "                ((0, N_TOKENS_AUDIO - len(arr_s)), (0, 0)),\n",
    "                constant_values=SEMANTIC_PAD_TOKEN,\n",
    "                mode=\"constant\",\n",
    "            )\n",
    "        arr = np.concatenate([arr_s, arr_c], axis=-1)\n",
    "        arr = arr.astype(np.uint16)\n",
    "        assert arr.shape == (N_TOKENS_AUDIO, SEMANTIC_N_CODEBOOKS + COARSE_N_CODEBOOKS)\n",
    "        new_meta = {\n",
    "            \"id\": meta_info[\"id\"],\n",
    "            \"start_s\": round(sem_start_idx / SEMANTIC_RATE_HZ, 2),\n",
    "            \"end_s\": round(sem_end_idx / SEMANTIC_RATE_HZ, 2),\n",
    "            \"original_duration_s\": round(len(semantic_arr) / SEMANTIC_RATE_HZ, 2),\n",
    "        }\n",
    "        if text is not None:\n",
    "            new_meta[\"text\"] = text\n",
    "            if private_text is not None and private_text != private_text:\n",
    "                new_meta[\"text_private\"] = private_text\n",
    "            new_meta[\"text_lang\"] = meta_info.get(\"lang\", \"\")\n",
    "            new_meta[\"text_aligned\"] = is_aligned\n",
    "            new_meta[\"dset_suffix\"] = (\n",
    "                \"lyrics\" if new_meta[\"text_lang\"] == \"en\" else \"lyrics_foreign\"\n",
    "            )\n",
    "        if \"tags\" in meta_info:\n",
    "            new_meta[\"tags\"] = meta_info[\"tags\"]\n",
    "        if \"private_tags\" in meta_info:\n",
    "            if \"tags\" in meta_info:\n",
    "                # verify that superset\n",
    "                assert len(set(meta_info[\"tags\"]) - set(meta_info[\"private_tags\"])) == 0\n",
    "            if (\n",
    "                \"tags\" not in meta_info\n",
    "                or meta_info[\"private_tags\"] != meta_info[\"tags\"]\n",
    "            ):\n",
    "                new_meta[\"tags_private\"] = meta_info[\"private_tags\"]\n",
    "        if \"original_id\" in meta_info:\n",
    "            new_meta[\"original_id\"] = meta_info[\"original_id\"]\n",
    "        if \"views\" in meta_info:\n",
    "            new_meta[\"views\"] = meta_info[\"views\"]\n",
    "        arr_list.append((arr, new_meta))\n",
    "        del arr_s, arr_c\n",
    "    return arr_list\n",
    "\n",
    "\n",
    "def _process_archives(\n",
    "    dset_name,\n",
    "    s3_semantic_archive_filepaths,\n",
    "    s3_coarse_archive_filepaths,\n",
    "    relevant_metas,\n",
    "):\n",
    "    semantic_archive = {}\n",
    "    s3_semantic_archive_filepaths = set(s3_semantic_archive_filepaths)\n",
    "    # print(len(s3_semantic_archive_filepaths))\n",
    "    for s3_semantic_archive_filepath in s3_semantic_archive_filepaths:\n",
    "        if not check_s3_file_exists(s3_semantic_archive_filepath):\n",
    "            continue\n",
    "        for k, v in read_from_s3(s3_semantic_archive_filepath, read_f=np.load).items():\n",
    "            semantic_archive[k] = v\n",
    "    coarse_archive = {}\n",
    "    s3_coarse_archive_filepaths = set(s3_coarse_archive_filepaths)\n",
    "    # print(len(s3_coarse_archive_filepaths))\n",
    "    for s3_coarse_archive_filepath in s3_coarse_archive_filepaths:\n",
    "        if not check_s3_file_exists(s3_coarse_archive_filepath):\n",
    "            continue\n",
    "        for k, v in read_from_s3(s3_coarse_archive_filepath, read_f=np.load).items():\n",
    "            coarse_archive[k] = v\n",
    "    semantic_uids, coarse_uids = set(semantic_archive.keys()), set(\n",
    "        coarse_archive.keys()\n",
    "    )\n",
    "    assert (len(semantic_uids) < 10 and len(coarse_uids) < 10) or (\n",
    "        len(semantic_uids & coarse_uids) / (len(semantic_uids) + len(coarse_uids)) > 0.1\n",
    "    )\n",
    "    assert len(semantic_uids & coarse_uids) > 0\n",
    "    arr_list = []\n",
    "    for uid in semantic_uids & coarse_uids:\n",
    "        if uid not in relevant_metas:\n",
    "            continue\n",
    "        semantic_arr = semantic_archive[uid]\n",
    "        coarse_arr = coarse_archive[uid]\n",
    "        if (\n",
    "            np.abs(\n",
    "                len(coarse_arr) / COARSE_RATE_HZ - len(semantic_arr) / SEMANTIC_RATE_HZ\n",
    "            )\n",
    "            > 0.1\n",
    "        ):\n",
    "            # skip if embeddings not roughly the same duration\n",
    "            continue\n",
    "        semantic_arr, coarse_arr = _trim_to_common(semantic_arr, coarse_arr)\n",
    "        assert len(coarse_arr) == len(semantic_arr) * COARSE_RATE_HZ / SEMANTIC_RATE_HZ\n",
    "        arr_list.extend(\n",
    "            _parse_arrays(dset_name, relevant_metas[uid], semantic_arr, coarse_arr)\n",
    "        )\n",
    "    del semantic_archive, coarse_archive\n",
    "    gc.collect()\n",
    "    return arr_list\n",
    "\n",
    "\n",
    "def _collect_uids(\n",
    "    s3_semantic_metas_filepaths,\n",
    "    s3_coarse_metas_filepaths,\n",
    "):\n",
    "    semantic_uids = []\n",
    "    for fp in s3_semantic_metas_filepaths:\n",
    "        semantic_uids.extend([m[\"id\"] for m in read_from_s3(fp, read_f=read_jsonl)])\n",
    "    coarse_uids = []\n",
    "    for fp in s3_coarse_metas_filepaths:\n",
    "        coarse_uids.extend([m[\"id\"] for m in read_from_s3(fp, read_f=read_jsonl)])\n",
    "    return set(semantic_uids) & set(coarse_uids)\n",
    "\n",
    "\n",
    "def _prep_data(\n",
    "    dataset,\n",
    "    njobs=5,\n",
    "    chunksize=10,\n",
    "    is_val=False,\n",
    "    n_offs=0,\n",
    "):\n",
    "    dset_name, dset_version, (start_idx, end_idx), n_sem, n_coarse = dataset\n",
    "    dset_type = \"val\" if is_val else \"tr\"\n",
    "    out_mm_filepath = os.path.join(OUT_DATA_DIR, f\"data_{dset_type}.bin\")\n",
    "    out_metas_filepath = os.path.join(OUT_DATA_DIR, f\"metas_{dset_type}.jsonl\")\n",
    "    tot_duration_dict = defaultdict(int)\n",
    "    n_chunks = int(np.ceil((end_idx - start_idx) / chunksize))\n",
    "    for idx_chunk in tqdm.tqdm(\n",
    "        funcy.chunks(chunksize, list(range(start_idx, end_idx))), total=n_chunks\n",
    "    ):\n",
    "        n_jobs = np.min([njobs, chunksize, len(idx_chunk)])\n",
    "        # collect relevant parts of meta file to avoid copying all to subprocesses\n",
    "        tmp_uid_chunks = Parallel(n_jobs=n_jobs, prefer=\"threads\")(\n",
    "            delayed(_collect_uids)(\n",
    "                [\n",
    "                    f\"s3://suno-data/datasets/bundles/{dset_version}/{dset_name}/{SEMANTIC_EMBED_DIR}/\"\n",
    "                    + f\"metas/part_{idx_idx}.jsonl\"\n",
    "                    for idx_idx in range(idx * n_sem, (idx + 1) * n_sem)\n",
    "                ],\n",
    "                [\n",
    "                    f\"s3://suno-data/datasets/bundles/{dset_version}/{dset_name}/{CODEC_EMBED_DIR}/\"\n",
    "                    + f\"metas/part_{idx_idx}.jsonl\"\n",
    "                    for idx_idx in range(idx * n_coarse, (idx + 1) * n_coarse)\n",
    "                ],\n",
    "            )\n",
    "            for idx in idx_chunk\n",
    "        )\n",
    "        uids_per_part = {idx: tmp_uid_chunks[n] for n, idx in enumerate(idx_chunk)}\n",
    "        # print(len(uids_per_part))\n",
    "        # collect data\n",
    "        encoded_arrays_list = Parallel(n_jobs=n_jobs, prefer=\"processes\")(\n",
    "            delayed(_process_archives)(\n",
    "                dset_name,\n",
    "                [\n",
    "                    f\"s3://suno-data/datasets/bundles/{dset_version}/{dset_name}/{SEMANTIC_EMBED_DIR}/\"\n",
    "                    + f\"part_{idx_idx}.npz\"\n",
    "                    for idx_idx in range(idx * n_sem, (idx + 1) * n_sem)\n",
    "                ],\n",
    "                [\n",
    "                    f\"s3://suno-data/datasets/bundles/{dset_version}/{dset_name}/{CODEC_EMBED_DIR}/\"\n",
    "                    + f\"part_{idx_idx}.npz\"\n",
    "                    for idx_idx in range(idx * n_coarse, (idx + 1) * n_coarse)\n",
    "                ],\n",
    "                {\n",
    "                    uid: meta_info_map[dset_name][uid]\n",
    "                    for uid in uids_per_part[idx]\n",
    "                    if uid in meta_info_map[dset_name]\n",
    "                },\n",
    "            )\n",
    "            for idx in idx_chunk\n",
    "        )\n",
    "        # print(len(encoded_arrays_list))\n",
    "        add_metas = []\n",
    "        for encoded_arrays in encoded_arrays_list:\n",
    "            # print(len(encoded_arrays))\n",
    "            to_write_len = np.sum([arr.size for arr, _ in encoded_arrays])\n",
    "            if to_write_len == 0:\n",
    "                continue\n",
    "            out_mm = np.memmap(\n",
    "                out_mm_filepath,\n",
    "                dtype=np.uint16,\n",
    "                mode=\"r+\",\n",
    "                shape=(n_offs + to_write_len,),\n",
    "            )\n",
    "            for arr, arr_meta in encoded_arrays:\n",
    "                out_mm[n_offs : n_offs + arr.size] = arr.reshape(\n",
    "                    -1,\n",
    "                )\n",
    "                n_offs += arr.size\n",
    "                dataset_str = dset_name\n",
    "                if \"dset_suffix\" in arr_meta:\n",
    "                    dataset_str += f\"_{arr_meta['dset_suffix']}\"\n",
    "                add_meta = {\n",
    "                    \"dataset\": dataset_str,\n",
    "                    \"id\": arr_meta[\"id\"],\n",
    "                    \"start_s\": round(arr_meta[\"start_s\"], 2),\n",
    "                    \"end_s\": round(arr_meta[\"end_s\"], 2),\n",
    "                    \"original_duration_s\": arr_meta[\"original_duration_s\"],\n",
    "                }\n",
    "                if \"original_id\" in arr_meta:\n",
    "                    add_meta[\"original_id\"] = arr_meta[\"original_id\"]\n",
    "                if \"tags\" in arr_meta:\n",
    "                    add_meta[\"tags\"] = arr_meta[\"tags\"]\n",
    "                if \"tags_private\" in arr_meta:\n",
    "                    add_meta[\"tags_private\"] = arr_meta[\"tags_private\"]\n",
    "                if \"text\" in arr_meta:\n",
    "                    add_meta[\"text\"] = arr_meta[\"text\"]\n",
    "                if \"text_private\" in arr_meta:\n",
    "                    add_meta[\"text_private\"] = arr_meta[\"text_private\"]\n",
    "                if \"text_lang\" in arr_meta:\n",
    "                    add_meta[\"text_lang\"] = arr_meta[\"text_lang\"]\n",
    "                if \"text_aligned\" in arr_meta:\n",
    "                    add_meta[\"text_aligned\"] = arr_meta[\"text_aligned\"]\n",
    "                if \"views\" in arr_meta:\n",
    "                    add_meta[\"views\"] = arr_meta[\"views\"]\n",
    "                tot_duration_dict[dataset_str] += (\n",
    "                    arr_meta[\"end_s\"] - arr_meta[\"start_s\"]\n",
    "                )\n",
    "                add_metas.append(add_meta)\n",
    "            # write it once\n",
    "            out_mm.flush()\n",
    "            del out_mm\n",
    "        write_jsonl(\n",
    "            add_metas,\n",
    "            os.path.join(out_metas_filepath),\n",
    "            do_append=bool(n_offs != 0),\n",
    "        )\n",
    "        del encoded_arrays_list\n",
    "        gc.collect()\n",
    "    for k, v in tot_duration_dict.items():\n",
    "        print(f\"{round(v / 60 / 60):,} hours of {k}\")\n",
    "    return n_offs\n",
    "\n",
    "\n",
    "def prep_data(\n",
    "    datasets,\n",
    "    is_val=False,\n",
    "    njobs=5,\n",
    "    chunksize=10,\n",
    "):\n",
    "    n_offs = 0\n",
    "    dset_type = \"val\" if is_val else \"tr\"\n",
    "    out_mm_filepath = os.path.join(OUT_DATA_DIR, f\"data_{dset_type}.bin\")\n",
    "    out_metas_filepath = os.path.join(OUT_DATA_DIR, f\"metas_{dset_type}.jsonl\")\n",
    "    out_info_filepath = os.path.join(OUT_DATA_DIR, f\"info_{dset_type}.json\")\n",
    "    out_mm = np.memmap(out_mm_filepath, dtype=np.uint16, mode=\"w+\", shape=(1,))\n",
    "    with open(out_metas_filepath, \"w\") as f:\n",
    "        f.write(\"\")\n",
    "    print(\"start prepare data\")\n",
    "    for dataset in datasets:\n",
    "        n_offs = _prep_data(\n",
    "            dataset,\n",
    "            njobs=njobs,\n",
    "            chunksize=chunksize,\n",
    "            is_val=is_val,\n",
    "            n_offs=n_offs,\n",
    "        )\n",
    "    datasets_info = {}\n",
    "    with open(out_metas_filepath) as f:\n",
    "        n = 0\n",
    "        for line in f:\n",
    "            line = line.strip()\n",
    "            if len(line) == 0:\n",
    "                continue\n",
    "            m = json.loads(line)\n",
    "            if m[\"dataset\"] not in datasets_info:\n",
    "                datasets_info[m[\"dataset\"]] = {\"idx_list\": []}\n",
    "            datasets_info[m[\"dataset\"]][\"idx_list\"].append(n)\n",
    "            n += 1\n",
    "    write_json(datasets_info, out_info_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f862095a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:45:51.219576Z",
     "start_time": "2023-12-10T23:45:51.172852Z"
    }
   },
   "outputs": [],
   "source": [
    "NJOBS = 40\n",
    "CHUNKSIZE = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3339fe31",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:46:03.446014Z",
     "start_time": "2023-12-10T23:45:51.221084Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start prepare data\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:12<00:00, 12.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "17 hours of imslp_lyrics_foreign\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# (start_idx, end_idx), n_archives_semantic, n_archives_coarse\n",
    "datasets = [\n",
    "    # (\"youtube_music\", \"v1\", (0, 1), 1, 1),\n",
    "#    (\"genius_hq\", \"v1\", (0, 10), 1, 1),\n",
    "#     (\"freesound\", \"v1\", (0, 2), 1, 1),\n",
    "     (\"imslp\", \"v1\", (0, 2), 1, 1),\n",
    "#     (\"jamendo\", \"v1\", (0, 1), 1, 1),\n",
    "#     (\"pond5_music\", \"v2\", (0, 1), 1, 1),\n",
    "#     (\"deezer\", \"v2\", (0, 1), 1, 1),\n",
    "#     (\"ytm_tagged\", \"v2\", (0, 1), 1, 1),\n",
    "#     (\"musescore\", \"v2\", (0, 1), 1, 1),\n",
    "]\n",
    "prep_data(\n",
    "    datasets,\n",
    "    is_val=True,\n",
    "    njobs=NJOBS,\n",
    "    chunksize=CHUNKSIZE,\n",
    ")\n",
    "# 26 hours of youtube_music\n",
    "#  3 hours of youtube_music_lyrics\n",
    "#  2 hours of youtube_music_lyrics_foreign\n",
    "#  4 hours of genius_hq\n",
    "#  9 hours of genius_hq_lyrics\n",
    "#  5 hours of genius_hq_lyrics_foreign\n",
    "#  1 hours of freesound\n",
    "# 29 hours of imslp\n",
    "# 36 hours of jamendo\n",
    "# 14 hours of pond5_music\n",
    "# 12 hours of deezer\n",
    "#  7 hours of deezer_lyrics\n",
    "#  8 hours of deezer_lyrics_foreign\n",
    "# 33 hours of ytm_tagged\n",
    "#  3 hours of musescore_lyrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a2dfa127",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:44.984384Z",
     "start_time": "2023-12-10T23:46:03.447241Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start prepare data\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 14/14 [07:39<00:00, 32.86s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8,039 hours of imslp_lyrics_foreign\n"
     ]
    }
   ],
   "source": [
    "# (start_idx, end_idx), n_archives_semantic, n_archives_coarse\n",
    "datasets = [\n",
    "#    (\"youtube_music\", \"v1\", (1, 4204), 1, 1),\n",
    "#    (\"genius_hq\", \"v1\", (10, 4000), 1, 1),\n",
    "#     (\"freesound\", \"v1\", (2, 1021), 1, 1),\n",
    "     (\"imslp\", \"v1\", (2, 558), 1, 1),\n",
    "#     (\"jamendo\", \"v1\", (1, 112), 1, 1),\n",
    "#     (\"pond5_music\", \"v2\", (1, 4138), 1, 1),\n",
    "#     (\"deezer\", \"v2\", (1, 1538), 1, 1),\n",
    "#     (\"ytm_tagged\", \"v2\", (1, 5545), 1, 1),\n",
    "#     (\"musescore\", \"v2\", (1, 33), 1, 1),\n",
    "]\n",
    "prep_data(\n",
    "    datasets,\n",
    "    is_val=False,\n",
    "    njobs=NJOBS,\n",
    "    chunksize=CHUNKSIZE,\n",
    ")\n",
    "# youtube_music: ~1.5h runtime\n",
    "#   113,879 hours of youtube_music\n",
    "#    12,287 hours of youtube_music_lyrics\n",
    "#    10,043 hours of youtube_music_lyrics_foreign\n",
    "# genius_hq: ~1.3h runtime\n",
    "#    17,615 hours of genius_hq\n",
    "#    41,544 hours of genius_hq_lyrics\n",
    "#    17,307 hours of genius_hq_lyrics_foreign\n",
    "# freesound: ~0.2h runtime\n",
    "#       410 hours of freesound\n",
    "# imslp: ~0.2h runtime\n",
    "#    19,514 hours of imslp\n",
    "# jamendo: ~0.1h runtime\n",
    "#     3,726 hours of jamendo\n",
    "# pond5_music: ~1.1h runtime\n",
    "#    62,117 hours of pond5_music\n",
    "# deezer: ~0.5h runtime\n",
    "#    12,287 hours of deezer\n",
    "#    10,175 hours of deezer_lyrics\n",
    "#     6,699 hours of deezer_lyrics_foreign\n",
    "# ytm_tagged: ~2h runtime\n",
    "#   152,162 hours of ytm_tagged\n",
    "# musescore: ~0.1h runtime\n",
    "#       103 hours of musescore_lyrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c572bd6f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.199541Z",
     "start_time": "2023-12-10T23:53:44.986035Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15G\t/app/suno/data/chirp_v2_finetune_v7_classical/data_tr.bin\r\n"
     ]
    }
   ],
   "source": [
    "!du -hs /app/suno/data/chirp_v2_finetune_v7_classical/data_tr.bin\n",
    "# 1003G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "999206cd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.205742Z",
     "start_time": "2023-12-10T23:53:45.202841Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DONE!\n"
     ]
    }
   ],
   "source": [
    "print(\"DONE!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2ae23b40",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.274517Z",
     "start_time": "2023-12-10T23:53:45.207131Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['imslp_lyrics_foreign']) 1\n"
     ]
    }
   ],
   "source": [
    "with open(\"/app/suno/data/chirp_v2_finetune_v7_classical/info_tr.json\", \"r\") as fp:\n",
    "    info_tr = json.load(fp)\n",
    "print(info_tr.keys(), len(info_tr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eb9bed5b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.456495Z",
     "start_time": "2023-12-10T23:53:45.275687Z"
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'BREAK' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[14], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mBREAK\u001b[49m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'BREAK' is not defined"
     ]
    }
   ],
   "source": [
    "BREAK"
   ]
  },
  {
   "cell_type": "raw",
   "id": "0793420f",
   "metadata": {},
   "source": [
    "# NEED: tokenizer_60k.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab57df8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.457718Z",
     "start_time": "2023-12-10T23:53:45.457710Z"
    }
   },
   "outputs": [],
   "source": [
    "# # verify\n",
    "# mm = np.memmap(\"/mnt/data/georg/data/chirp_v2/data_val.bin\", dtype=np.uint16, mode=\"r\")\n",
    "# metas = read_jsonl(\"/mnt/data/georg/data/chirp_v2/metas_val.jsonl\")\n",
    "# mm = mm.reshape(-1, 3008, 9)\n",
    "# assert(len(mm) == len(metas))\n",
    "# assert(mm[:100,:,0].min() >= 0)\n",
    "# assert(mm[:100,:,0].max() <= 4000)\n",
    "# assert(mm[:100,:,1:].min() >= 0)\n",
    "# assert(mm[:100,:,1:].max() <= 4096)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba4ee8dc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.458499Z",
     "start_time": "2023-12-10T23:53:45.458491Z"
    }
   },
   "outputs": [],
   "source": [
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/info_val.json s3://suno-data/georg/data/chirp_v2/\n",
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/metas_val.jsonl s3://suno-data/georg/data/chirp_v2/\n",
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/data_val.bin s3://suno-data/georg/data/chirp_v2/\n",
    "\n",
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/info_tr.json s3://suno-data/georg/data/chirp_v2/\n",
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/metas_tr.jsonl s3://suno-data/georg/data/chirp_v2/\n",
    "# !aws s3 cp /mnt/data/georg/data/chirp_v2/data_tr.bin s3://suno-data/georg/data/chirp_v2/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5bdf8ac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.459142Z",
     "start_time": "2023-12-10T23:53:45.459134Z"
    }
   },
   "outputs": [],
   "source": [
    "# !du -hs /mnt/data/georg/data/chirp_v2/data_tr.bin\n",
    "# # 1003G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbcbca32",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.459783Z",
     "start_time": "2023-12-10T23:53:45.459775Z"
    }
   },
   "outputs": [],
   "source": [
    "# # randomly listen to some stuff\n",
    "# from suno_utils.tasks.dac_2c import preload_models as preload_codec_models\n",
    "# from suno_utils.tasks.dac_2c import (\n",
    "#     encode as codec_encode,\n",
    "#     decode as codec_decode, \n",
    "#     EMBEDDING_RATE as CODEC_EMBEDDING_RATE,\n",
    "# )\n",
    "# _ = preload_codec_models(\"/mnt/data/georg/models/chirp_v2/dac_2c_25x8.pt\")\n",
    "# mm = np.memmap(\"/mnt/data/georg/data/chirp_v2/data_val.bin\", dtype=np.uint16, mode=\"r\")\n",
    "# mm = mm.reshape(-1, 3008, 9)\n",
    "# test_metas = read_jsonl(\"/mnt/data/georg/data/chirp_v2/metas_val.jsonl\")\n",
    "# test_info = read_json(\"/mnt/data/georg/data/chirp_v2/info_val.json\")\n",
    "# assert(len(test_metas) == len(mm))\n",
    "# idx_list = list(range(len(test_metas)))\n",
    "# #random.shuffle(idx_list)\n",
    "# #idx_list = [idx for idx in idx_list if \"text\" in test_metas[idx]]\n",
    "# print(len(mm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffa828e8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.460465Z",
     "start_time": "2023-12-10T23:53:45.460457Z"
    }
   },
   "outputs": [],
   "source": [
    "# show text and audio\n",
    "# idx = random.choice(idx_list)\n",
    "# idx = random.choice(test_info[\"pond5_music\"][\"idx_list\"])\n",
    "# idx = random.choice(test_info[\"genius_hq_lyrics\"][\"idx_list\"])\n",
    "# idx = random.choice(test_info[\"musescore_lyrics\"][\"idx_list\"])\n",
    "# idx = random.choice(test_info[\"deezer_lyrics\"][\"idx_list\"])\n",
    "# idx = random.choice(test_info[\"youtube_music_lyrics_foreign\"][\"idx_list\"])\n",
    "# idx_key = random.choice(list(test_info.keys()))\n",
    "# print(idx_key)\n",
    "# idx = random.choice(test_info[idx_key][\"idx_list\"])\n",
    "# assert(\"original_duration_s\" in test_metas[idx])\n",
    "# print(\"tags:\", test_metas[idx].get(\"tags\"))\n",
    "# arr = mm[idx,1:].copy().astype(np.int16)[:,1:]\n",
    "# pad_idx_arr = np.where(arr == 4096)[0]\n",
    "# if len(pad_idx_arr) > 0:\n",
    "#     arr = arr[:pad_idx_arr[0],:]\n",
    "# a = codec_decode(arr)\n",
    "# a.play()\n",
    "# print(\"text:\", test_metas[idx].get(\"text\"))\n",
    "# plt.plot(a.array_float[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7008ddc7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "067a561b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73c92bda",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "08665ec0",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "747f0363",
   "metadata": {},
   "source": [
    "#### get all tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73286512",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.461277Z",
     "start_time": "2023-12-10T23:53:45.461269Z"
    }
   },
   "outputs": [],
   "source": [
    "test_metas = read_jsonl(\"/mnt/data/georg/data/chirp_v2/metas_val.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ed7922",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.461960Z",
     "start_time": "2023-12-10T23:53:45.461952Z"
    }
   },
   "outputs": [],
   "source": [
    "tags = []\n",
    "for e in train_metas:\n",
    "    tags.extend(e.get(\"tags\", []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eed69a8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.462506Z",
     "start_time": "2023-12-10T23:53:45.462499Z"
    }
   },
   "outputs": [],
   "source": [
    "vc_tags = pd.Series(tags).value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf1c83e8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.463022Z",
     "start_time": "2023-12-10T23:53:45.463015Z"
    }
   },
   "outputs": [],
   "source": [
    "vc_tags.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f76f1e83",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.463839Z",
     "start_time": "2023-12-10T23:53:45.463831Z"
    }
   },
   "outputs": [],
   "source": [
    "vc_tags = vc_tags.to_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6766efb1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.464408Z",
     "start_time": "2023-12-10T23:53:45.464401Z"
    }
   },
   "outputs": [],
   "source": [
    "vc_tags.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fece8109",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.465164Z",
     "start_time": "2023-12-10T23:53:45.465157Z"
    }
   },
   "outputs": [],
   "source": [
    "a = vc_tags[vc_tags[\"count\"]>=3][\"count\"].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52e1a982",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.465873Z",
     "start_time": "2023-12-10T23:53:45.465866Z"
    }
   },
   "outputs": [],
   "source": [
    "test_metas = read_jsonl(\"/mnt/data/georg/data/chirp_v2/metas_val.jsonl\")\n",
    "test_info = read_json(\"/mnt/data/georg/data/chirp_v2/info_val.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6b9f19e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-10T23:53:45.466577Z",
     "start_time": "2023-12-10T23:53:45.466570Z"
    }
   },
   "outputs": [],
   "source": [
    "idx = random.choice(test_info[\"musescore_lyrics\"][\"idx_list\"])\n",
    "print(idx)\n",
    "m = test_metas[idx]\n",
    "re.sub(r\"\\[.*?\\]\", \"\", m[\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dffab2d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4531aac2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0e1aafb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143ac0a7",
   "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.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
