{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "08218f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "76fc0b2c",
   "metadata": {},
   "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": 9,
   "id": "ce43e8ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5164dde2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: add original_filepath on s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0e33ff3d",
   "metadata": {},
   "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",
    "\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",
    "# 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/georg/notebooks/gpt/chirp_v2/metadata/\"\n",
    "OUT_DATA_DIR = \"/mnt/data/georg/data/chirp_v2\"\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b30bfb77",
   "metadata": {},
   "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.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": 10,
   "id": "9c043b02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _verify_stuff(dset_name, meta_info):\n",
    "    if dset_name == \"genius_hq\":\n",
    "        assert(\n",
    "            \"private_text_segments\" in meta_info or \n",
    "            \"private_text\" in meta_info or\n",
    "            \"text_segments\" in meta_info or \n",
    "            \"text\" in meta_info\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",
    "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",
    "                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",
    "    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",
    "        for n in range(int(np.ceil((len(semantic_arr)-offs)/N_TOKENS_AUDIO))):\n",
    "            start_idx = offs + n*N_TOKENS_AUDIO\n",
    "            end_idx = min(len(semantic_arr), offs+(n+1)*N_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",
    "                start_idx, \n",
    "                end_idx,\n",
    "                meta_info.get(\"text\") if n == 0 else None,  # add text only to first piece\n",
    "                meta_info.get(\"private_text\") if n == 0 else None,  # add text only to first piece\n",
    "                False,\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",
    "    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 or \n",
    "            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[\"lang\"]\n",
    "            new_meta[\"text_aligned\"] = is_aligned\n",
    "            new_meta[\"dset_suffix\"] = \"lyrics\" if meta_info[\"lang\"] == \"en\" else \"lyrics_foreign\"  \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 \"tags\" not in meta_info or meta_info[\"private_tags\"] != meta_info[\"tags\"]:\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",
    "def _process_archives(\n",
    "    dset_name,\n",
    "    s3_semantic_archive_filepaths, \n",
    "    s3_coarse_archive_filepaths,\n",
    "    relevant_metas,\n",
    "):\n",
    "    semantic_archive = {}\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",
    "    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(coarse_archive.keys())\n",
    "    assert(\n",
    "        (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 np.abs(len(coarse_arr) / COARSE_RATE_HZ - len(semantic_arr) / SEMANTIC_RATE_HZ) > 0.1:\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(_parse_arrays(dset_name, relevant_metas[uid], semantic_arr, coarse_arr))\n",
    "    del semantic_archive, coarse_archive\n",
    "    gc.collect()\n",
    "    return arr_list\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",
    "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(funcy.chunks(chunksize, list(range(start_idx, end_idx))), total=n_chunks):\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",
    "        # 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",
    "        for encoded_arrays in encoded_arrays_list:\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, dtype=np.uint16, mode=\"r+\", 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(-1,)\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] += arr_meta[\"end_s\"] - arr_meta[\"start_s\"]\n",
    "                with open(out_metas_filepath, \"a\") as f:\n",
    "                    f.write(json.dumps(add_meta) + \"\\n\")\n",
    "            out_mm.flush()\n",
    "            del out_mm\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",
    "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",
    "    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": 18,
   "id": "3339fe31",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:37<00:00, 37.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 hours of youtube_music_lyrics\n",
      "26 hours of youtube_music\n",
      "2 hours of youtube_music_lyrics_foreign\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:33<00:00, 33.86s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9 hours of genius_hq_lyrics\n",
      "5 hours of genius_hq_lyrics_foreign\n",
      "4 hours of genius_hq\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:21<00:00, 21.74s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 hours of freesound\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:32<00:00, 32.23s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "29 hours of imslp\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:30<00:00, 30.76s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36 hours of jamendo\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:30<00:00, 30.12s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14 hours of pond5_music\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:34<00:00, 34.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8 hours of deezer_lyrics_foreign\n",
      "12 hours of deezer\n",
      "7 hours of deezer_lyrics\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:32<00:00, 32.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33 hours of ytm_tagged\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 1/1 [00:29<00:00, 29.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 hours of musescore_lyrics\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, 1), 1, 1),\n",
    "    (\"freesound\", \"v1\", (0, 2), 1, 1),\n",
    "    (\"imslp\", \"v1\", (0, 1), 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",
    ")\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": 19,
   "id": "a2dfa127",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 211/211 [1:24:19<00:00, 23.98s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "113,879 hours of youtube_music\n",
      "12,287 hours of youtube_music_lyrics\n",
      "10,043 hours of youtube_music_lyrics_foreign\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 216/216 [1:19:03<00:00, 21.96s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "41,544 hours of genius_hq_lyrics\n",
      "17,307 hours of genius_hq_lyrics_foreign\n",
      "17,615 hours of genius_hq\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████| 51/51 [15:19<00:00, 18.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "410 hours of freesound\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████| 28/28 [11:28<00:00, 24.61s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19,514 hours of imslp\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 6/6 [02:36<00:00, 26.01s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3,726 hours of jamendo\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 207/207 [1:13:15<00:00, 21.23s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "62,117 hours of pond5_music\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████| 77/77 [28:07<00:00, 21.92s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10,175 hours of deezer_lyrics\n",
      "12,287 hours of deezer\n",
      "6,699 hours of deezer_lyrics_foreign\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 278/278 [1:48:34<00:00, 23.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "152,162 hours of ytm_tagged\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████| 2/2 [00:48<00:00, 24.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "103 hours of musescore_lyrics\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\", (1, 4302), 1, 1),\n",
    "    (\"freesound\", \"v1\", (2, 1021), 1, 1),\n",
    "    (\"imslp\", \"v1\", (1, 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=20,\n",
    "    chunksize=20,\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": 13,
   "id": "2ab57df8",
   "metadata": {},
   "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": 23,
   "id": "ba4ee8dc",
   "metadata": {},
   "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": 25,
   "id": "f5bdf8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !du -hs /mnt/data/georg/data/chirp_v2/data_tr.bin\n",
    "# # 1003G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "cbcbca32",
   "metadata": {},
   "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": 2,
   "id": "ffa828e8",
   "metadata": {},
   "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": 23,
   "id": "73286512",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_metas = read_jsonl(\"/mnt/data/georg/data/chirp_v2/metas_val.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "a3ed7922",
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = []\n",
    "for e in train_metas:\n",
    "    tags.extend(e.get(\"tags\", []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "2eed69a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "vc_tags = pd.Series(tags).value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "bf1c83e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pop                2495781\n",
       "Rap                2431682\n",
       "Rock               1193692\n",
       "commute            1089137\n",
       "energy boosters     940473\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vc_tags.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "f76f1e83",
   "metadata": {},
   "outputs": [],
   "source": [
    "vc_tags = vc_tags.to_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "6766efb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>lonely castle.</th>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Strings.</th>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>asianbeatz-riff-1.</th>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>dnb loops 001 only guitar wet 180 bpm.</th>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>yuung.</th>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                        count\n",
       "lonely castle.                            0.0\n",
       "Strings.                                  0.0\n",
       "asianbeatz-riff-1.                        0.0\n",
       "dnb loops 001 only guitar wet 180 bpm.    0.0\n",
       "yuung.                                    0.0"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vc_tags.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "fece8109",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = vc_tags[vc_tags[\"count\"]>=3][\"count\"].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52e1a982",
   "metadata": {},
   "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": 21,
   "id": "b6b9f19e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7854\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" So, so you think you can tell heaven from hell, blue skies from pain? Can you tell a green field from a cold steel rail? A smile from a viel? Do you think you can tell? Did they get you to trade your heroes for ghosts? Hot ashfor trees, hot air for a cool breeze? Cold comfort for change Did you exchange a walkon part in the war for a lead role in a cage? How I wish, how I wish you were here, we're just two lost souls swiming in a fishbowl, year after year. Running over the same old ground, what have we found? The same old fears, wish you were here.\""
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
