{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "561dc275", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:16:50.867366Z", "start_time": "2024-04-09T13:16:50.738907Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "a10-yul-dev-node-444\r\n" ] } ], "source": [ "!echo $HOSTNAME" ] }, { "cell_type": "code", "execution_count": 2, "id": "250e8fae", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:16:51.354467Z", "start_time": "2024-04-09T13:16:51.352407Z" } }, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"" ] }, { "cell_type": "code", "execution_count": 3, "id": "ed2662f0", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:16:52.075541Z", "start_time": "2024-04-09T13:16:52.073628Z" } }, "outputs": [], "source": [ "# vocab: \n", "# 0-60_000 text\n", "# 1x0-3999 semantic\n", "# 12x0-2047 coarse\n", "\n", "# 4000 semantic pad token\n", "# 4001 semantic infer token\n", "# 2048 coarse pad token\n", "# 2049 coarse infer 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_25_2x4k & dac_2c_25_12):\n", "# s3://suno-data/datasets/bundles/\n", "# v1/youtube_music\n", "# v1/genius_hq\n", "# v1/jamendo\n", "# v1/imslp\n", "# v2/pond5_music\n", "# v2/deezer\n", "# v2/ytm_tagged\n", "# v3/discogs" ] }, { "cell_type": "code", "execution_count": 4, "id": "9c4df11a", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:16:53.199730Z", "start_time": "2024-04-09T13:16:52.778338Z" } }, "outputs": [], "source": [ "%matplotlib inline\n", "from matplotlib import pyplot as plt" ] }, { "cell_type": "code", "execution_count": 6, "id": "6b35b54b", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:17:23.225705Z", "start_time": "2024-04-09T13:17:23.216792Z" } }, "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 = 2048\n", "COARSE_N_CODEBOOKS = 12\n", "COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE\n", "COARSE_INFER_TOKEN = COARSE_CODEBOOK_SIZE + 1\n", "COARSE_VOCAB_SIZE = 2112\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 = 8704\n", "N_TOKENS_TEXT = 2560\n", "N_TOKENS_MEMMAP = 6016 # max 240s of audio\n", "# make sure we have enough space for shift 10\n", "assert(\n", " BLOCK_SIZE >= (\n", " N_TOKENS_TEXT + N_TOKENS_MEMMAP + \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_12\"\n", "\n", "# TODO: hopefully fast enough for multicore write\n", "METAS_DIR = \"/app/suno/data/chirp_v4/metadata\"\n", "OUT_DATA_DIR = \"/app/suno/data/chirp_v4/base\"\n", "# OUT_DATA_DIR = \"/mnt/localdisk/data/chirp_v4_test\" # no ssh set up to cluster\n", "os.makedirs(OUT_DATA_DIR, exist_ok=True)" ] }, { "cell_type": "code", "execution_count": 12, "id": "b120ac9f", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:19:07.598595Z", "start_time": "2024-04-09T13:19:07.467260Z" } }, "outputs": [], "source": [ "!ls /app/suno/data/chirp_v4/metadata/" ] }, { "cell_type": "code", "execution_count": 13, "id": "8a4257e2", "metadata": { "ExecuteTime": { "end_time": "2024-04-09T13:25:13.151724Z", "start_time": "2024-04-09T13:20:25.118518Z" } }, "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_v5.jsonl\"))},\n", "# \"youtube_music\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"youtube_music.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", " \"discogs\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"discogs.jsonl\"))},\n", "}" ] }, { "cell_type": "code", "execution_count": 7, "id": "c92c1319", "metadata": {}, "outputs": [], "source": [ "# TODO: if has text segments, still use first 4mins segment with full text as well" ] }, { "cell_type": "code", "execution_count": 19, "id": "1db468de", "metadata": {}, "outputs": [], "source": [ "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", " # 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", " int(round(m[\"vocal_start_s\"] * SEMANTIC_RATE_HZ))\n", " if m[\"vocal_start_s\"] is not None\n", " else None,\n", " min(\n", " len(semantic_arr),\n", " int(round(m[\"vocal_end_s\"] * SEMANTIC_RATE_HZ)),\n", " )\n", " if m[\"vocal_end_s\"] is not None\n", " else None,\n", " )\n", " )\n", " # if \"text\" in segmented text then add twice\n", " if \"text\" in meta_info:\n", " segments_info.append((\n", " 0, \n", " min(len(semantic_arr), N_TOKENS_MEMMAP),\n", " meta_info.get(\"text\") if n == 0 else None, # add text only to first piece\n", " None,\n", " None,\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", " # randomize if we don't have lyrics\n", " offs = random.randint(10*SEMANTIC_RATE_HZ, N_TOKENS_MEMMAP-1)\n", " segments_info.append((0, min(len(semantic_arr), offs), None, None, None))\n", " for n in range(int(np.ceil((len(semantic_arr)-offs)/N_TOKENS_MEMMAP))):\n", " start_idx = offs + n * N_TOKENS_MEMMAP\n", " end_idx = min(len(semantic_arr), offs+(n+1)*N_TOKENS_MEMMAP)\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", " None,\n", " None,\n", " ))\n", "\n", " arr_list = []\n", " for (\n", " sem_start_idx,\n", " sem_end_idx,\n", " text,\n", " vocal_start_idx,\n", " vocal_end_idx,\n", " ) in segments_info:\n", " if (\n", " sem_end_idx - sem_start_idx > N_TOKENS_MEMMAP\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_MEMMAP:\n", " arr_c = np.pad(\n", " arr_c,\n", " ((0, N_TOKENS_MEMMAP - 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_MEMMAP - 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_MEMMAP, 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", " \"vocal_start_s\": round(vocal_start_idx / SEMANTIC_RATE_HZ, 2)\n", " if vocal_start_idx is not None\n", " else None,\n", " \"vocal_end_s\": round(vocal_end_idx / SEMANTIC_RATE_HZ, 2)\n", " if vocal_end_idx is not None\n", " else None,\n", " }\n", " if text is not None:\n", " new_meta[\"text\"] = text\n", " new_meta[\"text_lang\"] = meta_info[\"lang\"]\n", " new_meta[\"dset_suffix\"] = (\n", " \"lyrics\" if meta_info[\"lang\"] == \"en\" else \"lyrics_foreign\"\n", " )\n", " if \"tags\" in meta_info:\n", " new_meta[\"tags\"] = meta_info[\"tags\"]\n", " if \"original_id\" in meta_info:\n", " new_meta[\"original_id\"] = meta_info[\"original_id\"]\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", "# print(len(relevant_metas))\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", " print(f\"missing {s3_semantic_archive_filepath}\")\n", " continue\n", " try:\n", " archive = {k: v for k, v in read_from_s3(s3_semantic_archive_filepath, read_f=np.load).items()}\n", " except:\n", " # corrupt archive\n", " print(f\"corrupt {s3_semantic_archive_filepath}\")\n", " continue\n", " for k, v in archive.items():\n", " semantic_archive[k] = v\n", " \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", " print(f\"missing {s3_coarse_archive_filepath}\")\n", " continue\n", " try:\n", " archive = {k: v for k, v in read_from_s3(s3_coarse_archive_filepath, read_f=np.load).items()}\n", " except:\n", " # corrupt archive\n", " print(f\"corrupt {s3_coarse_archive_filepath}\")\n", " continue\n", " for k, v in archive.items():\n", " coarse_archive[k] = v\n", "\n", " semantic_uids, coarse_uids = set(semantic_archive.keys()), set(\n", " coarse_archive.keys()\n", " )\n", " # removing this for now just incase (eg imslp)\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), \"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", " \"vocal_start_s\": round(arr_meta[\"vocal_start_s\"], 2)\n", " if arr_meta[\"vocal_start_s\"] is not None\n", " else None,\n", " \"vocal_end_s\": round(arr_meta[\"vocal_end_s\"], 2)\n", " if arr_meta[\"vocal_end_s\"] is not None\n", " else None,\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 \"text\" in arr_meta:\n", " add_meta[\"text\"] = arr_meta[\"text\"]\n", " if \"text_lang\" in arr_meta:\n", " add_meta[\"text_lang\"] = arr_meta[\"text_lang\"]\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": 20, "id": "2da8574b", "metadata": {}, "outputs": [], "source": [ "NJOBS = 32\n", "CHUNKSIZE = 32" ] }, { "cell_type": "code", "execution_count": 21, "id": "aa8c2cda", "metadata": {}, "outputs": [], "source": [ "# TODO: save n-audio-tokens in entry\n", "# TODO: batch data!!!" ] }, { "cell_type": "code", "execution_count": null, "id": "d726c704", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "start prepare data\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/1 [00:00= 0)\n", "assert(mm[:100,:,0].max() <= SEMANTIC_CODEBOOK_SIZE)\n", "assert(mm[:100,:,1:].min() >= 0)\n", "assert(mm[:100,:,1:].max() <= COARSE_CODEBOOK_SIZE)" ] }, { "cell_type": "code", "execution_count": null, "id": "846797cd", "metadata": {}, "outputs": [], "source": [ "# randomly listen to some stuff\n", "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", "\n", "_ = preload_codec_models(\"/app/suno/tony/v3/dac_2c_25x12.pt\")\n", "mm = np.memmap(os.path.join(OUT_DATA_DIR, \"data_val.bin\"), dtype=np.uint16, mode=\"r\")\n", "mm = mm.reshape(-1, N_TOKENS_MEMMAP, SEMANTIC_N_CODEBOOKS+COARSE_N_CODEBOOKS)\n", "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, \"metas_val.jsonl\"))\n", "test_info = read_json(os.path.join(OUT_DATA_DIR, \"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": "da0c7e34", "metadata": {}, "outputs": [], "source": [ "# show text and audio\n", "# idx = 171\n", "# idx = random.choice(idx_list)\n", "idx = random.choice(test_info[\"genius_hq_lyrics\"][\"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(\"dataset:\", test_metas[idx].get(\"dataset\"))\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 == COARSE_PAD_TOKEN)[0]\n", "if len(pad_idx_arr) > 0:\n", " arr = arr[: pad_idx_arr[0], :]\n", "a = codec_decode(arr)\n", "a.play(compress=False)\n", "print(\"text:\", test_metas[idx].get(\"text\"))\n", "plt.plot(a.array_float[0])" ] }, { "cell_type": "code", "execution_count": null, "id": "a2e60f8c", "metadata": {}, "outputs": [], "source": [ "test_metas[idx]" ] }, { "cell_type": "code", "execution_count": null, "id": "11715eb2", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "77cb7994", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "fd2297b7", "metadata": {}, "source": [ "## make smaller version" ] }, { "cell_type": "code", "execution_count": 26, "id": "92c8ea85", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "base base_small base_test base_tiny\tmetadata tokenizer_60k.json\r\n" ] } ], "source": [ "!ls /mnt/round-surf/data/chirp_v4_test/" ] }, { "cell_type": "code", "execution_count": 11, "id": "85dcfa84", "metadata": {}, "outputs": [], "source": [ "!mkdir -p /mnt/round-surf/data/chirp_v4_test/base_tiny" ] }, { "cell_type": "code", "execution_count": 12, "id": "9d8e56a3", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "total 7.9T\r\n", "drwxrwxr-x 2 georg georg 7 Mar 25 21:18 .\r\n", "drwxrwxr-x 6 georg georg 5 Apr 7 17:55 ..\r\n", "-rw-rw-r-- 1 georg georg 7.8T Mar 25 11:46 data_tr.bin\r\n", "-rw-rw-r-- 1 georg georg 766M Mar 23 22:38 data_val.bin\r\n", "-rw-rw-r-- 1 georg georg 512M Mar 25 11:51 info_tr.json\r\n", "-rw-rw-r-- 1 georg georg 30K Mar 23 22:38 info_val.json\r\n", "-rw-rw-r-- 1 georg georg 14G Mar 25 11:46 metas_tr.jsonl\r\n", "-rw-rw-r-- 1 georg georg 3.4M Mar 23 22:38 metas_val.jsonl\r\n", "-rw-r--r-- 1 georg georg 1.3M Mar 25 21:18 tokenizer_60k.json\r\n" ] } ], "source": [ "!ls -lah /mnt/round-surf/data/chirp_v4_test/base" ] }, { "cell_type": "code", "execution_count": 13, "id": "0fb02d9c", "metadata": {}, "outputs": [], "source": [ "!cp /mnt/round-surf/data/chirp_v4_test/base/tokenizer_60k.json /mnt/round-surf/data/chirp_v4_test/base_tiny/\n", "!cp /mnt/round-surf/data/chirp_v4_test/base/*val* /mnt/round-surf/data/chirp_v4_test/base_tiny/" ] }, { "cell_type": "code", "execution_count": 14, "id": "886c359b", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "total 770M\r\n", "drwxrwxr-x 2 georg georg 4 Apr 7 17:56 .\r\n", "drwxrwxr-x 6 georg georg 5 Apr 7 17:55 ..\r\n", "-rw-rw-r-- 1 georg georg 766M Apr 7 17:56 data_val.bin\r\n", "-rw-rw-r-- 1 georg georg 30K Apr 7 17:56 info_val.json\r\n", "-rw-rw-r-- 1 georg georg 3.4M Apr 7 17:56 metas_val.jsonl\r\n", "-rw-r--r-- 1 georg georg 1.3M Apr 7 17:56 tokenizer_60k.json\r\n" ] } ], "source": [ "!ls -lah /mnt/round-surf/data/chirp_v4_test/base_tiny" ] }, { "cell_type": "code", "execution_count": 15, "id": "d4632d44", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "from suno_utils.utils.text import read_jsonl, write_jsonl\n", "import json\n", "import tqdm" ] }, { "cell_type": "code", "execution_count": 16, "id": "b775951c", "metadata": {}, "outputs": [], "source": [ "metas = read_jsonl(\"/mnt/round-surf/data/chirp_v4_test/base/metas_tr.jsonl\")" ] }, { "cell_type": "code", "execution_count": 17, "id": "a65b4816", "metadata": {}, "outputs": [], "source": [ "with open(\"/mnt/round-surf/data/chirp_v4_test/base/info_tr.json\") as f:\n", " info = json.load(f)" ] }, { "cell_type": "code", "execution_count": 18, "id": "d57a905f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "42799402\n", "54775032\n" ] } ], "source": [ "print(len(info[\"ytm_mb\"][\"idx_list\"]))\n", "print(info[\"ytm_mb\"][\"idx_list\"][-1])" ] }, { "cell_type": "code", "execution_count": 19, "id": "39fcfc57", "metadata": {}, "outputs": [], "source": [ "n_keep = 30_000_000 # lines to keep" ] }, { "cell_type": "code", "execution_count": 20, "id": "f801fe76", "metadata": {}, "outputs": [], "source": [ "info[\"ytm_mb\"][\"idx_list\"] = info[\"ytm_mb\"][\"idx_list\"][:-(info[\"ytm_mb\"][\"idx_list\"][-1]-n_keep+1)]" ] }, { "cell_type": "code", "execution_count": 21, "id": "dbea5195", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "18024369\n", "29999999\n" ] } ], "source": [ "print(len(info[\"ytm_mb\"][\"idx_list\"]))\n", "assert(info[\"ytm_mb\"][\"idx_list\"][-1] == n_keep - 1)\n", "print(info[\"ytm_mb\"][\"idx_list\"][-1])" ] }, { "cell_type": "code", "execution_count": 22, "id": "838b4870", "metadata": {}, "outputs": [], "source": [ "with open(\"/mnt/round-surf/data/chirp_v4_test/base_tiny/info_tr.json\", \"w\") as f:\n", " json.dump(info, f)" ] }, { "cell_type": "code", "execution_count": null, "id": "33b1bba4", "metadata": {}, "outputs": [], "source": [ "write_jsonl(metas[:n_keep], \"/mnt/round-surf/data/chirp_v4_test/base_tiny/metas_tr.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "b1315d10", "metadata": { "scrolled": true }, "outputs": [], "source": [ "mm = np.memmap(\"/mnt/round-surf/data/chirp_v4_test/base/data_tr.bin\", dtype=np.uint16, mode=\"r\")\n", "mm = mm.reshape(-1, 6016, 13)\n", "len(mm)" ] }, { "cell_type": "code", "execution_count": 25, "id": "2ee6fdf5", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|████████████████████████████████████████████| 3000/3000 [4:18:57<00:00, 5.18s/it]\n" ] } ], "source": [ "# runtime ~4h for 30mm lines\n", "new_shape = (n_keep, 6016, 13)\n", "new_mm = np.memmap(\n", " \"/mnt/round-surf/data/chirp_v4_test/base_tiny/data_tr.bin\", dtype=np.uint16, mode=\"w+\", shape=new_shape\n", ")\n", "\n", "# Define the chunk size (number of rows per chunk)\n", "chunk_size = 10000 # Adjust this based on your memory constraints\n", "\n", "# Iterate over the original memmap in chunks and write to the new memmap\n", "for i in tqdm.tqdm(range(0, n_keep, chunk_size)):\n", " end = min(i + chunk_size, n_keep) # Ensure we don't go beyond 700k rows\n", " new_mm[i:end] = mm[i:end]\n", "\n", " # Flush changes to disk\n", " new_mm.flush()" ] }, { "cell_type": "code", "execution_count": null, "id": "a9154804", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "f8f94b90", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "954da7dd", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "c621c6c6", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "b836e1f2", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "5bc835c2", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "1148c0b9", "metadata": {}, "source": [ "### transfer" ] }, { "cell_type": "code", "execution_count": null, "id": "78149f3a", "metadata": {}, "outputs": [], "source": [ "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"mkdir -p /mnt/localdisk/data/test_base/\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"cp /app/suno/data/tokenizer_60k.json /mnt/localdisk/data/test_base\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"cp -r /app/suno/data/chirp_v4_test /mnt/localdisk/data/test_base\"" ] }, { "cell_type": "code", "execution_count": null, "id": "d2a7dfed", "metadata": {}, "outputs": [], "source": [ "## OLD: via network (doesnt work cross cluster)\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"mkdir -p /mnt/localdisk/data/test_base/\"\n", "\n", "# #\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/tokenizer_60k.json /mnt/localdisk/data/test_base/\"\n", "\n", "# #\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/data_val.bin /mnt/localdisk/data/test_base/\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/metas_val.jsonl /mnt/localdisk/data/test_base/\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/info_val.json /mnt/localdisk/data/test_base/\"\n", "\n", "# #\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/data_tr.bin /mnt/localdisk/data/test_base/\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/metas_tr.jsonl /mnt/localdisk/data/test_base/\"\n", "\n", "# pdsh -w compute-hpc-node-\\[42,57,88,93,136,188-189,294,311,361,381,479,527,549,573,652,768,18,126,149,237,243,290,308,553,622,694,775,788,905,951,973\\] \\\n", "# \"scp a10-yul-dev-node-960:/mnt/localdisk/data/chirp_v4_test/info_tr.json /mnt/localdisk/data/test_base/\"\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6100b384", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "a5ff0db4", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "20e1e6fa", "metadata": {}, "source": [ "## Playground" ] }, { "cell_type": "code", "execution_count": null, "id": "48d639b1", "metadata": {}, "outputs": [], "source": [ "# # TODO: check\n", "# /home/tony/Work/glockenspiel/sunoGPT/scripts/\n", "# from data_preparation_7b import *" ] }, { "cell_type": "code", "execution_count": null, "id": "f802a4fa", "metadata": {}, "outputs": [], "source": [ "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/genius_hq.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/youtube_music.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/jamendo.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/imslp.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/pond5_music.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/deezer.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/ytm_tagged.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/musescore.jsonl /app/suno/data/chirp_v3/metadata/\n", "# !aws s3 cp s3://suno-data/georg/data/chirp_v2_5/filtered_metas/ytm_mb.jsonl /app/suno/data/chirp_v3/metadata/" ] }, { "cell_type": "code", "execution_count": null, "id": "70b42c83", "metadata": {}, "outputs": [], "source": [ "# load manifests of IDs and text and tags etc\n", "# loading everything can take a while...4 mins\n", "meta_info_map = {\n", " \"genius_hq\": {m[\"id\"]: m for m in read_jsonl(os.path.join(\"/home/tony/Work/tony/FineTuning_chirp_v3_4min/metadata/\", \"genius_hq_v5.jsonl\"))},\n", " \"youtube_music\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"youtube_music.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", " # \"ytm_mb\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"ytm_mb.jsonl\"))}, # this is the large one in case we load full\n", "}" ] }, { "cell_type": "code", "execution_count": null, "id": "702768d1", "metadata": {}, "outputs": [], "source": [ "with open(\"/app/suno/data/preprocessing/meta_cutoff_freq.json\", \"r\") as fp:\n", " meta_cutoff_freq = json.load(fp)" ] }, { "cell_type": "code", "execution_count": null, "id": "496f0b7b", "metadata": {}, "outputs": [], "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", " (\"jamendo\", \"v1\", (0, 1), 1, 1),\n", " (\"imslp\", \"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", " (\"ytm_mb\", \"v2\", (0, 1), 1, 1),\n", "]\n", "prep_data(\n", " datasets,\n", " out_data_dir=OUT_DATA_DIR,\n", " meta_info_map=meta_info_map,\n", " meta_cutoff_freq=meta_cutoff_freq, # TODO: This arg should be removed in the longer run...\n", " is_val=True,\n", " njobs=NJOBS,\n", " chunksize=CHUNKSIZE,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "d933d6be", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "45fee227", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "566174e5", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "c2cf27ef", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "f2e5534d", "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" }, "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 }