{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "08218f19", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T18:57:59.185051Z", "start_time": "2023-12-16T18:57:59.183176Z" } }, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"" ] }, { "cell_type": "code", "execution_count": 2, "id": "76fc0b2c", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T18:57:59.821884Z", "start_time": "2023-12-16T18:57:59.820081Z" } }, "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-16T18:58:01.073063Z", "start_time": "2023-12-16T18:58:00.570736Z" } }, "outputs": [], "source": [ "%matplotlib inline\n", "from matplotlib import pyplot as plt" ] }, { "cell_type": "code", "execution_count": 4, "id": "5164dde2", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T18:58:01.988944Z", "start_time": "2023-12-16T18:58:01.987199Z" } }, "outputs": [], "source": [ "# TODO: add original_filepath on s3" ] }, { "cell_type": "code", "execution_count": 5, "id": "0e33ff3d", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T18:58:11.209615Z", "start_time": "2023-12-16T18:58:08.470696Z" } }, "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 (\n", " write_jsonl,\n", " read_jsonl,\n", " write_json,\n", " read_json,\n", " normalize_whitespace,\n", ")\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 BLOCK_SIZE >= (\n", " N_TOKENS_TEXT\n", " + N_TOKENS_AUDIO\n", " + SEMANTIC_N_CODEBOOKS * SEMANTIC_SHIFT_FACTOR\n", " + (COARSE_N_CODEBOOKS - 1) * COARSE_SHIFT_FACTOR\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_v13_full\"\n", "os.makedirs(OUT_DATA_DIR, exist_ok=True)" ] }, { "cell_type": "code", "execution_count": 6, "id": "b30bfb77", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T19:00:55.937315Z", "start_time": "2023-12-16T18:58:26.922901Z" } }, "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\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"genius_hq_v13.jsonl\"))\n", " },\n", " \"youtube_music\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"youtube_music.jsonl\"))\n", " },\n", " \"freesound\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"freesound.jsonl\"))\n", " },\n", " \"jamendo\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"jamendo.jsonl\"))\n", " },\n", " \"imslp\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"imslp.jsonl\"))},\n", " \"pond5_music\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"pond5_music_v12.jsonl\"))\n", " },\n", " \"deezer\": {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"deezer.jsonl\"))},\n", " \"ytm_tagged\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"ytm_tagged.jsonl\"))\n", " },\n", " \"musescore\": {\n", " m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"musescore.jsonl\"))\n", " },\n", "}" ] }, { "cell_type": "code", "execution_count": 7, "id": "9c043b02", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T19:08:40.007506Z", "start_time": "2023-12-16T19:08:39.982969Z" } }, "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", " 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", " (\n", " start_idx,\n", " end_idx,\n", " meta_info.get(\"text\")\n", " if n == 0\n", " else None, # 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", " 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[\"lang\"]\n", " new_meta[\"text_aligned\"] = is_aligned\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 \"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": 9, "id": "f862095a", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T19:08:55.128953Z", "start_time": "2023-12-16T19:08:55.126820Z" } }, "outputs": [], "source": [ "NJOBS = 60\n", "CHUNKSIZE = 60" ] }, { "cell_type": "code", "execution_count": 10, "id": "3339fe31", "metadata": { "ExecuteTime": { "end_time": "2023-12-16T19:13:56.799253Z", "start_time": "2023-12-16T19:08:57.500051Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "start prepare data\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:35<00:00, 35.65s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "32 hours of youtube_music\n", "0 hours of youtube_music_lyrics_foreign\n", "0 hours of youtube_music_lyrics\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:35<00:00, 35.68s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "7 hours of genius_hq_lyrics_foreign\n", "9 hours of genius_hq_lyrics\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:24<00:00, 24.67s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "1 hours of freesound\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:36<00:00, 36.00s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "29 hours of imslp\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:42<00:00, 42.03s/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.69s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "21 hours of pond5_music\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:35<00:00, 35.65s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "7 hours of deezer_lyrics\n", "12 hours of deezer\n", "8 hours of deezer_lyrics_foreign\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:30<00:00, 30.72s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "33 hours of ytm_tagged\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:28<00:00, 28.07s/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, 10), 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": 11, "id": "a2dfa127", "metadata": { "ExecuteTime": { "end_time": "2023-12-17T01:03:46.634954Z", "start_time": "2023-12-16T19:13:56.800704Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "start prepare data\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [1:23:42<00:00, 70.74s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "135,011 hours of youtube_music\n", "557 hours of youtube_music_lyrics_foreign\n", "641 hours of youtube_music_lyrics\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 72/72 [1:10:28<00:00, 58.73s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "42,815 hours of genius_hq_lyrics\n", "23,157 hours of genius_hq_lyrics_foreign\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 17/17 [07:53<00:00, 27.85s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "410 hours of freesound\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [11:49<00:00, 70.96s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "19,514 hours of imslp\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [03:05<00:00, 92.91s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "3,726 hours of jamendo\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 69/69 [44:38<00:00, 38.82s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "8,729 hours of pond5_music\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 26/26 [24:48<00:00, 57.26s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "10,175 hours of deezer_lyrics\n", "12,375 hours of deezer\n", "6,787 hours of deezer_lyrics_foreign\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 93/93 [1:40:56<00:00, 65.13s/it]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "152,162 hours of ytm_tagged\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:31<00:00, 32.00s/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", "# don't change the end_idx....X.x\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\", (10, 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": 12, "id": "c572bd6f", "metadata": { "ExecuteTime": { "end_time": "2023-12-17T01:03:48.015991Z", "start_time": "2023-12-17T01:03:46.636567Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "896G\t/app/suno/data/chirp_v2_finetune_v13_full/data_tr.bin\r\n" ] } ], "source": [ "!du -hs /app/suno/data/chirp_v2_finetune_v13_full/data_tr.bin\n", "# 1003G" ] }, { "cell_type": "code", "execution_count": 13, "id": "999206cd", "metadata": { "ExecuteTime": { "end_time": "2023-12-17T01:03:48.020692Z", "start_time": "2023-12-17T01:03:48.018144Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "DONE!\n" ] } ], "source": [ "print(\"DONE!\")" ] }, { "cell_type": "code", "execution_count": 19, "id": "2ae23b40", "metadata": { "ExecuteTime": { "end_time": "2023-12-17T18:56:57.352807Z", "start_time": "2023-12-17T18:56:55.910310Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "youtube_music\n", "youtube_music_lyrics_foreign\n", "youtube_music_lyrics\n", "genius_hq_lyrics\n", "genius_hq_lyrics_foreign\n", "freesound\n", "imslp\n", "jamendo\n", "pond5_music\n", "deezer_lyrics\n", "deezer\n", "deezer_lyrics_foreign\n", "ytm_tagged\n", "musescore_lyrics\n", "17768435\n", "92543.93229166667\n" ] } ], "source": [ "with open(\"/app/suno/data/chirp_v2_finetune_v13_full/info_tr.json\", \"r\") as fp:\n", " info_tr = json.load(fp)\n", "n_total = 0\n", "for k in info_tr:\n", " print(k)\n", " n_total += len(info_tr[k]['idx_list'])\n", "print(n_total)\n", "print(n_total / 4 / 8 / 6)" ] }, { "cell_type": "code", "execution_count": 16, "id": "eb9bed5b", "metadata": { "ExecuteTime": { "end_time": "2023-12-17T18:42:29.405825Z", "start_time": "2023-12-17T18:42:29.390856Z" } }, "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[16], 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": "code", "execution_count": null, "id": "4ad9d1f9", "metadata": { "ExecuteTime": { "end_time": "2023-12-13T06:16:08.119356Z", "start_time": "2023-12-13T06:16:08.119347Z" } }, "outputs": [], "source": [ "print(\"here\")" ] }, { "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-13T06:16:08.120278Z", "start_time": "2023-12-13T06:16:08.120269Z" } }, "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-13T06:16:08.121033Z", "start_time": "2023-12-13T06:16:08.121024Z" } }, "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-13T06:16:08.121872Z", "start_time": "2023-12-13T06:16:08.121863Z" } }, "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-13T06:16:08.122577Z", "start_time": "2023-12-13T06:16:08.122568Z" } }, "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-13T06:16:08.123326Z", "start_time": "2023-12-13T06:16:08.123317Z" } }, "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-13T06:16:08.124047Z", "start_time": "2023-12-13T06:16:08.124039Z" } }, "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-13T06:16:08.124773Z", "start_time": "2023-12-13T06:16:08.124764Z" } }, "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-13T06:16:08.125508Z", "start_time": "2023-12-13T06:16:08.125499Z" } }, "outputs": [], "source": [ "vc_tags = pd.Series(tags).value_counts()" ] }, { "cell_type": "code", "execution_count": null, "id": "bf1c83e8", "metadata": { "ExecuteTime": { "end_time": "2023-12-13T06:16:08.126090Z", "start_time": "2023-12-13T06:16:08.126082Z" } }, "outputs": [], "source": [ "vc_tags.head(5)" ] }, { "cell_type": "code", "execution_count": null, "id": "f76f1e83", "metadata": { "ExecuteTime": { "end_time": "2023-12-13T06:16:08.126777Z", "start_time": "2023-12-13T06:16:08.126769Z" } }, "outputs": [], "source": [ "vc_tags = vc_tags.to_frame()" ] }, { "cell_type": "code", "execution_count": null, "id": "6766efb1", "metadata": { "ExecuteTime": { "end_time": "2023-12-13T06:16:08.127662Z", "start_time": "2023-12-13T06:16:08.127654Z" } }, "outputs": [], "source": [ "vc_tags.tail()" ] }, { "cell_type": "code", "execution_count": null, "id": "fece8109", "metadata": { "ExecuteTime": { "end_time": "2023-12-13T06:16:08.128387Z", "start_time": "2023-12-13T06:16:08.128379Z" } }, "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-13T06:16:08.128967Z", "start_time": "2023-12-13T06:16:08.128959Z" } }, "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-13T06:16:08.129808Z", "start_time": "2023-12-13T06:16:08.129801Z" } }, "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": [ "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", " 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", " (\n", " start_idx,\n", " end_idx,\n", " meta_info.get(\"text\")\n", " if n == 0\n", " else None, # 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", " 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[\"lang\"]\n", " new_meta[\"text_aligned\"] = is_aligned\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 \"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", " 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(\n", " coarse_archive.keys()\n", " )\n", " # removing this for now just incase (eg imslp)\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 (\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", " # 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,\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", " 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", "\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)" ] } ], "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" }, "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 }