{
 "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": 21,
   "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_v19\"\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 (374496, 68)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\n",
    "    \"/home/tony/Data/Preference/auk_t0/interesting_clips_auk_t0_20250730_refactor.pkl\"\n",
    ")\n",
    "print(\"Preference data shape\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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_merged = pd.read_pickle(\"/home/tony/Data/Preference/auk_t0/fully_merged_auk_t0.pkl\")\n",
    "# # check how much df_merged is in df, and for cases not, fraction of is_deleted\n",
    "# # Check how many rows in df_merged are present in df (by 'id'), and for those not present, compute the fraction that are is_deleted\n",
    "\n",
    "# if \"id\" in df_merged.columns and \"id\" in df.columns:\n",
    "#     merged_ids = set(df_merged[\"id\"])\n",
    "#     df_ids = set(df[\"id\"])\n",
    "#     in_both = merged_ids & df_ids\n",
    "#     only_in_merged = merged_ids - df_ids\n",
    "\n",
    "#     print(f\"Total in df_merged: {len(merged_ids):,}\")\n",
    "#     print(f\"Total in df: {len(df_ids):,}\")\n",
    "#     print(f\"Rows in both: {len(in_both):,}\")\n",
    "#     print(f\"Rows only in df_merged: {len(only_in_merged):,}\")\n",
    "\n",
    "#     if only_in_merged:\n",
    "#         missing_df = df_merged[df_merged[\"id\"].isin(only_in_merged)]\n",
    "#         if \"is_deleted\" in missing_df.columns:\n",
    "#             frac_deleted = missing_df[\"is_deleted\"].mean()\n",
    "#             print(f\"Fraction of missing rows that are is_deleted: {frac_deleted:.4f}\")\n",
    "#             print(f\"Count is_deleted True: {missing_df['is_deleted'].sum()}, False: {(~missing_df['is_deleted']).sum()}\")\n",
    "#         else:\n",
    "#             print(\"No 'is_deleted' column in df_merged.\")\n",
    "#     else:\n",
    "#         print(\"No missing rows in df_merged compared to df.\")\n",
    "# else:\n",
    "#     print(\"No 'id' column in one of the dataframes.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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 (374496, 63)\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": 8,
   "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 (374496, 63)\n",
      "downloaded df (374496, 63)\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": 9,
   "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",
       "cover                 59338\n",
       "artist_consistency    34000\n",
       "                      31292\n",
       "extend                13094\n",
       "artist_cover          12536\n",
       "upload_extend          6596\n",
       "artist_extend          2086\n",
       "underpainting             2\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 9,
     "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": 10,
   "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    187248\n",
      "True        chirp-auk-t0    187248\n",
      "Name: count, dtype: int64\n",
      "before filter on model name (374496, 63)\n",
      "after filter on model name (374496, 63)\n",
      "is_public\n",
      "False    360519\n",
      "True      13977\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": 11,
   "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 (374496, 63)\n",
      "after filter on request id pairs (374496, 63)\n",
      "preference  model_name  \n",
      "False       chirp-auk-t0    187248\n",
      "True        chirp-auk-t0    187248\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": 12,
   "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 187248\n",
      "before removing duplicates (374496, 160)\n",
      "after removing duplicates (374496, 148)\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": 13,
   "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    134512\n",
       "2.0     52736\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 13,
     "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": 14,
   "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\"    2044\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": 15,
   "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 804 duplicated prompts 402 unique requests\n",
      "Found 171 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 374496\n",
      "After dedup user gen requests 374496\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": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.fillna({\"task\": \"\"}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.fillna({\"total_start_s\": 0}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count    374488.000000\n",
      "mean          2.191912\n",
      "std          28.723420\n",
      "min           0.000000\n",
      "25%           0.000000\n",
      "50%           0.000000\n",
      "75%           0.000000\n",
      "max        4129.592895\n",
      "Name: total_start_s, dtype: float64\n"
     ]
    }
   ],
   "source": [
    "# Ensure 'total_start_s' is numeric before describing\n",
    "df[\"total_start_s\"] = pd.to_numeric(df[\"total_start_s\"], errors=\"coerce\")\n",
    "print(df[\"total_start_s\"].describe())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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": [
      "10627\n",
      "good_continue_at\n",
      "True     374446\n",
      "False        50\n",
      "Name: count, dtype: int64\n",
      "\n",
      " Check some basics... \n",
      " preference\n",
      "False    187248\n",
      "True     187248\n",
      "Name: count, dtype: int64 model_name\n",
      "chirp-auk-t0    374496\n",
      "Name: count, dtype: int64 preference  model_name  \n",
      "False       chirp-auk-t0    187248\n",
      "True        chirp-auk-t0    187248\n",
      "Name: count, dtype: int64\n",
      "task\n",
      "                      246844\n",
      "cover                  59338\n",
      "artist_consistency     34000\n",
      "extend                 13094\n",
      "artist_cover           12536\n",
      "upload_extend           6596\n",
      "artist_extend           2086\n",
      "underpainting              2\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df[\"id\"] = df[\"str_id\"]\n",
    "df[\"total_start_s\"] = df[\"total_start_s\"].astype(float)\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.9990734213449542\n",
      "after infill duration 1.0\n",
      "neg_filter_reaction_play_count 1.0\n",
      "neg_filter_upvote_count 0.993\n"
     ]
    },
    {
     "ename": "KeyError",
     "evalue": "'norm_play_frac'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/indexes/base.py:3805\u001b[0m, in \u001b[0;36mIndex.get_loc\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m   3804\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 3805\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_loc\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcasted_key\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   3806\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
      "File \u001b[0;32mindex.pyx:167\u001b[0m, in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mindex.pyx:196\u001b[0m, in \u001b[0;36mpandas._libs.index.IndexEngine.get_loc\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mpandas/_libs/hashtable_class_helper.pxi:7081\u001b[0m, in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mpandas/_libs/hashtable_class_helper.pxi:7089\u001b[0m, in \u001b[0;36mpandas._libs.hashtable.PyObjectHashTable.get_item\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'norm_play_frac'",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[22], line 29\u001b[0m\n\u001b[1;32m     24\u001b[0m neg_filter_upvote_count \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m~\u001b[39mdf[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpreference\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;241m&\u001b[39m (df[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mupvote_count\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28mprint\u001b[39m(\n\u001b[1;32m     26\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mneg_filter_upvote_count\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     27\u001b[0m     \u001b[38;5;28mround\u001b[39m(neg_filter_upvote_count\u001b[38;5;241m.\u001b[39msum() \u001b[38;5;241m/\u001b[39m total_negative, \u001b[38;5;241m4\u001b[39m),\n\u001b[1;32m     28\u001b[0m )\n\u001b[0;32m---> 29\u001b[0m neg_filter_norm_play_frac \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m~\u001b[39mdf[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpreference\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;241m&\u001b[39m (\u001b[43mdf\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnorm_play_frac\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m3.1\u001b[39m)\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28mprint\u001b[39m(\n\u001b[1;32m     31\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mneg_filter_norm_play_frac\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     32\u001b[0m     \u001b[38;5;28mround\u001b[39m(neg_filter_norm_play_frac\u001b[38;5;241m.\u001b[39msum() \u001b[38;5;241m/\u001b[39m total_negative, \u001b[38;5;241m4\u001b[39m),\n\u001b[1;32m     33\u001b[0m )\n\u001b[1;32m     34\u001b[0m neg_filter_continues \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m~\u001b[39mdf[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpreference\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;241m&\u001b[39m (\n\u001b[1;32m     35\u001b[0m     df[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_continue_and_start_continue_at\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39misna()\n\u001b[1;32m     36\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/frame.py:4102\u001b[0m, in \u001b[0;36mDataFrame.__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m   4100\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcolumns\u001b[38;5;241m.\u001b[39mnlevels \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m   4101\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_getitem_multilevel(key)\n\u001b[0;32m-> 4102\u001b[0m indexer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcolumns\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_loc\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   4103\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m is_integer(indexer):\n\u001b[1;32m   4104\u001b[0m     indexer \u001b[38;5;241m=\u001b[39m [indexer]\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/indexes/base.py:3812\u001b[0m, in \u001b[0;36mIndex.get_loc\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m   3807\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(casted_key, \u001b[38;5;28mslice\u001b[39m) \u001b[38;5;129;01mor\u001b[39;00m (\n\u001b[1;32m   3808\u001b[0m         \u001b[38;5;28misinstance\u001b[39m(casted_key, abc\u001b[38;5;241m.\u001b[39mIterable)\n\u001b[1;32m   3809\u001b[0m         \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(x, \u001b[38;5;28mslice\u001b[39m) \u001b[38;5;28;01mfor\u001b[39;00m x \u001b[38;5;129;01min\u001b[39;00m casted_key)\n\u001b[1;32m   3810\u001b[0m     ):\n\u001b[1;32m   3811\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m InvalidIndexError(key)\n\u001b[0;32m-> 3812\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m(key) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[1;32m   3813\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[1;32m   3814\u001b[0m     \u001b[38;5;66;03m# If we have a listlike key, _check_indexing_error will raise\u001b[39;00m\n\u001b[1;32m   3815\u001b[0m     \u001b[38;5;66;03m#  InvalidIndexError. Otherwise we fall through and re-raise\u001b[39;00m\n\u001b[1;32m   3816\u001b[0m     \u001b[38;5;66;03m#  the TypeError.\u001b[39;00m\n\u001b[1;32m   3817\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_indexing_error(key)\n",
      "\u001b[0;31mKeyError\u001b[0m: 'norm_play_frac'"
     ]
    }
   ],
   "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",
    "play_duration_filter = (df[\"preference\"]) & (\n",
    "    df[\"duration_rel_diff\"] / df[\"duration\"] < 0.5\n",
    ")\n",
    "print(\n",
    "    \"after filter on play duration\",\n",
    "    round(play_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.25)\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",
    "    & play_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": [
    {
     "ename": "NameError",
     "evalue": "name 'unique_requests' 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[23], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m df_slice \u001b[38;5;241m=\u001b[39m df[df[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrequest_id\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39misin(\u001b[38;5;28mset\u001b[39m(\u001b[43munique_requests\u001b[49m))]\u001b[38;5;241m.\u001b[39mcopy()\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\n\u001b[1;32m      3\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mos\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mbasename(OUT_DATA_DIR)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m requests\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m      4\u001b[0m     df_slice[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrequest_id\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnunique(),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[38;5;66;03m# f\"n pro users {df_slice[df_slice['is_pro_user']]['user_id'].nunique()}\",\u001b[39;00m\n\u001b[1;32m     12\u001b[0m )\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# 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\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;66;03m# 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\u001b[39;00m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;66;03m# 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\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;66;03m# auk_t0_v14_refactor requests 64215 clips 128430 total khrs 8.523; N gpus for 1000 iters 8.027; 4 gpus for x iters 2006.719; n unique users 48442\u001b[39;00m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;66;03m# 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\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'unique_requests' is not defined"
     ]
    }
   ],
   "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 / 2 / 4:.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 64215 clips 128430 total khrs 8.523; N gpus for 1000 iters 8.027; 4 gpus for x iters 2006.719; n unique users 48442\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [],
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": [],
   "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": null,
   "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": [],
   "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": null,
   "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": [],
   "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": null,
   "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": [],
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": null,
   "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": null,
   "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": [],
   "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": null,
   "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": [],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 8 / 8 / 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": [],
   "source": [
    "import shutil\n",
    "\n",
    "# Basic file copy\n",
    "shutil.copy(\n",
    "    \"/home/tony/Work/tony/Preference/make_dataset_auk_t0_refactor.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
}
