{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check /home/minz/llm_tagging/data/genius_llm_metas_r3_enhanced.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:32.757700Z",
     "start_time": "2024-05-16T20:55:31.651146Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1565046/1029180021.py:1: DeprecationWarning: \n",
      "Pyarrow will become a required dependency of pandas in the next major release of pandas (pandas 3.0),\n",
      "(to allow more performant data types, such as the Arrow string type, and better interoperability with other libraries)\n",
      "but was not found to be installed on your system.\n",
      "If this would cause problems for you,\n",
      "please provide us feedback at https://github.com/pandas-dev/pandas/issues/54466\n",
      "        \n",
      "  import pandas as pd\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "from suno_utils.utils.s3 import download_s3_files\n",
    "import sys\n",
    "from collections import defaultdict\n",
    "from suno_utils.utils.text import (\n",
    "    write_jsonl,\n",
    "    read_jsonl,\n",
    "    write_json,\n",
    "    read_json,\n",
    ")\n",
    "import shutil\n",
    "import ast\n",
    "from preference_helper import *\n",
    "\n",
    "sys.path.insert(0, \"/home/tony/Work/glockenspiel/sunoGPT/scripts/\")\n",
    "\n",
    "from preference_data_preparation_4min import *\n",
    "import numpy as np\n",
    "\n",
    "pd.set_option(\"display.max_rows\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:32.793025Z",
     "start_time": "2024-05-16T20:55:32.759394Z"
    }
   },
   "outputs": [],
   "source": [
    "INPUT_DATA_DIR = \"/app/suno/data/dpo/top_mix\"\n",
    "OUT_DATA_DIR = \"/app/suno/data/dpo/13v_real_v1/\"\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "shutil.copyfile(\n",
    "    \"/app/suno/data/dpo/7v_v1_full/tokenizer_60k.json\",\n",
    "    os.path.join(OUT_DATA_DIR, \"tokenizer_60k.json\"),\n",
    ")\n",
    "NPZ_DIR = \"/app/suno/data/dpo/gen_cycle_npz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:32.913341Z",
     "start_time": "2024-05-16T20:55:32.794133Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "174090\n",
      "174090\n"
     ]
    }
   ],
   "source": [
    "converted_paths = os.listdir(NPZ_DIR)\n",
    "print(len(converted_paths))\n",
    "converted_paths = set([f.replace(\".npz\", \"\") for f in converted_paths])\n",
    "print(len(converted_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:32.916861Z",
     "start_time": "2024-05-16T20:55:32.915273Z"
    }
   },
   "outputs": [],
   "source": [
    "input_dataset = \"tr\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:37.839239Z",
     "start_time": "2024-05-16T20:55:32.917928Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "173936\n",
      "173936\n"
     ]
    }
   ],
   "source": [
    "with open(os.path.join(INPUT_DATA_DIR, f\"gen_{input_dataset}.json\"), \"r\") as fp:\n",
    "    val_inputs = json.load(fp)\n",
    "print(len(val_inputs))\n",
    "\n",
    "val_real_mm = np.memmap(\n",
    "    os.path.join(INPUT_DATA_DIR, f\"data_{input_dataset}.bin\"), dtype=np.uint16, mode=\"r\"\n",
    ")\n",
    "val_real_metas = read_jsonl(\n",
    "    os.path.join(INPUT_DATA_DIR, f\"metas_{input_dataset}.jsonl\")\n",
    ")\n",
    "val_real_info = read_json(os.path.join(INPUT_DATA_DIR, f\"info_{input_dataset}.json\"))\n",
    "val_real_mm = val_real_mm.reshape(-1, 3008, 13)\n",
    "print(len(val_real_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:37.842247Z",
     "start_time": "2024-05-16T20:55:37.840649Z"
    }
   },
   "outputs": [],
   "source": [
    "# val_inputs = val_inputs[:100000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:38.059698Z",
     "start_time": "2024-05-16T20:55:37.843344Z"
    }
   },
   "outputs": [],
   "source": [
    "gen_input_to_real_metas = {}\n",
    "for i, test_input in enumerate(val_inputs):\n",
    "    gen_input_to_real_metas[test_input[\"id\"]] = val_real_metas[i]\n",
    "    gen_input_to_real_metas[test_input[\"id\"]][\"lookup_index\"] = i"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Downloads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:38.139453Z",
     "start_time": "2024-05-16T20:55:38.061110Z"
    }
   },
   "outputs": [],
   "source": [
    "s3_ids = [test_input[\"id\"].replace(\"_gen_\", \"_gen_cycle_\") for test_input in val_inputs]\n",
    "s3_paths = [f\"s3://suno-data-uploads/studio/uploads/{s3_id}.npz\" for s3_id in s3_ids]\n",
    "local_paths = [f\"{NPZ_DIR}/{s3_id}.npz\" for s3_id in s3_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:38.420409Z",
     "start_time": "2024-05-16T20:55:38.140902Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jobs to be done 2 2\n"
     ]
    }
   ],
   "source": [
    "finished_paths = os.listdir(NPZ_DIR)\n",
    "finished_paths_set = set(finished_paths)\n",
    "unfinished_s3_paths = [\n",
    "    path for path in s3_paths if os.path.basename(path) not in finished_paths_set\n",
    "]\n",
    "unfinished_paths = [\n",
    "    path for path in local_paths if os.path.basename(path) not in finished_paths_set\n",
    "]\n",
    "print(\"jobs to be done\", len(unfinished_paths), len(unfinished_s3_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:39.241455Z",
     "start_time": "2024-05-16T20:55:38.421785Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.23it/s]\n"
     ]
    }
   ],
   "source": [
    "_ = download_s3_files(unfinished_s3_paths, unfinished_paths, n_cores=48)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:45.347324Z",
     "start_time": "2024-05-16T20:55:39.243423Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "174092\n",
      "174092\n"
     ]
    }
   ],
   "source": [
    "converted_paths = os.listdir(NPZ_DIR)\n",
    "print(len(converted_paths))\n",
    "converted_paths = set([f.replace(\".npz\", \"\") for f in converted_paths])\n",
    "print(len(converted_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:45.433719Z",
     "start_time": "2024-05-16T20:55:45.348661Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "173936\n"
     ]
    }
   ],
   "source": [
    "val_inputs = [\n",
    "    test_input\n",
    "    for test_input in val_inputs\n",
    "    if test_input[\"id\"].replace(\"_gen_\", \"_gen_cycle_\") in converted_paths\n",
    "]\n",
    "print(len(val_inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:45.438105Z",
     "start_time": "2024-05-16T20:55:45.436525Z"
    }
   },
   "outputs": [],
   "source": [
    "# BREAK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:45.498897Z",
     "start_time": "2024-05-16T20:55:45.439148Z"
    }
   },
   "outputs": [],
   "source": [
    "def reshift(arr):\n",
    "    sem_start_idx = 0\n",
    "    sem_end_idx = len(arr) - 1\n",
    "    semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS]\n",
    "    coarse_arr = arr[:, SEMANTIC_N_CODEBOOKS:]\n",
    "\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",
    "    assert not (sem_end_idx > len(semantic_arr) or coarse_end_idx > len(coarse_arr))\n",
    "\n",
    "    # get array segments\n",
    "    arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy()\n",
    "    arr_c = coarse_arr[coarse_start_idx:coarse_end_idx, :].copy()\n",
    "    assert arr_s.max() <= SEMANTIC_PAD_TOKEN\n",
    "    assert arr_c.max() <= COARSE_PAD_TOKEN\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",
    "    return arr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actually make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T20:55:45.559079Z",
     "start_time": "2024-05-16T20:55:45.500092Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'a7833ca4-3f8d-44b1-9cd6-17cb694343da_gen_0',\n",
       " 'prompt_text': 'Letra de \"Você Vai Entender\" por 1Kilo[Verso 1]\\nEu me esquivei de todas marteladas\\nE entre todas opções tô na mais arriscada\\nMeus irmãos na luta iluminando cada passo\\nE mesmo quando me sentia só, havia alguém ali\\nEu peço proteção a todos os meus\\nOuço reclamações e o mundo continua igual\\nA força que eu peço a ti\\nA calma que eu peço a ti\\nA benção que eu peço a ti\\nÓ Deus, eu sei que cê vai entender\\nA força que eu peço a ti\\nA calma que eu peço a ti\\nA benção que eu peço a ti\\nÓ Deus, eu sei que cê vai entender\\n[Refrão]\\nÓ Deus, que palhaçada tem aí pra mim?\\nQue desapego vou ter que aprender?\\nO que é que mais vou precisar passar?\\nQuem mais vou precisar esquecer?\\nEu tô Morgado, tô meio deprê\\nE se te ofende quando eu fervo se eu escuto um papo errado\\nÉ que tá foda de resolver tudo que tem rolado\\nManda o teu papo e se adianta, bro',\n",
       " 'metadata': {'tags': 'Brasil, Em Português, Rap'},\n",
       " 'model_name': 'chirp-v3-engine-v0',\n",
       " 'title': 'a7833ca4-3f8d-44b1-9cd6-17cb694343da_gen_0'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_inputs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.556Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█████████████████▌                                                                                     | 29607/173936 [13:16<1:06:12, 36.33it/s]"
     ]
    }
   ],
   "source": [
    "dset_type = input_dataset\n",
    "out_mmap_path = os.path.join(OUT_DATA_DIR, f\"data_{dset_type}.bin\")\n",
    "out_metas_path = os.path.join(OUT_DATA_DIR, f\"meta_{dset_type}.jsonl\")\n",
    "out_info_filepath = os.path.join(OUT_DATA_DIR, f\"info_{dset_type}.json\")\n",
    "\n",
    "# gather the data\n",
    "_ = np.memmap(out_mmap_path, dtype=np.uint16, mode=\"w+\", shape=(1,))\n",
    "n_offs = 0\n",
    "tot_duration_dict = defaultdict(int)\n",
    "datasets_info = defaultdict(dict)\n",
    "n = 0\n",
    "for i, row in tqdm.tqdm(enumerate(val_inputs), total=len(val_inputs)):\n",
    "    # we need to alternate between preference: neg, pos\n",
    "    preference = False\n",
    "    local_path = f\"{NPZ_DIR}/{row['id'].replace('_gen_', '_gen_cycle_')}.npz\"\n",
    "    if not os.path.exists(local_path):\n",
    "        # print(row, local_path)\n",
    "        raise ValueError()\n",
    "    try:\n",
    "        arr = np.load(local_path)[\"v3.0_raw\"]\n",
    "    except Exception as e:\n",
    "        print(local_path)\n",
    "        raise e\n",
    "    assert arr.shape[0] <= 3000\n",
    "    assert arr.shape[1] == 13\n",
    "    arr_duration = arr.shape[0] / 25\n",
    "    # print(arr.shape)\n",
    "    arr = reshift(arr)\n",
    "    # print(\"after shift and pad\", arr.shape)\n",
    "    arr = arr.reshape(\n",
    "        -1,\n",
    "    )\n",
    "    # print(arr.shape)\n",
    "    out_mm = np.memmap(\n",
    "        out_mmap_path,\n",
    "        dtype=np.uint16,\n",
    "        mode=\"r+\",\n",
    "        shape=(n_offs + arr.size,),\n",
    "    )\n",
    "    out_mm[n_offs : n_offs + arr.size] = arr\n",
    "    # print(f\"offset is: {n_offs}\")\n",
    "    # break\n",
    "    # write it once\n",
    "    out_mm.flush()\n",
    "    del out_mm\n",
    "    add_metas = []\n",
    "    add_meta = {\n",
    "        \"dataset\": f\"perference_{int(preference)}\",\n",
    "        \"id\": row[\"id\"],  # this is the row s3_id\n",
    "        \"start_s\": 0,\n",
    "        \"end_s\": None,  # for full clips we do know it has an edding, other wise, we don't know\n",
    "        \"original_duration_s\": None,\n",
    "        \"vocal_start_s\": None,  # these are unfortunately missing for now\n",
    "        \"vocal_end_s\": None,  # these are unfortunately missing for now\n",
    "        \"tags\": [row[\"metadata\"][\"tags\"]],  # tags is a list, do you know :)\n",
    "        \"text\": row[\"prompt_text\"],\n",
    "    }\n",
    "    add_metas.append(add_meta)\n",
    "    tot_duration_dict[preference] += arr_duration\n",
    "    write_jsonl(\n",
    "        add_metas,\n",
    "        os.path.join(out_metas_path),\n",
    "        do_append=bool(n_offs != 0),\n",
    "    )\n",
    "    if \"idx_list\" not in datasets_info[add_meta[\"dataset\"]]:\n",
    "        datasets_info[add_meta[\"dataset\"]][\"idx_list\"] = [n]\n",
    "    else:\n",
    "        datasets_info[add_meta[\"dataset\"]][\"idx_list\"].append(n)\n",
    "    n_offs += arr.size\n",
    "    n += 1\n",
    "    # let's do the real music here\n",
    "    preference = True\n",
    "    real_meta = gen_input_to_real_metas[row[\"id\"]]\n",
    "    real_meta_index = real_meta[\"lookup_index\"]\n",
    "    assert real_meta.get(\"text\", \"\") == row[\"prompt_text\"]\n",
    "    real_arr = val_real_mm[real_meta_index, :].copy()\n",
    "    # print(real_arr.shape)\n",
    "    real_arr = reshift(real_arr)\n",
    "    # print(\"after shift and pad\", real_arr.shape)\n",
    "\n",
    "    real_arr = real_arr.reshape(\n",
    "        -1,\n",
    "    )\n",
    "    out_mm = np.memmap(\n",
    "        out_mmap_path,\n",
    "        dtype=np.uint16,\n",
    "        mode=\"r+\",\n",
    "        shape=(n_offs + real_arr.size,),\n",
    "    )\n",
    "    out_mm[n_offs : n_offs + real_arr.size] = real_arr\n",
    "    out_mm.flush()\n",
    "    del out_mm\n",
    "    arr_duration = real_meta.get(\"end_s\", 0) - real_meta.get(\"start_s\", 0)\n",
    "    add_metas = []\n",
    "    add_meta = {\n",
    "        \"dataset\": f\"perference_{int(preference)}\",\n",
    "        \"id\": real_meta[\"id\"],  # this is the row s3_id\n",
    "        \"start_s\": 0,\n",
    "        \"end_s\": None,  # for full clips we do know it has an edding, other wise, we don't know\n",
    "        \"original_duration_s\": None,\n",
    "        \"vocal_start_s\": None,  # these are unfortunately missing for now\n",
    "        \"vocal_end_s\": None,  # these are unfortunately missing for now\n",
    "        \"tags\": [row[\"metadata\"][\"tags\"]],  # tags is a list, do you know :)\n",
    "        \"text\": row[\"prompt_text\"],\n",
    "    }\n",
    "    add_metas.append(add_meta)\n",
    "    tot_duration_dict[preference] += arr_duration\n",
    "    write_jsonl(\n",
    "        add_metas,\n",
    "        os.path.join(out_metas_path),\n",
    "        do_append=bool(n_offs != 0),\n",
    "    )\n",
    "    if \"idx_list\" not in datasets_info[add_meta[\"dataset\"]]:\n",
    "        datasets_info[add_meta[\"dataset\"]][\"idx_list\"] = [n]\n",
    "    else:\n",
    "        datasets_info[add_meta[\"dataset\"]][\"idx_list\"].append(n)\n",
    "    n_offs += arr.size\n",
    "    n += 1\n",
    "\n",
    "write_json(datasets_info, out_info_filepath)\n",
    "print(f\"Total {n} clips\")\n",
    "for k, v in tot_duration_dict.items():\n",
    "    print(f\"{round(v / 60 / 60):,} hours of {k}\")\n",
    "print(f\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-29T19:46:47.549860Z",
     "start_time": "2024-01-29T19:46:47.548015Z"
    }
   },
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.557Z"
    }
   },
   "outputs": [],
   "source": [
    "# verify\n",
    "mm = np.memmap(os.path.join(OUT_DATA_DIR, f\"data_val.bin\"), dtype=np.uint16, mode=\"r\")\n",
    "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"meta_val.jsonl\"))\n",
    "test_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_val.json\"))\n",
    "mm = mm.reshape(-1, 6016, 13)\n",
    "assert len(mm) == len(test_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() <= 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.557Z"
    }
   },
   "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_stream_to_full_audio as codec_decode,\n",
    "#     EMBEDDING_RATE as CODEC_EMBEDDING_RATE,\n",
    "#     decode as decode\n",
    "# )\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "# _ = preload_codec_models(\"/app/suno/tony/v3/dac_2c_25x12.pt\", device=\"cuda\")\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,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.558Z"
    }
   },
   "outputs": [],
   "source": [
    "# idx = random.choice(test_info[\"perference_0\"][\"idx_list\"])\n",
    "# assert \"original_duration_s\" in test_metas[idx]\n",
    "# # positive index should be shifted by 1\n",
    "# pos_idx = idx + 1\n",
    "# print(\n",
    "#     \"tags:\",\n",
    "#     test_metas[idx].get(\"tags\") == test_metas[pos_idx].get(\"tags\"),\n",
    "#     test_metas[idx].get(\"tags\"),\n",
    "# )\n",
    "# arr = mm[idx, 1:].copy().astype(np.int16)[:, 1:]\n",
    "# pos_arr = mm[pos_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",
    "# pos_pad_idx_arr = np.where(pos_arr == COARSE_PAD_TOKEN)[0]\n",
    "# if len(pos_pad_idx_arr) > 0:\n",
    "#     pos_arr = pos_arr[: pos_pad_idx_arr[0], :]\n",
    "# a = decode(arr)\n",
    "# print(\"negative example\", test_metas[idx])\n",
    "# a.play(compress=False)\n",
    "# pos_a = decode(pos_arr)\n",
    "# print(\"positive example\", test_metas[pos_idx])\n",
    "# pos_a.play(compress=False)\n",
    "# print(\n",
    "#     \"text:\",\n",
    "#     test_metas[idx].get(\"text\") == test_metas[pos_idx].get(\"text\"),\n",
    "#     test_metas[idx].get(\"text\"),\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.558Z"
    }
   },
   "outputs": [],
   "source": [
    "# val_df[val_df[\"tags\"] == 'a vibrant blend of experimental jazz fusion, drum-and-bass and swagger fuzzed-out guitars']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.559Z"
    }
   },
   "outputs": [],
   "source": [
    "# from collections import Counter\n",
    "# c = Counter()\n",
    "# for _, row in df_slice.iterrows():\n",
    "#     # print(row[\"metadata\"])\n",
    "#     for k in ast.literal_eval(row[\"metadata\"]).keys():\n",
    "#         c[k] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.559Z"
    }
   },
   "outputs": [],
   "source": [
    "# original_npz_path = f\"/app/suno/data/dpo/7b_npz/{test_metas[idx]['id']}.npz\"\n",
    "# original_npz_path = \"/app/suno/data/dpo/7b_npz/729c3011-f672-4ccd-8d82-1cbf2b52ff69.npz\"\n",
    "# original_arr = np.load(original_npz_path)[\"v2_raw\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.559Z"
    }
   },
   "outputs": [],
   "source": [
    "def validation_on_metas(input_metas):\n",
    "    total_bad = 0\n",
    "    total_good = 0\n",
    "    for idx in range(len(input_metas)):\n",
    "        if idx % 2 == 0:\n",
    "            pos_idx = idx + 1\n",
    "            if input_metas[idx].get(\"tags\") != input_metas[pos_idx].get(\"tags\"):\n",
    "                # print(test_metas[idx].get(\"text\") == test_metas[pos_idx].get(\"text\"), test_metas[idx].get(\"tags\"), test_metas[pos_idx].get(\"tags\"))\n",
    "                total_bad += 1\n",
    "            else:\n",
    "                total_good += 1\n",
    "    print(total_good, total_bad)\n",
    "    return\n",
    "\n",
    "\n",
    "validation_on_metas(test_metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-17T00:32:08.486813Z",
     "start_time": "2024-05-17T00:32:08.460129Z"
    }
   },
   "outputs": [],
   "source": [
    "train_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_tr.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-17T00:32:08.552057Z",
     "start_time": "2024-05-17T00:32:08.548024Z"
    }
   },
   "outputs": [],
   "source": [
    "n_neg_tr = train_info[\"perference_0\"][\"idx_list\"]\n",
    "n_pos_tr = train_info[\"perference_1\"][\"idx_list\"]\n",
    "assert len(n_pos_tr) == len(n_neg_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-17T00:32:08.697717Z",
     "start_time": "2024-05-17T00:32:08.695940Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total samples 347872\n"
     ]
    }
   ],
   "source": [
    "total_iters = len(n_neg_tr) + len(n_pos_tr)\n",
    "print(\"total samples\", total_iters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-17T00:32:08.989719Z",
     "start_time": "2024-05-17T00:32:08.987998Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 epoch per batch 4, total 3623.6666666666665\n"
     ]
    }
   ],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 8 / 3 / 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.561Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm && sbatch sbatch_dpo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.561Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm && sbatch sbatch_ipo_real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-05-16T20:55:31.562Z"
    }
   },
   "outputs": [],
   "source": [
    "# converted_paths = os.listdir(\"/app/suno/data/dpo/7b_recycle_npz\")\n",
    "\n",
    "# # don't run this unless you kill some job accidentally\n",
    "# for s3_processed_file in tqdm.tqdm(converted_paths):\n",
    "#     processed_file = os.path.basename(s3_processed_file)\n",
    "#     processed_file_path = os.path.join(\"/app/suno/data/dpo/7b_recycle_npz\", processed_file)\n",
    "#     # could have been removed already\n",
    "#     if os.path.exists(processed_file_path):\n",
    "#         file_size = os.stat(processed_file_path).st_size\n",
    "#         # print(processed_file_path, file_size)\n",
    "#         # break\n",
    "#         if file_size < 2000:\n",
    "#             print(processed_file_path, file_size)\n",
    "#             os.remove(processed_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env_dev",
   "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.15"
  },
  "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": 4
}
