{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:41.381084Z",
     "iopub.status.busy": "2025-06-05T14:35:41.380857Z",
     "iopub.status.idle": "2025-06-05T14:35:41.393493Z",
     "shell.execute_reply": "2025-06-05T14:35:41.393081Z",
     "shell.execute_reply.started": "2025-06-05T14:35:41.381061Z"
    }
   },
   "outputs": [],
   "source": [
    "# setup autoload\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:21.040680Z",
     "start_time": "2024-05-16T13:58:19.777010Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:41.395105Z",
     "iopub.status.busy": "2025-06-05T14:35:41.394986Z",
     "iopub.status.idle": "2025-06-05T14:35:43.878474Z",
     "shell.execute_reply": "2025-06-05T14:35:43.877918Z",
     "shell.execute_reply.started": "2025-06-05T14:35:41.395092Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "import ast\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "from collections import defaultdict\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from preference_data_preparation_auk import *\n",
    "from preference_helper import *\n",
    "from sklearn.model_selection import train_test_split\n",
    "from suno_utils.utils.s3 import download_s3_files\n",
    "from suno_utils.utils.text import read_json, read_jsonl, write_json, write_jsonl\n",
    "from tqdm import tqdm\n",
    "\n",
    "pd.set_option(\"display.max_rows\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)\n",
    "\n",
    "# setup autoload\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "\n",
    "def custom_parse(x):\n",
    "    try:\n",
    "        return json.loads(x)\n",
    "    except:\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:21.082172Z",
     "start_time": "2024-05-16T13:58:21.041926Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:43.880437Z",
     "iopub.status.busy": "2025-06-05T14:35:43.880310Z",
     "iopub.status.idle": "2025-06-05T14:35:43.948421Z",
     "shell.execute_reply": "2025-06-05T14:35:43.947954Z",
     "shell.execute_reply.started": "2025-06-05T14:35:43.880423Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N_TOKENS_AUDIO 12000\n"
     ]
    }
   ],
   "source": [
    "OUT_DATA_DIR = \"/app2/suno/data/dpo/auk_t0_v20\"\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "shutil.copyfile(\n",
    "    \"/app/suno/data/dpo/7v_v20_full/tokenizer_60k.json\",\n",
    "    os.path.join(OUT_DATA_DIR, \"tokenizer_60k.json\"),\n",
    ")\n",
    "NPZ_DIR = \"/app2/suno/data/dpo/auk_t0_npz\"\n",
    "N_TOKENS_AUDIO = 25 * 8 * 60\n",
    "print(\"N_TOKENS_AUDIO\", N_TOKENS_AUDIO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import glob\n",
    "# import os\n",
    "# from tqdm import tqdm\n",
    "\n",
    "# # Load the base dataframe\n",
    "# base_file = \"/home/tony/Data/Preference/auk_t0/fully_merged_auk_t0.pkl\"\n",
    "# df = pd.read_pickle(base_file)\n",
    "\n",
    "# base_time_file = base_file  # \"/home/tony/Data/Preference/auk_t0/interesting_clips_auk_t0_20250606.pkl\"\n",
    "# base_ctime = os.path.getctime(base_time_file)\n",
    "# # print(f\"Base file ctime: {base_ctime}\")\n",
    "\n",
    "# # Find all pkl files in the directory with same name pattern\n",
    "# pkl_files = glob.glob(\"/home/tony/Data/Preference/auk_t0/interesting_clips_*.pkl\")\n",
    "\n",
    "# # Filter files that are newer than the base file and print debug info\n",
    "# newer_files = []\n",
    "# for f in pkl_files:\n",
    "#     f_ctime = os.path.getctime(f)\n",
    "#     if f_ctime > base_ctime:\n",
    "#         newer_files.append(f)\n",
    "#         print(f\"Found newer file: {os.path.basename(f)}\")\n",
    "#         print(f\"  Base file ctime: {base_ctime}\")\n",
    "#         print(f\"  File ctime: {f_ctime}\")\n",
    "#         print(f\"  Difference: {f_ctime - base_ctime} seconds\")\n",
    "\n",
    "# newer_files.sort(key=lambda x: os.path.getctime(x))\n",
    "\n",
    "# print(f\"Found {len(newer_files)} newer files to process\")\n",
    "\n",
    "# # Process each newer file\n",
    "# for pkl_file in tqdm(newer_files, desc=\"Loading newer pickle files\"):\n",
    "#     print(f\"\\nProcessing {os.path.basename(pkl_file)}\")\n",
    "#     prev_size = len(df)\n",
    "#     temp_df = pd.read_pickle(pkl_file)\n",
    "#     new_size = len(temp_df)\n",
    "\n",
    "#     # Convert datetime columns if they exist\n",
    "#     for col in [\"created_at\", \"updated_at\"]:\n",
    "#         if col in temp_df.columns:\n",
    "#             temp_df[col] = pd.to_datetime(temp_df[col], utc=True)\n",
    "\n",
    "#     # Handle duplicates based on id\n",
    "#     if \"id\" in temp_df.columns:\n",
    "#         df = pd.concat([df, temp_df], ignore_index=True)\n",
    "#         df = df.drop_duplicates(subset=[\"id\"], keep=\"last\")\n",
    "#     else:\n",
    "#         df = pd.concat([df, temp_df], ignore_index=True)\n",
    "\n",
    "#     # Print size statistics\n",
    "#     current_size = len(df)\n",
    "#     net_increase = current_size - prev_size\n",
    "#     print(f\"Previous size: {prev_size:,}\")\n",
    "#     print(f\"New input size: {new_size:,}\")\n",
    "#     print(f\"Current total size: {current_size:,}\")\n",
    "#     print(f\"Net increase: {net_increase:,}\")\n",
    "\n",
    "# print(\"\\nFinal dataframe shape:\", df.shape)\n",
    "# print(\n",
    "#     \"Unique ids:\",\n",
    "#     df[\"id\"].nunique() if \"id\" in df.columns else \"No id column\",\n",
    "# )\n",
    "# df.to_pickle(\"/home/tony/Data/Preference/auk_t0/fully_merged_auk_t0.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:53.962528Z",
     "start_time": "2024-05-16T13:58:21.105919Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:43.950183Z",
     "iopub.status.busy": "2025-06-05T14:35:43.950058Z",
     "iopub.status.idle": "2025-06-05T14:35:46.603501Z",
     "shell.execute_reply": "2025-06-05T14:35:46.602927Z",
     "shell.execute_reply.started": "2025-06-05T14:35:43.950169Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preference data shape (340122, 90)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\n",
    "    \"/home/tony/Data/Preference/auk_t0/interesting_clips_auk_t0_20250728.pkl\"\n",
    ")\n",
    "print(\"Preference data shape\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total hoot cer scores: 306236\n"
     ]
    }
   ],
   "source": [
    "# clip_id_to_cer = {}\n",
    "with open(\"/home/tony/Data/Preference/auk_t0/hoot_cer.json\", \"r\") as file:\n",
    "    clip_id_to_cer = json.load(file)\n",
    "print(\"Total hoot cer scores:\", len(clip_id_to_cer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "340122it [00:42, 7975.66it/s] \n"
     ]
    }
   ],
   "source": [
    "# find all the hoot jsons in the json dir\n",
    "JSON_DIR = \"/app2/suno/data/dpo/auk_t0_json/\"\n",
    "for _, row in tqdm(df.iterrows()):\n",
    "    clip_id = row[\"s3_id\"]\n",
    "    if clip_id in clip_id_to_cer:\n",
    "        continue\n",
    "    hoot_json_path = os.path.join(JSON_DIR, f\"{clip_id}_hoot.json\")\n",
    "    if not os.path.exists(hoot_json_path):\n",
    "        clip_id_to_cer[clip_id] = 1.0\n",
    "        continue\n",
    "    with open(os.path.join(JSON_DIR, f\"{clip_id}_hoot.json\"), \"r\") as f:\n",
    "        data = json.load(f)\n",
    "    for data_dict in data:\n",
    "        if \"hoot_cer\" in data_dict:\n",
    "            clip_id_to_cer[clip_id] = data_dict[\"hoot_cer\"]\n",
    "            break\n",
    "\n",
    "# add the cer to the df\n",
    "df[\"cer\"] = df[\"s3_id\"].map(clip_id_to_cer)\n",
    "\n",
    "# BREAK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# update the hoot cer cache\n",
    "with open(\"/home/tony/Data/Preference/auk_t0/hoot_cer.json\", \"w\") as file:\n",
    "    json.dump(clip_id_to_cer, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.fillna({\"cer\": 1})\n",
    "df[\"cer_diff\"] = df[\"cer\"].diff()\n",
    "df = df.fillna({\"cer_diff\": 0})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:46.605521Z",
     "iopub.status.busy": "2025-06-05T14:35:46.605393Z",
     "iopub.status.idle": "2025-06-05T14:35:46.620570Z",
     "shell.execute_reply": "2025-06-05T14:35:46.620130Z",
     "shell.execute_reply.started": "2025-06-05T14:35:46.605507Z"
    }
   },
   "outputs": [],
   "source": [
    "# df[\"created_at\"] = pd.to_datetime(df[\"created_at\"], utc=True)\n",
    "# cutoff_date = pd.to_datetime(\"2025-05-15\", utc=True)\n",
    "# # cutoff_date = pd.to_datetime(\"2025-04-17\", utc=True)\n",
    "# print(df.shape, df[df[\"created_at\"] > cutoff_date].shape)\n",
    "# df = df[(df[\"created_at\"] > cutoff_date)].copy()\n",
    "# print(\"after date cut\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:46.622306Z",
     "iopub.status.busy": "2025-06-05T14:35:46.622176Z",
     "iopub.status.idle": "2025-06-05T14:35:46.816997Z",
     "shell.execute_reply": "2025-06-05T14:35:46.816439Z",
     "shell.execute_reply.started": "2025-06-05T14:35:46.622292Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "after dropna (340122, 87)\n"
     ]
    }
   ],
   "source": [
    "df = df.dropna(axis=1, how=\"all\")\n",
    "df[\"s3_id\"] = df[\"id\"].astype(str)\n",
    "print(\"after dropna\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.199480Z",
     "start_time": "2024-05-16T13:58:53.963687Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:46.817692Z",
     "iopub.status.busy": "2025-06-05T14:35:46.817546Z",
     "iopub.status.idle": "2025-06-05T14:35:50.559461Z",
     "shell.execute_reply": "2025-06-05T14:35:50.558720Z",
     "shell.execute_reply.started": "2025-06-05T14:35:46.817677Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "455152\n",
      "455152\n",
      "pre-downloaded df (340122, 87)\n",
      "downloaded df (340122, 87)\n"
     ]
    }
   ],
   "source": [
    "converted_paths = os.listdir(NPZ_DIR)\n",
    "print(len(converted_paths))\n",
    "\n",
    "converted_paths = set([f.replace(\".npz\", \"\") for f in converted_paths])\n",
    "print(len(converted_paths))\n",
    "\n",
    "print(\"pre-downloaded df\", df.shape)\n",
    "df[df[\"s3_id\"].isin(converted_paths)].shape\n",
    "df = df[df[\"s3_id\"].isin(converted_paths)].copy()\n",
    "print(\"downloaded df\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.467253Z",
     "start_time": "2024-05-16T13:58:56.207647Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:50.560340Z",
     "iopub.status.busy": "2025-06-05T14:35:50.560172Z",
     "iopub.status.idle": "2025-06-05T14:35:50.587770Z",
     "shell.execute_reply": "2025-06-05T14:35:50.587208Z",
     "shell.execute_reply.started": "2025-06-05T14:35:50.560323Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "task\n",
       "                      222668\n",
       "cover                  55018\n",
       "artist_consistency     30994\n",
       "artist_cover           11798\n",
       "extend                 11704\n",
       "upload_extend           6020\n",
       "artist_extend           1918\n",
       "underpainting              2\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"task\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LET's do the data prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.592883Z",
     "start_time": "2024-05-16T13:58:56.470781Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:50.588559Z",
     "iopub.status.busy": "2025-06-05T14:35:50.588398Z",
     "iopub.status.idle": "2025-06-05T14:35:50.688406Z",
     "shell.execute_reply": "2025-06-05T14:35:50.687815Z",
     "shell.execute_reply.started": "2025-06-05T14:35:50.588544Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preference  model_name  \n",
      "False       chirp-auk-t0    170061\n",
      "True        chirp-auk-t0    170061\n",
      "Name: count, dtype: int64\n",
      "before filter on model name (340122, 87)\n",
      "after filter on model name (340122, 87)\n",
      "is_public\n",
      "False    327486\n",
      "True      12636\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "## for 13b this is easy for now\n",
    "print(df.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "print(\"before filter on model name\", df.shape)\n",
    "df = df[df[\"model_name\"].isin([\"chirp-auk-t0\"])]\n",
    "print(\"after filter on model name\", df.shape)\n",
    "print(df[\"is_public\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.909539Z",
     "start_time": "2024-05-16T13:58:56.595736Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:50.689384Z",
     "iopub.status.busy": "2025-06-05T14:35:50.689234Z",
     "iopub.status.idle": "2025-06-05T14:35:50.867060Z",
     "shell.execute_reply": "2025-06-05T14:35:50.866366Z",
     "shell.execute_reply.started": "2025-06-05T14:35:50.689370Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before filter on request id pairs (340122, 87)\n",
      "after filter on request id pairs (340122, 87)\n",
      "preference  model_name  \n",
      "False       chirp-auk-t0    170061\n",
      "True        chirp-auk-t0    170061\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(\"before filter on request id pairs\", df.shape)\n",
    "df = df[\n",
    "    df[\"request_id\"].isin(\n",
    "        df[\"request_id\"].value_counts().index[df[\"request_id\"].value_counts() == 2]\n",
    "    )\n",
    "]\n",
    "print(\"after filter on request id pairs\", df.shape)\n",
    "print(df.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "assert df.shape[0] == df[\"request_id\"].nunique() * 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:35:50.868033Z",
     "iopub.status.busy": "2025-06-05T14:35:50.867771Z",
     "iopub.status.idle": "2025-06-05T14:36:18.830004Z",
     "shell.execute_reply": "2025-06-05T14:36:18.829246Z",
     "shell.execute_reply.started": "2025-06-05T14:35:50.867967Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unique_requests 170061\n",
      "before removing duplicates (340122, 183)\n",
      "after removing duplicates (340122, 176)\n"
     ]
    }
   ],
   "source": [
    "# Let's use the old selection for now -- for quality assurance\n",
    "# expand the metadata columns -- this takes forever...~ 6 mins\n",
    "# test_slice = df[\"metadata\"].apply(lambda x: ast.literal_eval(str(x)))\n",
    "# test_slice = df[\"metadata\"].apply(lambda x: custom_parse(x))\n",
    "test_slice = df[\"metadata\"]  # .apply(lambda x: json.loads(x))\n",
    "test_slice_series = test_slice.apply(pd.Series)\n",
    "df = pd.concat([df, test_slice_series], axis=1, join=\"inner\")\n",
    "print(\"unique_requests\", df[\"request_id\"].nunique())\n",
    "# remove the duplicates\n",
    "print(\"before removing duplicates\", df.shape)\n",
    "df = df.loc[:, ~df.columns.duplicated()].copy()\n",
    "print(\"after removing duplicates\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:18.830948Z",
     "iopub.status.busy": "2025-06-05T14:36:18.830768Z",
     "iopub.status.idle": "2025-06-05T14:36:19.142965Z",
     "shell.execute_reply": "2025-06-05T14:36:19.142338Z",
     "shell.execute_reply.started": "2025-06-05T14:36:18.830930Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pos_diff_preference\n",
       "1.0    124003\n",
       "2.0     46058\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = df.sort_values(by=[\"request_id\", \"preference\", \"diff_preference\"])\n",
    "df[\"pos_diff_preference\"] = df[\"diff_preference\"].diff()\n",
    "# df[\"cer_diff_preference\"] = df[\"cer\"].diff()\n",
    "df[df[\"preference\"]][\"pos_diff_preference\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:19.143842Z",
     "iopub.status.busy": "2025-06-05T14:36:19.143676Z",
     "iopub.status.idle": "2025-06-05T14:36:19.268563Z",
     "shell.execute_reply": "2025-06-05T14:36:19.267877Z",
     "shell.execute_reply.started": "2025-06-05T14:36:19.143826Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive param_experiment\n",
      "mask_control_slider    1883\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(\"positive\", df[df[\"preference\"]][\"param_experiment\"].value_counts())\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:19.269458Z",
     "iopub.status.busy": "2025-06-05T14:36:19.269284Z",
     "iopub.status.idle": "2025-06-05T14:36:20.963563Z",
     "shell.execute_reply": "2025-06-05T14:36:20.962833Z",
     "shell.execute_reply.started": "2025-06-05T14:36:19.269440Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 788 duplicated prompts 394 unique requests\n",
      "Found 166 request_ids with duplicate prompts but not highest play counts in their group\n",
      "['acbb68b7-f59d-45f3-b624-8f8b69e4ac6b', '4dd77060-5e38-42b5-86f5-8e76b12eee4c', '41bf36a1-2d43-4246-92a2-472da7099828', '0249eb98-0678-431e-8993-b6a8518a7cb0', '26723fd0-383f-4696-813c-225131a92491', '6d1c2b46-d25c-4a37-9956-e765d5f3eb82', '1d511b27-cc68-4dc5-b176-5ee0e824cb05', '84a0c0f4-0d9c-4d07-b8d1-47b825f0e7f8', 'ddc1f82e-089c-41ce-b884-3442889ae4a8', '75ef4d3b-3fc4-4df7-b349-74dc3f47490f']\n",
      "Before dedup user gen requests 340122\n",
      "After dedup user gen requests 340122\n"
     ]
    }
   ],
   "source": [
    "# Find duplicated prompts with count > 2\n",
    "duplicate_entries = df.groupby(\n",
    "    [\"user_id\", \"prompt_text\", \"tags\", \"task\", \"edited_clip_id\"]\n",
    ").filter(lambda x: len(x) > 2)\n",
    "print(\n",
    "    \"Found\",\n",
    "    len(duplicate_entries),\n",
    "    \"duplicated prompts\",\n",
    "    len(duplicate_entries[\"request_id\"].unique()),\n",
    "    \"unique requests\",\n",
    ")\n",
    "\n",
    "# Group by user_id, prompt_text, and tags to find duplicate prompt groups\n",
    "prompt_groups = duplicate_entries.groupby(\n",
    "    [\"user_id\", \"prompt_text\", \"tags\", \"task\", \"edited_clip_id\"]\n",
    ")\n",
    "\n",
    "# For each prompt group, find the request_id with the highest total reaction_play_count\n",
    "low_play_count_request_ids = []\n",
    "for prompt_key, prompt_group in prompt_groups:\n",
    "    # Get the sum of reaction_play_count for each request_id in this group\n",
    "    request_play_counts = prompt_group.groupby(\"request_id\")[\n",
    "        \"reaction_play_count\"\n",
    "    ].sum()\n",
    "\n",
    "    # Find the max play count in this group\n",
    "    max_play_count = request_play_counts.max()\n",
    "\n",
    "    # Add request_ids that don't have the max play count to our filter list\n",
    "    lower_play_count_request_ids = request_play_counts[\n",
    "        request_play_counts < max_play_count\n",
    "    ].index.tolist()\n",
    "    low_play_count_request_ids.extend(lower_play_count_request_ids)\n",
    "\n",
    "# Display the filtered request IDs\n",
    "print(\n",
    "    f\"Found {len(low_play_count_request_ids)} request_ids with duplicate prompts but not highest play counts in their group\"\n",
    ")\n",
    "print(\n",
    "    low_play_count_request_ids[:10]\n",
    "    if len(low_play_count_request_ids) > 10\n",
    "    else low_play_count_request_ids\n",
    ")\n",
    "print(\"Before dedup user gen requests\", df.shape[0])\n",
    "# df = df[~df[\"request_id\"].isin(low_play_count_request_ids)]\n",
    "print(\"After dedup user gen requests\", df.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3421218/815349087.py:1: FutureWarning: A value is trying to be set on a copy of a DataFrame or Series through chained assignment using an inplace method.\n",
      "The behavior will change in pandas 3.0. This inplace method will never work because the intermediate object on which we are setting values always behaves as a copy.\n",
      "\n",
      "For example, when doing 'df[col].method(value, inplace=True)', try using 'df.method({col: value}, inplace=True)' or df[col] = df[col].method(value) instead, to perform the operation inplace on the original object.\n",
      "\n",
      "\n",
      "  df[\"task\"].fillna(\"\", inplace=True)\n"
     ]
    }
   ],
   "source": [
    "df[\"task\"].fillna(\"\", inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:40.799375Z",
     "start_time": "2024-05-16T13:59:36.394236Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:20.964466Z",
     "iopub.status.busy": "2025-06-05T14:36:20.964282Z",
     "iopub.status.idle": "2025-06-05T14:36:22.022063Z",
     "shell.execute_reply": "2025-06-05T14:36:22.021326Z",
     "shell.execute_reply.started": "2025-06-05T14:36:20.964450Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9567\n",
      "good_continue_at\n",
      "True     340078\n",
      "False        44\n",
      "Name: count, dtype: int64\n",
      "\n",
      " Check some basics... \n",
      " preference\n",
      "False    170061\n",
      "True     170061\n",
      "Name: count, dtype: int64 model_name\n",
      "chirp-auk-t0    340122\n",
      "Name: count, dtype: int64 preference  model_name  \n",
      "False       chirp-auk-t0    170061\n",
      "True        chirp-auk-t0    170061\n",
      "Name: count, dtype: int64\n",
      "task\n",
      "                      222668\n",
      "cover                  55018\n",
      "artist_consistency     30994\n",
      "artist_cover           11798\n",
      "extend                 11704\n",
      "upload_extend           6020\n",
      "artist_extend           1918\n",
      "underpainting              2\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"id\"] = df[\"str_id\"]\n",
    "# get the original duration of the clips, if they are concacted\n",
    "df[\"original_duration_s\"] = df[\"total_start_s\"] + df[\"duration\"]\n",
    "# classify the continue at behavoirs by the duration choice\n",
    "audio_prompt_id_to_continue_at = {}\n",
    "\n",
    "for _, row in df[~df[\"continued_parent\"].isna()].iterrows():\n",
    "    audio_prompt_id = row[\"continued_parent\"]\n",
    "    if audio_prompt_id not in audio_prompt_id_to_continue_at:\n",
    "        audio_prompt_id_to_continue_at[audio_prompt_id] = row[\"continue_at\"]\n",
    "    else:\n",
    "        # pick the max\n",
    "        audio_prompt_id = max(\n",
    "            audio_prompt_id_to_continue_at[audio_prompt_id], row[\"continue_at\"]\n",
    "        )\n",
    "print(len(audio_prompt_id_to_continue_at))\n",
    "df[\"has_continue_and_start_continue_at\"] = df[\"id\"].apply(\n",
    "    lambda x: audio_prompt_id_to_continue_at.get(x)\n",
    ")\n",
    "# we want continue at to be at most of the clip...\n",
    "df[\"good_continue_at\"] = (\n",
    "    (df[\"has_continue_and_start_continue_at\"] / df[\"duration\"]) > 0.9\n",
    ") | df[\"has_continue_and_start_continue_at\"].isna()\n",
    "print(df[\"good_continue_at\"].value_counts())\n",
    "\n",
    "\n",
    "print(\n",
    "    \"\\n Check some basics... \\n\",\n",
    "    df[\"preference\"].value_counts(),\n",
    "    df[\"model_name\"].value_counts(),\n",
    "    df.groupby([\"preference\"])[\"model_name\"].value_counts(),\n",
    ")\n",
    "\n",
    "df = df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "df[\"duration_rel_diff\"] = df[\"duration\"].diff()\n",
    "df[\"play_rel_diff\"] = df[\"reaction_play_count\"].diff()\n",
    "print(df[\"task\"].value_counts())\n",
    "\n",
    "df[\"post_infill_duration\"] = (\n",
    "    df[\"duration\"]\n",
    "    + df[\"infill_context_end_s\"]\n",
    "    - df[\"infill_context_start_s\"]\n",
    "    - df[\"include_future_s\"]\n",
    "    - df[\"include_history_s\"]\n",
    "    - df[\"infill_dur_s\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.035167Z",
     "start_time": "2024-05-16T13:59:40.801098Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:22.023244Z",
     "iopub.status.busy": "2025-06-05T14:36:22.022980Z",
     "iopub.status.idle": "2025-06-05T14:36:22.558373Z",
     "shell.execute_reply": "2025-06-05T14:36:22.557641Z",
     "shell.execute_reply.started": "2025-06-05T14:36:22.023226Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "after duration 0.9990621012460235\n",
      "after infill duration 1.0\n",
      "neg_filter_reaction_play_count 1.0\n",
      "neg_filter_upvote_count 0.9928\n",
      "neg_filter_norm_play_frac 1.0\n",
      "neg_filter_continues 1.0\n",
      "----------------\n",
      "pos_filter_continues 0.9997\n",
      "pos_filter_reaction_play_count 1.0\n",
      "pos_filter_relative_play_count 0.9611\n",
      "pos_filter_cer_diff_preference 1.0\n",
      "pos_filter_bad_flags 0.9999\n",
      "after filter on play counts 0.9953\n",
      "after filter on higher quality 0.3371\n",
      "after filter on play duration 0.6998\n",
      "after filter on task 0.6891\n",
      "after filter on cer diff 0.9213\n",
      "----------------\n",
      "negative 168670 positive 37334\n",
      "----------------\n",
      "total pair requests 170061  --> selected pair requests 37064 frac 0.218  --> total intitial users 96466\n"
     ]
    }
   ],
   "source": [
    "normal_pos_play_count = 3\n",
    "# this is lower, cause a concat is probably already ensuring that it is good\n",
    "concat_pos_play_count = 1\n",
    "# this is a filter on the concated clip\n",
    "concat_total_play_count = 3\n",
    "\n",
    "all_fitlers = (df[\"duration\"] >= 10) & (df[\"duration\"] <= 480)\n",
    "print(\"after duration\", all_fitlers.sum() / df.shape[0])\n",
    "infill_duration_filter = ~df[\"task\"].isin(\n",
    "    [\n",
    "        \"infill\",\n",
    "        \"infill_intro\",\n",
    "        \"infill_outro\",\n",
    "    ]\n",
    ") | (df[\"post_infill_duration\"] <= 460)\n",
    "print(\"after infill duration\", infill_duration_filter.sum() / df.shape[0])\n",
    "# negative fitlers\n",
    "total_negative = df[~df[\"preference\"]].shape[0]\n",
    "neg_filter_reaction_play_count = (~df[\"preference\"]) & (df[\"reaction_play_count\"] >= 1)\n",
    "print(\n",
    "    \"neg_filter_reaction_play_count\",\n",
    "    round(neg_filter_reaction_play_count.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_upvote_count = (~df[\"preference\"]) & (df[\"upvote_count\"] == 0)\n",
    "print(\n",
    "    \"neg_filter_upvote_count\",\n",
    "    round(neg_filter_upvote_count.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_norm_play_frac = (~df[\"preference\"]) & (df[\"norm_play_frac\"] <= 3.1)\n",
    "print(\n",
    "    \"neg_filter_norm_play_frac\",\n",
    "    round(neg_filter_norm_play_frac.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_continues = (~df[\"preference\"]) & (\n",
    "    df[\"has_continue_and_start_continue_at\"].isna()\n",
    ")\n",
    "print(\n",
    "    \"neg_filter_continues\",\n",
    "    round(neg_filter_continues.sum() / total_negative, 4),\n",
    ")\n",
    "\n",
    "neg_filter_selection_mask = (\n",
    "    all_fitlers\n",
    "    & infill_duration_filter\n",
    "    & neg_filter_reaction_play_count\n",
    "    & neg_filter_upvote_count\n",
    "    & neg_filter_norm_play_frac\n",
    "    & neg_filter_continues\n",
    ")\n",
    "\n",
    "print(\"----------------\")\n",
    "total_positive = df[df[\"preference\"]].shape[0]\n",
    "assert total_positive == total_negative\n",
    "pos_filter_continues = (df[\"preference\"]) & (df[\"good_continue_at\"])\n",
    "print(\"pos_filter_continues\", round(pos_filter_continues.sum() / total_positive, 4))\n",
    "pos_filter_reaction_play_count = (df[\"preference\"]) & (df[\"reaction_play_count\"] >= 1)\n",
    "print(\n",
    "    \"pos_filter_reaction_play_count\",\n",
    "    round(pos_filter_reaction_play_count.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_relative_play_count = (df[\"preference\"]) & (df[\"play_rel_diff\"] >= 0)\n",
    "print(\n",
    "    \"pos_filter_relative_play_count\",\n",
    "    round(pos_filter_relative_play_count.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_cer_diff_preference = (\n",
    "    df[\n",
    "        \"preference\"\n",
    "    ]  # & (df[\"pos_diff_preference\"] == 2) # & (df[\"cer_diff_preference\"] < 0.5) & (df[\"cer\"] < 0.99)\n",
    ")\n",
    "print(\n",
    "    \"pos_filter_cer_diff_preference\",\n",
    "    round(pos_filter_cer_diff_preference.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_bad_flags = (\n",
    "    (df[\"preference\"]) & (df[\"flag_count\"] == 0) & (df[\"dislike_count\"] == 0)\n",
    ")\n",
    "print(\n",
    "    \"pos_filter_bad_flags\",\n",
    "    round(pos_filter_bad_flags.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_play_counts = (df[\"preference\"]) & (\n",
    "    (\n",
    "        (df[\"part_of_concat\"])\n",
    "        & (df[\"reaction_play_count\"] >= concat_pos_play_count)\n",
    "        & (df[\"concat_play_counts\"] >= concat_total_play_count)\n",
    "    )\n",
    "    | (\n",
    "        (~df[\"part_of_concat\"]) & (df[\"reaction_play_count\"] >= normal_pos_play_count)\n",
    "        # & (df[\"norm_play_frac\"] >= 2.1)  # this is a bit of a luxury cut...\n",
    "    )\n",
    "    | (df[\"task\"].isin([\"infill\", \"infill_intro\", \"infill_outro\"]))\n",
    ")\n",
    "print(\n",
    "    \"after filter on play counts\",\n",
    "    round(pos_filter_play_counts.sum() / total_positive, 4),\n",
    ")\n",
    "high_quality_tasks_filter = (\n",
    "    (\n",
    "        df[\"task\"].isin(\n",
    "            [\n",
    "                \"cover\",\n",
    "                \"upload_extend\",\n",
    "                \"cover_extend\",\n",
    "                \"artist_cover\",\n",
    "                \"artist_consistency\",\n",
    "                \"artist_extend\",\n",
    "                \"playlist_condition\",\n",
    "                \"overpainting\",\n",
    "                \"underpainting\",\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    & (\n",
    "        (df[\"upvote_count\"] >= 1)  # (df[\"upvote_count\"] >= 1)\n",
    "        | (df[\"reaction_play_count\"] >= 10)\n",
    "        | (df[\"concat_play_counts\"] >= 10)\n",
    "    )\n",
    "    & (\n",
    "        (df[\"part_of_concat\"])\n",
    "        | (\n",
    "            (~df[\"part_of_concat\"])\n",
    "            & (df[\"norm_play_frac\"] >= 5.1)  # this is a bit of a luxury cut...\n",
    "            & (\n",
    "                df[\"norm_play_frac\"] >= df[\"reaction_play_count\"] / 3\n",
    "            )  # play duration is not low on average\n",
    "        )\n",
    "    )\n",
    ")\n",
    "medium_quality_tasks_filter = (\n",
    "    (\n",
    "        df[\"task\"].isin(\n",
    "            [\n",
    "                \"extend\",\n",
    "                \"\",\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    & (\n",
    "        (df[\"upvote_count\"] >= 1)  # (df[\"upvote_count\"] >= 1)\n",
    "        | (df[\"reaction_play_count\"] >= 5)\n",
    "        | (df[\"concat_play_counts\"] >= 5)\n",
    "    )\n",
    "    & (\n",
    "        (df[\"part_of_concat\"])\n",
    "        | (\n",
    "            (~df[\"part_of_concat\"])\n",
    "            & (df[\"norm_play_frac\"] >= 2.1)  # this is a bit of a luxury cut...\n",
    "            & (\n",
    "                df[\"norm_play_frac\"] >= df[\"reaction_play_count\"] / 3\n",
    "            )  # play duration is not low on average\n",
    "        )\n",
    "    )\n",
    ")\n",
    "pos_filter_higher_quality = (df[\"preference\"]) & (\n",
    "    high_quality_tasks_filter | medium_quality_tasks_filter\n",
    ")\n",
    "print(\n",
    "    \"after filter on higher quality\",\n",
    "    round(pos_filter_higher_quality.sum() / total_positive, 4),\n",
    ")\n",
    "\n",
    "user_gen_filter = (\n",
    "    df[\"user_n_clips\"] >= 4\n",
    ")  # user needs to have genereated at least 100 over the time period\n",
    "\n",
    "clip_duration_filter = (df[\"preference\"]) & (\n",
    "    (df[\"duration_rel_diff\"] / df[\"duration\"] < 0.0) | (df[\"duration\"] <= 240)\n",
    ")\n",
    "print(\n",
    "    \"after filter on play duration\",\n",
    "    round(clip_duration_filter.sum() / total_positive, 4),\n",
    ")\n",
    "\n",
    "task_filter = (df[\"preference\"]) & (\n",
    "    df[\"task\"].isin(\n",
    "        [\n",
    "            \"\",\n",
    "            \"extend\",\n",
    "        ]\n",
    "    )\n",
    ")\n",
    "print(\n",
    "    \"after filter on task\",\n",
    "    round(task_filter.sum() / total_positive, 4),\n",
    ")\n",
    "\n",
    "cer_diff_filter = (df[\"preference\"]) & (df[\"cer_diff\"] < 0.20)\n",
    "print(\n",
    "    \"after filter on cer diff\",\n",
    "    round(cer_diff_filter.sum() / total_positive, 4),\n",
    ")\n",
    "print(\"----------------\")\n",
    "pos_filter_selectin_mask = (\n",
    "    (df[\"preference\"])  # get basics aligned\n",
    "    & all_fitlers\n",
    "    & infill_duration_filter\n",
    "    & pos_filter_continues\n",
    "    & pos_filter_reaction_play_count\n",
    "    & pos_filter_relative_play_count\n",
    "    & pos_filter_cer_diff_preference\n",
    "    & pos_filter_bad_flags\n",
    "    & pos_filter_play_counts\n",
    "    & pos_filter_higher_quality\n",
    "    & user_gen_filter\n",
    "    & clip_duration_filter\n",
    "    # & task_filter\n",
    "    & cer_diff_filter\n",
    ")\n",
    "print(\n",
    "    \"negative\",\n",
    "    sum(neg_filter_selection_mask),\n",
    "    \"positive\",\n",
    "    sum(pos_filter_selectin_mask),\n",
    ")\n",
    "\n",
    "neg_filter_requests = df[neg_filter_selection_mask][\"request_id\"].unique()\n",
    "pos_filter_requests = df[pos_filter_selectin_mask][\"request_id\"].unique()\n",
    "# looking for very strong signal here:\n",
    "# listen to the positive/negative more than once\n",
    "# disliked one of the clips\n",
    "unique_requests = set(pos_filter_requests).intersection(neg_filter_requests)\n",
    "print(\"----------------\")\n",
    "print(\n",
    "    \"total pair requests\",\n",
    "    df[\"request_id\"].nunique(),\n",
    "    \" --> selected pair requests\",\n",
    "    len(unique_requests),\n",
    "    f\"frac {len(unique_requests) / df['request_id'].nunique():.3f}\",\n",
    "    \" --> total intitial users\",\n",
    "    df[\"user_id\"].nunique(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.250737Z",
     "start_time": "2024-05-16T13:59:41.036434Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:22.559393Z",
     "iopub.status.busy": "2025-06-05T14:36:22.559111Z",
     "iopub.status.idle": "2025-06-05T14:36:22.853329Z",
     "shell.execute_reply": "2025-06-05T14:36:22.852571Z",
     "shell.execute_reply.started": "2025-06-05T14:36:22.559376Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "auk_t0_v20 requests 37064 clips 74128 total khrs 4.603; N gpus for 1000 iters 4.633; 4 gpus for x iters 386.083; n unique users 29012\n"
     ]
    }
   ],
   "source": [
    "df_slice = df[df[\"request_id\"].isin(set(unique_requests))].copy()\n",
    "print(\n",
    "    f\"{os.path.basename(OUT_DATA_DIR)} requests\",\n",
    "    df_slice[\"request_id\"].nunique(),\n",
    "    \"clips\",\n",
    "    df_slice.shape[0],\n",
    "    f\"total khrs {sum(df_slice['duration'] / 3600 / 1000):.3f};\",\n",
    "    f\"N gpus for 1000 iters {df_slice.shape[0] / 8 / 2 / 1000:.3f};\",\n",
    "    f\"4 gpus for x iters {df_slice.shape[0] / 8 / 4 / 6:.3f};\",\n",
    "    f\"n unique users {df_slice['user_id'].nunique()}\",\n",
    "    # f\"n pro users {df_slice[df_slice['is_pro_user']]['user_id'].nunique()}\",\n",
    ")\n",
    "# auk_mix_t1_v2 requests 102002 clips 204004 total khrs 9.191; N gpus for 1000 iters 12.750; 4 gpus for x iters 3187.562; n unique users 36408 n pro users 34038\n",
    "# auk_t1_v1 requests 9179 clips 18358 total khrs 0.854; N gpus for 1000 iters 1.147; 4 gpus for x iters 286.844; n unique users 6288 n pro users 6275\n",
    "# auk_t1_v2 requests 40903 clips 81806 total khrs 3.864; N gpus for 1000 iters 5.113; 4 gpus for x iters 1278.219; n unique users 21079 n pro users 20966\n",
    "# auk_t1_v3 requests 102015 clips 204030 total khrs 9.703; N gpus for 1000 iters 12.752; 4 gpus for x iters 3187.969; n unique users 42837 n pro users 42462\n",
    "# auk_t1_v4 requests 211452 clips 422904 total khrs 20.216; N gpus for 1000 iters 26.431; 4 gpus for x iters 6607.875; n unique users 71182 n pro users 70059\n",
    "# auk_t0_v1 requests 14775 clips 29550 total khrs 1.937; N gpus for 1000 iters 1.847; 4 gpus for x iters 461.719; n unique users 12832 n pro users 12637\n",
    "# auk_t0_v2 requests 9906 clips 19812 total khrs 1.306; N gpus for 1000 iters 1.238; 4 gpus for x iters 309.562; n unique users 8854 n pro users 8645\n",
    "# auk_t0_v3 requests 8293 clips 16586 total khrs 1.080; N gpus for 1000 iters 1.037; 4 gpus for x iters 259.156; n unique users 7537 n pro users 7464\n",
    "# auk_t0_v4 requests 26662 clips 53324 total khrs 3.505; N gpus for 1000 iters 3.333; 4 gpus for x iters 833.188; n unique users 21954 n pro users 21379\n",
    "# auk_t0_v5 requests 33141 clips 66282 total khrs 4.371; N gpus for 1000 iters 4.143; 4 gpus for x iters 1035.656; n unique users 26504 n pro users 25502\n",
    "# auk_t0_v6 requests 41330 clips 82660 total khrs 5.459; N gpus for 1000 iters 5.166; 4 gpus for x iters 1291.562; n unique users 31865 n pro users 30443\n",
    "# auk_t0_v7 requests 54058 clips 108116 total khrs 7.141; N gpus for 1000 iters 6.757; 4 gpus for x iters 1689.312; n unique users 40495 n pro users 38099\n",
    "# auk_t0_v8 requests 71198 clips 142396 total khrs 9.426; N gpus for 1000 iters 8.900; 4 gpus for x iters 2224.938; n unique users 51040 n pro users 48268\n",
    "# auk_t0_v10 requests 78123 clips 156246 total khrs 10.347; N gpus for 1000 iters 9.765; 4 gpus for x iters 2441.344; n unique users 54758 n pro users 51411\n",
    "# auk_t0_v11 requests 85259 clips 170518 total khrs 11.314; N gpus for 1000 iters 10.657; 4 gpus for x iters 2664.344; n unique users 59110 n pro users 55487\n",
    "# auk_t0_v12 requests 107970 clips 215940 total khrs 14.392; N gpus for 1000 iters 13.496; 4 gpus for x iters 3374.062; n unique users 72100 n pro users 67607\n",
    "# auk_t0_v13 requests 88625 clips 177250 total khrs 11.765; N gpus for 1000 iters 11.078; 4 gpus for x iters 2769.531; n unique users 61083 n pro users 57392\n",
    "# auk_t0_v14 refactor requests 27418 clips 54836 total khrs 3.359; N gpus for 1000 iters 3.427; 4 gpus for x iters 856.812; n unique users 22125\n",
    "# auk_t0_v14 requests 55934 clips 111868 total khrs 7.456; N gpus for 1000 iters 6.992; 4 gpus for x iters 1747.938; n unique users 34010\n",
    "# auk_t0_v15 requests 46654 clips 93308 total khrs 6.190; N gpus for 1000 iters 5.832; 4 gpus for x iters 1457.938; n unique users 35646\n",
    "# auk_t0_v16 requests 45946 clips 91892 total khrs 6.112; N gpus for 1000 iters 5.743; 4 gpus for x iters 1435.812; n unique users 33751\n",
    "# auk_t0_v17 requests 32253 clips 64506 total khrs 4.363; N gpus for 1000 iters 4.032; 4 gpus for x iters 1007.906; n unique users 24545\n",
    "# auk_t0_v18 requests 50600 clips 101200 total khrs 6.829; N gpus for 1000 iters 6.325; 4 gpus for x iters 527.083; n unique users 36764\n",
    "# auk_t0_v19 requests 37064 clips 74128 total khrs 4.603; N gpus for 1000 iters 4.633; 4 gpus for x iters 386.083; n unique users 29012"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.277006Z",
     "start_time": "2024-05-16T13:59:41.252105Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:22.854325Z",
     "iopub.status.busy": "2025-06-05T14:36:22.854043Z",
     "iopub.status.idle": "2025-06-05T14:36:22.894128Z",
     "shell.execute_reply": "2025-06-05T14:36:22.893495Z",
     "shell.execute_reply.started": "2025-06-05T14:36:22.854308Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive in playlist (8548, 183)\n",
      "Task and Public/Private Distribution:\n",
      "\n",
      "Task counts and fractions:\n",
      "task\n",
      "                      59114\n",
      "extend                 4836\n",
      "cover                  4674\n",
      "artist_consistency     3370\n",
      "artist_cover           1186\n",
      "upload_extend           480\n",
      "artist_extend           468\n",
      "Name: count, dtype: int64\n",
      "task\n",
      "                      0.797458\n",
      "extend                0.065239\n",
      "cover                 0.063053\n",
      "artist_consistency    0.045462\n",
      "artist_cover          0.015999\n",
      "upload_extend         0.006475\n",
      "artist_extend         0.006313\n",
      "Name: count, dtype: float64\n",
      "\n",
      "Public/Private counts and fractions:\n",
      "is_public\n",
      "False    69787\n",
      "True      4341\n",
      "Name: count, dtype: int64\n",
      "is_public\n",
      "False    0.941439\n",
      "True     0.058561\n",
      "Name: count, dtype: float64\n"
     ]
    }
   ],
   "source": [
    "test_mask = (df_slice[\"preference\"]) & (\n",
    "    (df_slice[\"is_in_playlist\"]) | (df_slice[\"concat_in_playlist\"])\n",
    ")\n",
    "print(\"positive in playlist\", df_slice[test_mask].shape)\n",
    "print(\"Task and Public/Private Distribution:\")\n",
    "task_counts = df_slice[\"task\"].value_counts()\n",
    "public_counts = df_slice[\"is_public\"].value_counts()\n",
    "\n",
    "print(\"\\nTask counts and fractions:\")\n",
    "print(task_counts)\n",
    "print(task_counts / len(df_slice))\n",
    "\n",
    "print(\"\\nPublic/Private counts and fractions:\")\n",
    "print(public_counts)\n",
    "print(public_counts / len(df_slice))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:22.915790Z",
     "iopub.status.busy": "2025-06-05T14:36:22.915502Z",
     "iopub.status.idle": "2025-06-05T14:36:22.944859Z",
     "shell.execute_reply": "2025-06-05T14:36:22.944332Z",
     "shell.execute_reply.started": "2025-06-05T14:36:22.915773Z"
    }
   },
   "outputs": [],
   "source": [
    "df_slice[\"npz_path\"] = df_slice[\"s3_id\"].map(lambda x: f\"{NPZ_DIR}/{x}.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T14:00:20.866354Z",
     "start_time": "2024-05-16T14:00:12.443344Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T14:36:22.945849Z",
     "iopub.status.busy": "2025-06-05T14:36:22.945449Z",
     "iopub.status.idle": "2025-06-05T14:36:23.148506Z",
     "shell.execute_reply": "2025-06-05T14:36:23.147816Z",
     "shell.execute_reply.started": "2025-06-05T14:36:22.945833Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(74128, 184)\n",
      "task\n",
      "                      59114\n",
      "extend                 4836\n",
      "cover                  4674\n",
      "artist_consistency     3370\n",
      "artist_cover           1186\n",
      "upload_extend           480\n",
      "artist_extend           468\n",
      "Name: count, dtype: int64\n"
     ]
    },
    {
     "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[26], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(df_slice\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28mprint\u001b[39m(df_slice[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtask\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mvalue_counts())\n\u001b[0;32m----> 6\u001b[0m \u001b[43mBREAK\u001b[49m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'BREAK' is not defined"
     ]
    }
   ],
   "source": [
    "# df_slice.to_pickle(\n",
    "#     \"/home/tony/Data/Preference/auk_t0/fully_merged_auk_t0_final_slice.pkl\"\n",
    "# )\n",
    "print(df_slice.shape)\n",
    "print(df_slice[\"task\"].value_counts())\n",
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Need to kick out the ones has gpt prompt -- these are pairs with different text inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\n",
    "#     \"Before filtering by user_id and task\",\n",
    "#     df_slice.shape[0],\n",
    "#     \"user_id unique:\",\n",
    "#     df_slice[\"user_id\"].nunique(),\n",
    "# )\n",
    "\n",
    "# # Create a copy to avoid fragmentation warning\n",
    "# df_slice = df_slice.copy()\n",
    "\n",
    "# # Calculate score for each row: reaction_play_count + 5 if preference is True, else 0\n",
    "# score_values = (\n",
    "#     df_slice[\"reaction_play_count\"] + (5 * df_slice[\"upvote_count\"].astype(int))\n",
    "# ) * df_slice[\"preference\"].astype(int)\n",
    "\n",
    "# # Use pd.concat to add the score column efficiently\n",
    "# df_slice = pd.concat(\n",
    "#     [df_slice, pd.DataFrame({\"score\": score_values}, index=df_slice.index)], axis=1\n",
    "# )\n",
    "\n",
    "# # Group by user_id and task, then for each group find the request_id with highest score\n",
    "# best_request_ids = []\n",
    "# for (user_id, task), group in tqdm(\n",
    "#     df_slice.groupby([\"user_id\", \"task\"]), desc=\"Processing user_id and task groups\"\n",
    "# ):\n",
    "#     # Get the request_id with the highest score in this group\n",
    "#     best_request_id = group.loc[group[\"score\"].idxmax(), \"request_id\"]\n",
    "#     best_request_ids.append(best_request_id)\n",
    "\n",
    "# # Filter df_slice to keep only the best request_ids for each user_id, task combination\n",
    "# df_slice = df_slice[df_slice[\"request_id\"].isin(best_request_ids)].copy()\n",
    "\n",
    "# print(\n",
    "#     \"After filtering by user_id and task\",\n",
    "#     df_slice.shape[0],\n",
    "#     \"user_id unique:\",\n",
    "#     df_slice[\"user_id\"].nunique(),\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\n",
    "#     \"/home/tony/Data/Preference/auk_t0/fully_merged_auk_t0_final_slice_successful_downloads.json\",\n",
    "#     \"r\",\n",
    "# ) as f:\n",
    "#     downloaded_clip_ids = json.load(f)\n",
    "\n",
    "# print(df_slice.shape, len(downloaded_clip_ids))\n",
    "# # trim off the tails so that we only keep the downloaded clip that still have audios\n",
    "# df_slice = df[df[\"s3_id\"].isin(set(downloaded_clip_ids))].copy()\n",
    "# print(df_slice.shape)\n",
    "# print(df_slice[\"task\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932296Z",
     "start_time": "2024-05-16T13:59:41.932287Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:30:57.442526Z",
     "iopub.status.busy": "2025-06-05T15:30:57.442109Z",
     "iopub.status.idle": "2025-06-05T15:30:58.387182Z",
     "shell.execute_reply": "2025-06-05T15:30:58.386533Z",
     "shell.execute_reply.started": "2025-06-05T15:30:57.442505Z"
    }
   },
   "outputs": [],
   "source": [
    "# don't have continue at\n",
    "df_slice[\"request_id\"] = df_slice[\"request_id\"].astype(str)\n",
    "# df_slice[\"npz_path\"] = df_slice[\"npz_path\"].apply(lambda x: str(x).replace(\"_npz\", \"_npz/\"))\n",
    "# df_slice[df_slice[\"continue_at\"].isna()][\"request_id\"].nunique(), df_slice[\"request_id\"].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T15:30:58.388363Z",
     "iopub.status.busy": "2025-06-05T15:30:58.388104Z",
     "iopub.status.idle": "2025-06-05T15:30:58.630112Z",
     "shell.execute_reply": "2025-06-05T15:30:58.629368Z",
     "shell.execute_reply.started": "2025-06-05T15:30:58.388346Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(74128, 184)\n",
      "(74128, 184)\n",
      "(74128, 184)\n"
     ]
    }
   ],
   "source": [
    "print(df_slice.shape)\n",
    "df_slice = df_slice[df_slice[\"request_id\"].apply(lambda x: len(x) > 3)]\n",
    "print(df_slice.shape)\n",
    "# df_slice = df_slice[df_slice[\"is_pro_user\"]].copy()\n",
    "print(df_slice.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932966Z",
     "start_time": "2024-05-16T13:59:41.932957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:30:58.631235Z",
     "iopub.status.busy": "2025-06-05T15:30:58.630851Z",
     "iopub.status.idle": "2025-06-05T15:30:58.657681Z",
     "shell.execute_reply": "2025-06-05T15:30:58.657070Z",
     "shell.execute_reply.started": "2025-06-05T15:30:58.631217Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37064\n"
     ]
    }
   ],
   "source": [
    "final_filtered_requests = df_slice[\"request_id\"].astype(str).unique()\n",
    "# final_filtered_requests = df_slice[df_slice[\"is_pro_user\"]][\"request_id\"].astype(str).unique()\n",
    "print(len(final_filtered_requests))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934277Z",
     "start_time": "2024-05-16T13:59:41.934268Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:30:58.659069Z",
     "iopub.status.busy": "2025-06-05T15:30:58.658807Z",
     "iopub.status.idle": "2025-06-05T15:30:58.987107Z",
     "shell.execute_reply": "2025-06-05T15:30:58.986401Z",
     "shell.execute_reply.started": "2025-06-05T15:30:58.659052Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36693 371\n",
      "(73386, 184) (742, 184)\n"
     ]
    }
   ],
   "source": [
    "train_requests, val_requests = train_test_split(\n",
    "    sorted(list(final_filtered_requests)), test_size=0.01, random_state=42\n",
    ")\n",
    "print(len(train_requests), len(val_requests))\n",
    "\n",
    "train_df = df_slice[df_slice[\"request_id\"].astype(str).isin(set(train_requests))].copy()\n",
    "val_df = df_slice[df_slice[\"request_id\"].astype(str).isin(set(val_requests))].copy()\n",
    "train_df = train_df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "train_df = train_df  # .reset_index()\n",
    "val_df = val_df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "val_df = val_df  # .reset_index()\n",
    "train_df = train_df.reset_index(drop=True)\n",
    "val_df = val_df.reset_index(drop=True)\n",
    "print(train_df.shape, val_df.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actually make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.935620Z",
     "start_time": "2024-05-16T13:59:41.935613Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:30:58.988343Z",
     "iopub.status.busy": "2025-06-05T15:30:58.987867Z",
     "iopub.status.idle": "2025-06-05T15:31:01.321350Z",
     "shell.execute_reply": "2025-06-05T15:31:01.320674Z",
     "shell.execute_reply.started": "2025-06-05T15:30:58.988323Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 73386/73386 [00:02<00:00, 29116.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4,557 hours of 73386 clips, 4.586625 nodes, 286.6640625 iters\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "total_duration = 0\n",
    "for i, row in tqdm(train_df.iterrows(), total=len(train_df)):\n",
    "    # we need to alternate between preference: neg, pos\n",
    "    # print(i, row)\n",
    "    try:\n",
    "        assert row[\"preference\"] == (i % 2 == 1)\n",
    "        total_duration += row[\"duration\"]\n",
    "    except Exception as E:\n",
    "        print(i, row)\n",
    "        print(E)\n",
    "        raise ValueError()\n",
    "\n",
    "print(\n",
    "    f\"{round(total_duration / 60 / 60):,} hours of {train_df.shape[0]} clips, {train_df.shape[0] / 8 / 2 / 1000} nodes, {train_df.shape[0] / 8 / 4 / 8} iters\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936268Z",
     "start_time": "2024-05-16T13:59:41.936260Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:31:01.322192Z",
     "iopub.status.busy": "2025-06-05T15:31:01.322028Z",
     "iopub.status.idle": "2025-06-05T15:31:08.288287Z",
     "shell.execute_reply": "2025-06-05T15:31:08.287725Z",
     "shell.execute_reply.started": "2025-06-05T15:31:01.322176Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t_data_memmap is set to: 12000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 742/742 [00:06<00:00, 110.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 742 clips, 0 different prompts, 0 different tags, 0 different negative tags\n",
      "28 hours of False\n",
      "24 hours of True\n",
      "cover: 4.7 hours\n",
      "gen: 39.2 hours\n",
      "artist_cover: 1.3 hours\n",
      "extend: 3.9 hours\n",
      "artist_consistency: 1.8 hours\n",
      "artist_extend: 1.5 hours\n",
      "\n",
      "--- Gender Distribution ---\n",
      "  unspecified: 742 (100.0%)\n",
      "\n",
      "--- Negative Tags Usage ---\n",
      "  has_neg_tags: 50 (6.7%)\n",
      "  no_neg_tags: 692 (93.3%)\n",
      "\n",
      "--- Control Slider Usage ---\n",
      "  has_control_slider: 162 (21.8% of clips)\n",
      "  no_control_slider: 580 (78.2% of clips)\n",
      "Done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "make_dataset(\n",
    "    val_df, OUT_DATA_DIR, is_val=True, npz_dir=NPZ_DIR, t_data_memmap=N_TOKENS_AUDIO\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T15:31:08.289058Z",
     "iopub.status.busy": "2025-06-05T15:31:08.288896Z",
     "iopub.status.idle": "2025-06-05T15:31:08.305755Z",
     "shell.execute_reply": "2025-06-05T15:31:08.305277Z",
     "shell.execute_reply.started": "2025-06-05T15:31:08.289040Z"
    }
   },
   "outputs": [],
   "source": [
    "# test_npz = np.load(\"/app/suno/data/dpo/30b_npz/26d19085-18da-4701-af43-122684543891.npz\")\n",
    "# for k in test_npz.keys():\n",
    "#     print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936964Z",
     "start_time": "2024-05-16T13:59:41.936957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:31:08.306426Z",
     "iopub.status.busy": "2025-06-05T15:31:08.306268Z",
     "iopub.status.idle": "2025-06-05T15:41:43.336389Z",
     "shell.execute_reply": "2025-06-05T15:41:43.335819Z",
     "shell.execute_reply.started": "2025-06-05T15:31:08.306410Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t_data_memmap is set to: 12000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 73386/73386 [10:46<00:00, 113.52it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 73386 clips, 11 different prompts, 1 different tags, 0 different negative tags\n",
      "2,801 hours of False\n",
      "2,419 hours of True\n",
      "gen: 3803.1 hours\n",
      "extend: 374.9 hours\n",
      "cover: 480.7 hours\n",
      "artist_cover: 141.4 hours\n",
      "artist_consistency: 369.4 hours\n",
      "artist_extend: 50.9 hours\n",
      "\n",
      "--- Gender Distribution ---\n",
      "  unspecified: 73,386 (100.0%)\n",
      "\n",
      "--- Negative Tags Usage ---\n",
      "  has_neg_tags: 5,110 (7.0%)\n",
      "  no_neg_tags: 68,276 (93.0%)\n",
      "\n",
      "--- Control Slider Usage ---\n",
      "  has_control_slider: 14,192 (19.3% of clips)\n",
      "  no_control_slider: 59,194 (80.7% of clips)\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "make_dataset(\n",
    "    train_df, OUT_DATA_DIR, is_val=False, npz_dir=NPZ_DIR, t_data_memmap=N_TOKENS_AUDIO\n",
    ")"
   ]
  },
  {
   "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": 37,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.937879Z",
     "start_time": "2024-05-16T13:59:41.937870Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:43.337171Z",
     "iopub.status.busy": "2025-06-05T15:41:43.337009Z",
     "iopub.status.idle": "2025-06-05T15:41:44.164497Z",
     "shell.execute_reply": "2025-06-05T15:41:44.163956Z",
     "shell.execute_reply.started": "2025-06-05T15:41:43.337153Z"
    }
   },
   "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, N_TOKENS_AUDIO, 1)\n",
    "assert len(mm) == len(test_metas)\n",
    "assert mm[:100, :, 0].min() >= 0\n",
    "assert mm[:100, :, 0].max() <= 4000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.166386Z",
     "iopub.status.busy": "2025-06-05T15:41:44.166144Z",
     "iopub.status.idle": "2025-06-05T15:41:44.181961Z",
     "shell.execute_reply": "2025-06-05T15:41:44.181482Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.166368Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({None: 598, 'extend': 60, 'cover': 44, 'artist_consistency': 14, 'artist_extend': 14, 'artist_cover': 12})\n"
     ]
    }
   ],
   "source": [
    "task_counts = Counter()\n",
    "for test_meta in test_metas:\n",
    "    task_counts[test_meta.get(\"task\")] += 1\n",
    "print(task_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.938629Z",
     "start_time": "2024-05-16T13:59:41.938621Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.182615Z",
     "iopub.status.busy": "2025-06-05T15:41:44.182465Z",
     "iopub.status.idle": "2025-06-05T15:41:44.195327Z",
     "shell.execute_reply": "2025-06-05T15:41:44.194899Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.182599Z"
    }
   },
   "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/data/dpo/models/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": 40,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939205Z",
     "start_time": "2024-05-16T13:59:41.939198Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.196030Z",
     "iopub.status.busy": "2025-06-05T15:41:44.195896Z",
     "iopub.status.idle": "2025-06-05T15:41:44.207526Z",
     "shell.execute_reply": "2025-06-05T15:41:44.207102Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.196016Z"
    }
   },
   "outputs": [],
   "source": [
    "# import random\n",
    "# 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(\"\\n negative example \\n\", test_metas[idx])\n",
    "# a.play(compress=False)\n",
    "# pos_a = decode(pos_arr)\n",
    "# print(\"\\n positive example \\n\", 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": 41,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939977Z",
     "start_time": "2024-05-16T13:59:41.939969Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.208113Z",
     "iopub.status.busy": "2025-06-05T15:41:44.207980Z",
     "iopub.status.idle": "2025-06-05T15:41:44.219428Z",
     "shell.execute_reply": "2025-06-05T15:41:44.218994Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.208100Z"
    }
   },
   "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": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.940610Z",
     "start_time": "2024-05-16T13:59:41.940603Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.220020Z",
     "iopub.status.busy": "2025-06-05T15:41:44.219887Z",
     "iopub.status.idle": "2025-06-05T15:41:44.231497Z",
     "shell.execute_reply": "2025-06-05T15:41:44.231067Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.220006Z"
    }
   },
   "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": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941167Z",
     "start_time": "2024-05-16T13:59:41.941159Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.232084Z",
     "iopub.status.busy": "2025-06-05T15:41:44.231953Z",
     "iopub.status.idle": "2025-06-05T15:41:44.243609Z",
     "shell.execute_reply": "2025-06-05T15:41:44.243086Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.232070Z"
    }
   },
   "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": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941801Z",
     "start_time": "2024-05-16T13:59:41.941793Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.244320Z",
     "iopub.status.busy": "2025-06-05T15:41:44.244183Z",
     "iopub.status.idle": "2025-06-05T15:41:44.257069Z",
     "shell.execute_reply": "2025-06-05T15:41:44.256604Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.244306Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "371 0\n"
     ]
    }
   ],
   "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": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.942520Z",
     "start_time": "2024-05-16T13:59:41.942511Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.257764Z",
     "iopub.status.busy": "2025-06-05T15:41:44.257623Z",
     "iopub.status.idle": "2025-06-05T15:41:44.273196Z",
     "shell.execute_reply": "2025-06-05T15:41:44.272753Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.257750Z"
    }
   },
   "outputs": [],
   "source": [
    "train_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_tr.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.943072Z",
     "start_time": "2024-05-16T13:59:41.943065Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.273803Z",
     "iopub.status.busy": "2025-06-05T15:41:44.273669Z",
     "iopub.status.idle": "2025-06-05T15:41:44.288031Z",
     "shell.execute_reply": "2025-06-05T15:41:44.287596Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.273789Z"
    }
   },
   "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)\n",
    "# make sure they are offset by 1 and exactly 1\n",
    "for i, j in zip(n_neg_tr, n_pos_tr):\n",
    "    assert i == j - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.944246Z",
     "start_time": "2024-05-16T13:59:41.944237Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.288633Z",
     "iopub.status.busy": "2025-06-05T15:41:44.288497Z",
     "iopub.status.idle": "2025-06-05T15:41:44.300353Z",
     "shell.execute_reply": "2025-06-05T15:41:44.299902Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.288619Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total samples 73386 (73386, 184) train_df_len 73386\n"
     ]
    }
   ],
   "source": [
    "total_iters = len(n_neg_tr) + len(n_pos_tr)\n",
    "print(\"total samples\", total_iters, train_df.shape, \"train_df_len\", len(train_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945249Z",
     "start_time": "2024-05-16T13:59:41.945241Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.300947Z",
     "iopub.status.busy": "2025-06-05T15:41:44.300814Z",
     "iopub.status.idle": "2025-06-05T15:41:44.312734Z",
     "shell.execute_reply": "2025-06-05T15:41:44.312282Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.300933Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 epoch per batch 4, total 458.6625\n"
     ]
    }
   ],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 5 / 8 / 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945972Z",
     "start_time": "2024-05-16T13:59:41.945964Z"
    },
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.313327Z",
     "iopub.status.busy": "2025-06-05T15:41:44.313187Z",
     "iopub.status.idle": "2025-06-05T15:41:44.324692Z",
     "shell.execute_reply": "2025-06-05T15:41:44.324168Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.313312Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm/bluejay && sbatch sbatch_ipo_bluejay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-06-05T15:41:44.325366Z",
     "iopub.status.busy": "2025-06-05T15:41:44.325232Z",
     "iopub.status.idle": "2025-06-05T15:41:44.344490Z",
     "shell.execute_reply": "2025-06-05T15:41:44.344032Z",
     "shell.execute_reply.started": "2025-06-05T15:41:44.325352Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cache kept!\n"
     ]
    }
   ],
   "source": [
    "import shutil\n",
    "\n",
    "# Basic file copy\n",
    "shutil.copy(\n",
    "    \"/home/tony/Work/tony/Preference/make_dataset_auk_t0.ipynb\",\n",
    "    os.path.join(OUT_DATA_DIR, \"make_dataset.ipynb\"),\n",
    ")\n",
    "print(\"Cache kept!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# some gymathtics loading prev data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}
