{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ffdbda4b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:08:21.448320Z",
     "start_time": "2023-10-16T18:08:21.443831Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd41c10c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:08:23.489683Z",
     "start_time": "2023-10-16T18:08:21.450916Z"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import json\n",
    "import numpy as np\n",
    "import tqdm\n",
    "import torch\n",
    "import funcy\n",
    "import time\n",
    "import gc\n",
    "from scipy.io import wavfile\n",
    "import tempfile\n",
    "import collections\n",
    "from collections import defaultdict\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "from suno_utils.utils.s3 import _apply_mp\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.data_loader import load_audio_mp\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl, write_json, read_json\n",
    "from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3\n",
    "from suno_utils.audio.conversion import convert_audio_files\n",
    "\n",
    "SAMPLE_RATE = 24_000\n",
    "EMBEDDING_RATE = 25\n",
    "N_CODEBOOKS = 8\n",
    "\n",
    "IN_DATA_DIR = \"/app/suno/data/mert_25hz_speech_short\"\n",
    "IN_AUDIO_DIR = os.path.join(IN_DATA_DIR, \"audio\")\n",
    "IN_TSV_DIR = os.path.join(IN_DATA_DIR, \"audio_tsv\")\n",
    "IN_LABEL_DIR = os.path.join(IN_DATA_DIR, \"label\")\n",
    "\n",
    "OUT_DATA_DIR = \"/app/suno/data/mert_25hz_short\"\n",
    "OUT_AUDIO_DIR = os.path.join(OUT_DATA_DIR, \"audio\")\n",
    "OUT_TSV_DIR = os.path.join(OUT_DATA_DIR, \"audio_tsv\")\n",
    "OUT_LABEL_DIR = os.path.join(OUT_DATA_DIR, \"label\")\n",
    "OUT_TEMP_DIR = os.path.join(OUT_DATA_DIR, \"temp\")\n",
    "\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "os.makedirs(OUT_AUDIO_DIR, exist_ok=True)\n",
    "os.makedirs(OUT_TSV_DIR, exist_ok=True)\n",
    "os.makedirs(OUT_LABEL_DIR, exist_ok=True)\n",
    "os.makedirs(OUT_TEMP_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "195c03f2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:08:23.832386Z",
     "start_time": "2023-10-16T18:08:23.830545Z"
    }
   },
   "outputs": [],
   "source": [
    "MIN_SIZE = SAMPLE_RATE * 5\n",
    "MAX_SIZE = SAMPLE_RATE * 12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c729d86d",
   "metadata": {},
   "source": [
    "# File loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "42fa9e8a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:21.056573Z",
     "start_time": "2023-10-16T18:24:21.053998Z"
    }
   },
   "outputs": [],
   "source": [
    "data_split = \"train\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "66631638",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:31.931633Z",
     "start_time": "2023-10-16T18:24:21.325001Z"
    }
   },
   "outputs": [],
   "source": [
    "# load current data\n",
    "tsv_info = []\n",
    "with open(os.path.join(OUT_TSV_DIR, f\"{data_split}.tsv\"), \"r\") as f:\n",
    "    for line in f.read().strip().split(\"\\n\"):\n",
    "        if len(line.strip()) == 0:\n",
    "            continue\n",
    "        tsv_info.append(line.strip().split(\"\\t\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cad52faa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:31.936490Z",
     "start_time": "2023-10-16T18:24:31.934168Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7085919\n"
     ]
    }
   ],
   "source": [
    "print(len(tsv_info))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "741d5c23",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:32.062809Z",
     "start_time": "2023-10-16T18:24:31.938316Z"
    }
   },
   "outputs": [],
   "source": [
    "tsv_data_path = tsv_info[0]\n",
    "tsv_data = tsv_info[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "98814c0c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:34.906543Z",
     "start_time": "2023-10-16T18:24:32.064941Z"
    }
   },
   "outputs": [],
   "source": [
    "speech_tsv_info = []\n",
    "with open(os.path.join(IN_TSV_DIR, f\"{data_split}.tsv\"), \"r\") as f:\n",
    "    for line in f.read().strip().split(\"\\n\"):\n",
    "        if len(line.strip()) == 0:\n",
    "            continue\n",
    "        speech_tsv_info.append(line.strip().split(\"\\t\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1feab3db",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:34.928711Z",
     "start_time": "2023-10-16T18:24:34.908590Z"
    }
   },
   "outputs": [],
   "source": [
    "speech_tsv_data_path = speech_tsv_info[0]\n",
    "speech_tsv_data = speech_tsv_info[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c0a194fc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:34.977791Z",
     "start_time": "2023-10-16T18:24:34.930760Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1479973\n"
     ]
    }
   ],
   "source": [
    "print(len(speech_tsv_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "bc441e6a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:35.119024Z",
     "start_time": "2023-10-16T18:24:34.979454Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8565891\n"
     ]
    }
   ],
   "source": [
    "new_tsv_data = tsv_data + speech_tsv_data\n",
    "print(len(new_tsv_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "e1e94768",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:24:36.903723Z",
     "start_time": "2023-10-16T18:24:35.120852Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(\n",
    "    os.path.join(OUT_DATA_DIR, \"audio_speech_tsv\", f\"{data_split}.tsv\"), \"w\"\n",
    ") as f:\n",
    "    f.write(OUT_AUDIO_DIR + \"\\n\")\n",
    "    for fn in new_tsv_data:\n",
    "        f.write(\"\\t\".join(fn)+ \"\\n\") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4defb74e",
   "metadata": {},
   "source": [
    "## Slice the codebook labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "a0b384f7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-16T18:29:56.680192Z",
     "start_time": "2023-10-16T18:24:36.905738Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "codebook: 0\n",
      "codebook: 1\n",
      "codebook: 2\n",
      "codebook: 3\n",
      "codebook: 4\n",
      "codebook: 5\n",
      "codebook: 6\n",
      "codebook: 7\n"
     ]
    }
   ],
   "source": [
    "for codebook in range(N_CODEBOOKS):\n",
    "    print(\"codebook:\", codebook)\n",
    "    with open(os.path.join(OUT_LABEL_DIR, f\"{data_split}.codec_{codebook}.npy\"), \"rb\") as f_out:\n",
    "        original_array = np.load(f_out)\n",
    "    with open(os.path.join(IN_LABEL_DIR, f\"{data_split}.codec_{codebook}.npy\"), \"rb\") as f_out:\n",
    "        speech_array = np.load(f_out)\n",
    "    # print(original_array.shape, speech_array.shape)\n",
    "    output_array = np.vstack([original_array, speech_array])\n",
    "    with open(os.path.join(OUT_DATA_DIR, \"speech_label\", f\"{data_split}.codec_{codebook}.npy\"), \"wb\") as f_out:\n",
    "        np.save(f_out, output_array)\n",
    "    # print(output_array.shape)\n",
    "    assert output_array.shape[0] == len(new_tsv_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40c20906",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3880618e",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a48e0ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    " OMP_NUM_THREADS=6 python -u /home/tony/Work/glockenspiel/mert_training/src/fairseq/fairseq_cli/hydra_train.py \\\n",
    "    --config-dir '/home/tony/Work/glockenspiel/mert_training/mert_fairseq/config/pretrain' \\\n",
    "    --config-name 'custom_95M' \\\n",
    "    common.user_dir='/home/tony/Work/glockenspiel/mert_training/mert_fairseq' \\\n",
    "    common.wandb_project='mert_test' \\\n",
    "    checkpoint.save_dir='/app/suno/checkpoints/mert_mus_speech_test_valid' \\\n",
    "    distributed_training.distributed_rank=0 \\\n",
    "    distributed_training.distributed_world_size=4  \\\n",
    "    distributed_training.nprocs_per_node=4 \\\n",
    "    optimization.update_freq='[4]' \\\n",
    "    distributed_training.distributed_init_method=\"tcp://127.0.0.2:39683\" \\\n",
    "    task.data='/app/suno/data/mert_25hz_short/audio_speech_tsv' \\\n",
    "    task.label_dir='/app/suno/data/mert_25hz_short/speech_label' \\\n",
    "    task.labels='[\"codec_0\", \"codec_1\"]' \\\n",
    "    task.do_valid_set_test=true \\\n",
    "    dataset.num_workers=6 \\\n",
    "    dataset.max_tokens=1800000 \\\n",
    "    dataset.disable_validation=true \\\n",
    "    model.label_rate=25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0194f166",
   "metadata": {},
   "outputs": [],
   "source": [
    "# on A100\n",
    "OMP_NUM_THREADS=6 python -u /home/tony/Work/glockenspiel/mert_training/src/fairseq/fairseq_cli/hydra_train.py \\\n",
    "    --config-dir '/home/tony/Work/glockenspiel/mert_training/mert_fairseq/config/pretrain' \\\n",
    "    --config-name 'custom_95M' \\\n",
    "    common.user_dir='/home/tony/Work/glockenspiel/mert_training/mert_fairseq' \\\n",
    "    common.wandb_project='mert_test' \\\n",
    "    checkpoint.save_dir='/app/suno/checkpoints/mert_mus_speech_test' \\\n",
    "    distributed_training.distributed_rank=0 \\\n",
    "    distributed_training.distributed_world_size=8  \\\n",
    "    distributed_training.nprocs_per_node=8 \\\n",
    "    optimization.update_freq='[1]' \\\n",
    "    distributed_training.distributed_init_method=\"tcp://127.0.0.1:39683\" \\\n",
    "    task.data='/app/suno/data/mert_25hz_short/audio_speech_tsv' \\\n",
    "    task.label_dir='/app/suno/data/mert_25hz_short/speech_label' \\\n",
    "    task.labels='[\"codec_0\", \"codec_1\"]' \\\n",
    "    dataset.num_workers=6 \\\n",
    "    dataset.max_tokens=4000000 \\\n",
    "    dataset.disable_validation=false \\\n",
    "    model.label_rate=25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65299565",
   "metadata": {},
   "outputs": [],
   "source": [
    "OMP_NUM_THREADS=6 python -u /home/tony/Work/glockenspiel/mert_training/src/fairseq/fairseq_cli/hydra_train.py \\\n",
    "    --config-dir '/home/tony/Work/glockenspiel/mert_training/mert_fairseq/config/pretrain' \\\n",
    "    --config-name 'custom_95M_2' \\\n",
    "    common.user_dir='/home/tony/Work/glockenspiel/mert_training/mert_fairseq' \\\n",
    "    common.wandb_project='mert_test' \\\n",
    "    checkpoint.save_dir='/app/suno/checkpoints/mert_mus_speech_test_nocqt_4code' \\\n",
    "    distributed_training.distributed_rank=0 \\\n",
    "    distributed_training.distributed_world_size=8  \\\n",
    "    distributed_training.nprocs_per_node=8 \\\n",
    "    optimization.update_freq='[1]' \\\n",
    "    distributed_training.distributed_init_method=\"tcp://127.0.0.1:39683\" \\\n",
    "    task.data='/app/suno/data/mert_25hz_short/audio_speech_tsv' \\\n",
    "    task.label_dir='/app/suno/data/mert_25hz_short/speech_label' \\\n",
    "    task.labels='[\"codec_0\", \"codec_1\", \"codec_2\", \"codec_3\"]' \\\n",
    "    dataset.num_workers=6 \\\n",
    "    dataset.max_tokens=4000000 \\\n",
    "    dataset.disable_validation=false \\\n",
    "    model.label_rate=25"
   ]
  }
 ],
 "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
}