{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:37:29.690267Z",
     "iopub.status.busy": "2025-11-14T03:37:29.690135Z",
     "iopub.status.idle": "2025-11-14T03:37:29.703226Z",
     "shell.execute_reply": "2025-11-14T03:37:29.702817Z",
     "shell.execute_reply.started": "2025-11-14T03:37:29.690252Z"
    }
   },
   "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-11-14T03:37:29.703806Z",
     "iopub.status.busy": "2025-11-14T03:37:29.703669Z",
     "iopub.status.idle": "2025-11-14T03:37:32.081806Z",
     "shell.execute_reply": "2025-11-14T03:37:32.081257Z",
     "shell.execute_reply.started": "2025-11-14T03:37:29.703792Z"
    }
   },
   "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-11-14T03:37:32.082541Z",
     "iopub.status.busy": "2025-11-14T03:37:32.082339Z",
     "iopub.status.idle": "2025-11-14T03:37:32.141274Z",
     "shell.execute_reply": "2025-11-14T03:37:32.140806Z",
     "shell.execute_reply.started": "2025-11-14T03:37:32.082525Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N_TOKENS_AUDIO 12000\n"
     ]
    }
   ],
   "source": [
    "OUT_DATA_DIR = \"/app2/suno/data/dpo/crow_t1_c1c2_v1\"\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/crow_t1_npz\"\n",
    "N_TOKENS_AUDIO = 25 * 8 * 60\n",
    "print(\"N_TOKENS_AUDIO\", N_TOKENS_AUDIO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:41:17.318765Z",
     "iopub.status.busy": "2025-11-14T03:41:17.318417Z",
     "iopub.status.idle": "2025-11-14T03:41:20.989731Z",
     "shell.execute_reply": "2025-11-14T03:41:20.989130Z",
     "shell.execute_reply.started": "2025-11-14T03:41:17.318747Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before dropna (132026, 93)\n",
      "after dropna (132026, 87)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\n",
    "    \"/home/tony/Data/Preference/crow_t1/interesting_clips_crow_t1_c12_20251113.pkl\"\n",
    ")\n",
    "print(\"before dropna\", df.shape)\n",
    "df = df.dropna(axis=1, how=\"all\")\n",
    "print(\"after dropna\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:41:20.990757Z",
     "iopub.status.busy": "2025-11-14T03:41:20.990597Z",
     "iopub.status.idle": "2025-11-14T03:41:21.010647Z",
     "shell.execute_reply": "2025-11-14T03:41:21.010202Z",
     "shell.execute_reply.started": "2025-11-14T03:41:20.990742Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "preference\n",
       "False    66013\n",
       "True     66013\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"preference\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.199480Z",
     "start_time": "2024-05-16T13:58:53.963687Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T03:41:21.011380Z",
     "iopub.status.busy": "2025-11-14T03:41:21.011232Z",
     "iopub.status.idle": "2025-11-14T03:49:02.541589Z",
     "shell.execute_reply": "2025-11-14T03:49:02.540792Z",
     "shell.execute_reply.started": "2025-11-14T03:41:21.011365Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18002649\n",
      "18002649\n",
      "pre-downloaded df (132026, 87)\n",
      "downloaded df (132026, 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": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.467253Z",
     "start_time": "2024-05-16T13:58:56.207647Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:02.542539Z",
     "iopub.status.busy": "2025-11-14T03:49:02.542366Z",
     "iopub.status.idle": "2025-11-14T03:49:03.432188Z",
     "shell.execute_reply": "2025-11-14T03:49:03.431504Z",
     "shell.execute_reply.started": "2025-11-14T03:49:02.542521Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "task\n",
       "                      73888\n",
       "cover                 37026\n",
       "artist_consistency     8216\n",
       "artist_cover           5148\n",
       "extend                 2788\n",
       "playlist_condition     2222\n",
       "upload_extend           816\n",
       "overpainting            638\n",
       "underpainting           458\n",
       "artist_extend           440\n",
       "infill                  386\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 8,
     "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": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.592883Z",
     "start_time": "2024-05-16T13:58:56.470781Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:03.434213Z",
     "iopub.status.busy": "2025-11-14T03:49:03.433685Z",
     "iopub.status.idle": "2025-11-14T03:49:04.283734Z",
     "shell.execute_reply": "2025-11-14T03:49:04.283056Z",
     "shell.execute_reply.started": "2025-11-14T03:49:03.434192Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preference  model_name      \n",
      "False       chirp-crow-t1       65980\n",
      "            chirp-crow-t1-c2       22\n",
      "            chirp-crow-t1-c1       11\n",
      "True        chirp-crow-t1-c2    37832\n",
      "            chirp-crow-t1-c1    28181\n",
      "Name: count, dtype: int64\n",
      "before filter on model name (132026, 87)\n",
      "after filter on model name (132026, 87)\n",
      "is_public\n",
      "False    128192\n",
      "True       3834\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-crow-t1\", \"chirp-crow-t2\", \"chirp-crow-t1-c1\", \"chirp-crow-t1-c2\"])]\n",
    "# df = df[df[\"model_name\"].isin([\"chirp-v3p5-engine-t-6\"])]\n",
    "print(\"after filter on model name\", df.shape)\n",
    "print(df[\"is_public\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.909539Z",
     "start_time": "2024-05-16T13:58:56.595736Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:04.284555Z",
     "iopub.status.busy": "2025-11-14T03:49:04.284394Z",
     "iopub.status.idle": "2025-11-14T03:49:04.520677Z",
     "shell.execute_reply": "2025-11-14T03:49:04.519997Z",
     "shell.execute_reply.started": "2025-11-14T03:49:04.284539Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before filter on request id pairs (132026, 87)\n",
      "after filter on request id pairs (132026, 87)\n",
      "preference  model_name      \n",
      "False       chirp-crow-t1       65980\n",
      "            chirp-crow-t1-c2       22\n",
      "            chirp-crow-t1-c1       11\n",
      "True        chirp-crow-t1-c2    37832\n",
      "            chirp-crow-t1-c1    28181\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": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:04.521513Z",
     "iopub.status.busy": "2025-11-14T03:49:04.521355Z",
     "iopub.status.idle": "2025-11-14T03:49:34.913226Z",
     "shell.execute_reply": "2025-11-14T03:49:34.912637Z",
     "shell.execute_reply.started": "2025-11-14T03:49:04.521497Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unique_requests 66013\n",
      "before removing duplicates (132026, 199)\n",
      "after removing duplicates (132026, 192)\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: custom_parse(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": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:34.913957Z",
     "iopub.status.busy": "2025-11-14T03:49:34.913802Z",
     "iopub.status.idle": "2025-11-14T03:49:35.168040Z",
     "shell.execute_reply": "2025-11-14T03:49:35.167539Z",
     "shell.execute_reply.started": "2025-11-14T03:49:34.913941Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pos_diff_preference\n",
       "1.0    47782\n",
       "2.0    18231\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 12,
     "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": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:35.168746Z",
     "iopub.status.busy": "2025-11-14T03:49:35.168600Z",
     "iopub.status.idle": "2025-11-14T03:49:35.269074Z",
     "shell.execute_reply": "2025-11-14T03:49:35.268547Z",
     "shell.execute_reply.started": "2025-11-14T03:49:35.168730Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive param_experiment\n",
      "cfg_steps_480    1\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": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:35.269798Z",
     "iopub.status.busy": "2025-11-14T03:49:35.269650Z",
     "iopub.status.idle": "2025-11-14T03:49:36.331832Z",
     "shell.execute_reply": "2025-11-14T03:49:36.331274Z",
     "shell.execute_reply.started": "2025-11-14T03:49:35.269782Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 456 duplicated prompts 228 unique requests\n",
      "Found 112 request_ids with duplicate prompts but not highest play counts in their group\n",
      "['c4a314d0-c30b-4f2e-b5e5-38a10d746bf4', 'd2ea14dd-918d-4340-b0e0-33da6883ff5d', 'a00c0b12-6d02-407c-9211-c23e436363f3', '7776efdc-a54f-4a45-a068-5a8860f64097', 'cba89e4d-de0a-457f-a07f-45d49fd39cf9', '998c470f-1e75-4286-8ad9-24d7d70acede', '655d2205-c639-48fe-bf1b-7b9b86ff061b', '38392fc2-06a1-497f-bfe2-afc17dd38e25', 'a841e2ca-ad25-4e5c-8f4a-2260abb792ad', '043059bf-85c0-4efc-bd62-5da415732363']\n",
      "Before dedup user gen requests 132026\n",
      "After dedup user gen requests 132026\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": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:40.799375Z",
     "start_time": "2024-05-16T13:59:36.394236Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T03:49:36.332570Z",
     "iopub.status.busy": "2025-11-14T03:49:36.332419Z",
     "iopub.status.idle": "2025-11-14T03:49:37.021453Z",
     "shell.execute_reply": "2025-11-14T03:49:37.020889Z",
     "shell.execute_reply.started": "2025-11-14T03:49:36.332554Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2183\n",
      "good_continue_at\n",
      "True     132013\n",
      "False        13\n",
      "Name: count, dtype: int64\n",
      "\n",
      " Check some basics... \n",
      " preference\n",
      "False    66013\n",
      "True     66013\n",
      "Name: count, dtype: int64 model_name\n",
      "chirp-crow-t1       65980\n",
      "chirp-crow-t1-c2    37854\n",
      "chirp-crow-t1-c1    28192\n",
      "Name: count, dtype: int64 preference  model_name      \n",
      "False       chirp-crow-t1       65980\n",
      "            chirp-crow-t1-c2       22\n",
      "            chirp-crow-t1-c1       11\n",
      "True        chirp-crow-t1-c2    37832\n",
      "            chirp-crow-t1-c1    28181\n",
      "Name: count, dtype: int64\n",
      "task\n",
      "                      73888\n",
      "cover                 37026\n",
      "artist_consistency     8216\n",
      "artist_cover           5148\n",
      "extend                 2788\n",
      "playlist_condition     2222\n",
      "upload_extend           816\n",
      "overpainting            638\n",
      "underpainting           458\n",
      "artist_extend           440\n",
      "infill                  386\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": 76,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.035167Z",
     "start_time": "2024-05-16T13:59:40.801098Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:16.106123Z",
     "iopub.status.busy": "2025-11-14T04:25:16.105791Z",
     "iopub.status.idle": "2025-11-14T04:25:16.543389Z",
     "shell.execute_reply": "2025-11-14T04:25:16.542786Z",
     "shell.execute_reply.started": "2025-11-14T04:25:16.106104Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "after duration 0.9962280156938785\n",
      "after infill duration 0.9999166830775756\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.9998\n",
      "pos_filter_reaction_play_count 1.0\n",
      "pos_filter_relative_play_count 0.9697\n",
      "pos_filter_cer_diff_preference 1.0\n",
      "pos_filter_bad_flags 0.9998\n",
      "after filter on play counts 0.9887\n",
      "after filter on higher quality 0.2724\n",
      "----------------\n",
      "negative 65274 positive 16618\n",
      "----------------\n",
      "total pair requests 66013  --> selected pair requests 16488 frac 0.250  --> total intitial users 52844\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) & (df[\"clip_type\"] != \"preview\")\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\"] <= 239)\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",
    "                \"extend\",\n",
    "                \"artist_consistency\",\n",
    "                \"artist_extend\",\n",
    "                \"playlist_condition\",\n",
    "                \"overpainting\",\n",
    "                \"underpainting\",\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\"] >= 3.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 -- huh only small difference with this cut\n",
    "        )\n",
    "    )\n",
    ")\n",
    "medium_quality_tasks_filter = (\n",
    "    df[\"task\"].isin(\n",
    "        [\n",
    "            \"infill\",\n",
    "            \"infill_intro\",\n",
    "            \"infill_outro\",\n",
    "            \"stem_condition\",\n",
    "        ]\n",
    "    )\n",
    ") & (  # let more infill through only in this case...\n",
    "    (\n",
    "        df[\"upvote_count\"] >= 1\n",
    "    )  # (df[\"upvote_count\"] >= 1)  (df[\"pos_diff_preference\"] == 2)\n",
    "    | (df[\"reaction_play_count\"] >= 1)\n",
    "    | (df[\"concat_play_counts\"] >= 1)\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\"] >= 100\n",
    ")  # user needs to have genereated at least 100 over the time period\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",
    ")\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": 77,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.250737Z",
     "start_time": "2024-05-16T13:59:41.036434Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:16.544572Z",
     "iopub.status.busy": "2025-11-14T04:25:16.544400Z",
     "iopub.status.idle": "2025-11-14T04:25:16.708000Z",
     "shell.execute_reply": "2025-11-14T04:25:16.707422Z",
     "shell.execute_reply.started": "2025-11-14T04:25:16.544556Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "task\n",
      "                      18072\n",
      "cover                  8340\n",
      "artist_consistency     2218\n",
      "extend                 1558\n",
      "artist_cover           1166\n",
      "playlist_condition      554\n",
      "infill                  334\n",
      "artist_extend           234\n",
      "upload_extend           208\n",
      "overpainting            180\n",
      "underpainting           112\n",
      "Name: count, dtype: int64\n",
      "crow_t1_c1c2_v1 requests 16488 clips 32976 total khrs 1.697; N gpus for 1000 iters 2.061; 4 gpus for x iters 515.250; n unique users 15152 n pro users 14720\n"
     ]
    }
   ],
   "source": [
    "df_slice = df[df[\"request_id\"].isin(set(unique_requests))].copy()\n",
    "print(df_slice[\"task\"].value_counts())\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_t1_v9 requests 547743 clips 1095486 total khrs 56.567; N gpus for 1000 iters 68.468; 4 gpus for x iters 17116.969; n unique users 129354 n pro users 125126\n",
    "# auk_t1_v13 requests 627646 clips 1255292 total khrs 64.530; N gpus for 1000 iters 78.456; 4 gpus for x iters 19613.938; n unique users 139873 n pro users 134738\n",
    "# auk_t1_v17 requests 494902 clips 989804 total khrs 51.445; N gpus for 1000 iters 61.863; 4 gpus for x iters 15465.688; n unique users 114240 n pro users 109708\n",
    "# auk_t1_v19 requests 740881 clips 1481762 total khrs 76.133; N gpus for 1000 iters 92.610; 4 gpus for x iters 23152.531; n unique users 154570 n pro users 147261\n",
    "# auk_t1_v24 requests 717745 clips 1435490 total khrs 74.515; N gpus for 1000 iters 89.718; 4 gpus for x iters 22429.531; n unique users 139344 n pro users 130740\n",
    "# auk_t1_v29 requests 1097586 clips 2195172 total khrs 112.614; N gpus for 1000 iters 137.198; 4 gpus for x iters 34299.562; n unique users 193138 n pro users 179429\n",
    "# auk_t1_v30 requests 1224422 clips 2448844 total khrs 125.330; N gpus for 1000 iters 153.053; 4 gpus for x iters 38263.188; n unique users 203138 n pro users 186758\n",
    "# auk_t1_v31 requests 389265 clips 778530 total khrs 40.761; N gpus for 1000 iters 48.658; 4 gpus for x iters 12164.531; n unique users 85198 n pro users 78445\n",
    "# auk_t1_v33 requests 1285261 clips 2570522 total khrs 131.585; N gpus for 1000 iters 160.658; 4 gpus for x iters 40164.406; n unique users 208932 n pro users 191380\n",
    "# auk_t1_v33 requests 1086639 clips 2173278 total khrs 110.641; N gpus for 1000 iters 135.830; 4 gpus for x iters 33957.469; n unique users 197293 n pro users 180968 -- play dur from /3 to /2\n",
    "# auk_t1_v33 requests 806877 clips 1613754 total khrs 81.964; N gpus for 1000 iters 100.860; 4 gpus for x iters 25214.906; n unique users 156679 n pro users 144253 -- filter to web\n",
    "# auk_t1_v37 requests 1339132 clips 2678264 total khrs 137.055; N gpus for 1000 iters 167.392; 4 gpus for x iters 41847.875; n unique users 214313 n pro users 196815\n",
    "# auk_t1_v38 requests 979451 clips 1958902 total khrs 102.038; N gpus for 1000 iters 122.431; 4 gpus for x iters 30607.844; n unique users 170773 n pro users 156737\n",
    "# auk_t1_v43 requests 184775 clips 369550 total khrs 19.089; N gpus for 1000 iters 23.097; 4 gpus for x iters 5774.219; n unique users 60652 n pro users 59126\n",
    "# auk_t1_v45 requests 1176216 clips 2352432 total khrs 122.102; N gpus for 1000 iters 147.027; 4 gpus for x iters 36756.750; n unique users 176751 n pro users 163097\n",
    "# auk_t1_v48 requests 294407 clips 588814 total khrs 29.992; N gpus for 1000 iters 36.801; 4 gpus for x iters 9200.219; n unique users 80991 n pro users 78825\n",
    "# bluejay_t1_v1 requests 47405 clips 94810 total khrs 5.515; N gpus for 1000 iters 5.926; 4 gpus for x iters 1481.406; n unique users 24344 n pro users 24194\n",
    "# bluejay_t1_v2 requests 101265 clips 202530 total khrs 11.753; N gpus for 1000 iters 12.658; 4 gpus for x iters 3164.531; n unique users 44574 n pro users 44024\n",
    "# bluejay_t1_v3 requests 144765 clips 289530 total khrs 16.756; N gpus for 1000 iters 18.096; 4 gpus for x iters 4523.906; n unique users 58421 n pro users 57449\n",
    "# bluejay_t1_v5 requests 179289 clips 358578 total khrs 20.764; N gpus for 1000 iters 22.411; 4 gpus for x iters 5602.781; n unique users 67850 n pro users 66572\n",
    "# bluejay_t1_v7 requests 282305 clips 564610 total khrs 32.692; N gpus for 1000 iters 35.288; 4 gpus for x iters 8822.031; n unique users 93023 n pro users 90936\n",
    "# bluejay_t1_v9 requests 540570 clips 1081140 total khrs 62.455; N gpus for 1000 iters 67.571; 4 gpus for x iters 16892.812; n unique users 145720 n pro users 142027\n",
    "# bluejay_t1_v11 requests 184433 clips 368866 total khrs 20.955; N gpus for 1000 iters 23.054; 4 gpus for x iters 5763.531; n unique users 52149 n pro users 51402\n",
    "# bluejay_t1_v12 requests 293821 clips 587642 total khrs 33.820; N gpus for 1000 iters 36.728; 4 gpus for x iters 9181.906; n unique users 98138 n pro users 95838\n",
    "# bluejay_t1_v13 requests 355678 clips 711356 total khrs 40.908; N gpus for 1000 iters 44.460; 4 gpus for x iters 11114.938; n unique users 111544 n pro users 108497\n",
    "# bluejay_t1_v17 requests 415411 clips 830822 total khrs 48.043; N gpus for 1000 iters 51.926; 4 gpus for x iters 12981.594; n unique users 124572 n pro users 120683\n",
    "# bluejay_t1_v24 requests 486884 clips 973768 total khrs 56.292; N gpus for 1000 iters 60.861; 4 gpus for x iters 15215.125; n unique users 138727 n pro users 133829\n",
    "# bluejay_t1_v26 requests 336220 clips 672440 total khrs 38.800; N gpus for 1000 iters 42.028; 4 gpus for x iters 10506.875; n unique users 80832 n pro users 78854\n",
    "# bluejay_t1_v28 requests 683709 clips 1367418 total khrs 78.507; N gpus for 1000 iters 85.464; 4 gpus for x iters 21365.906; n unique users 172776 n pro users 164584\n",
    "# bluejay_t1_v28 requests 662650 clips 1325300 total khrs 76.539; N gpus for 1000 iters 82.831; 4 gpus for x iters 20707.812; n unique users 170192 n pro users 162222\n",
    "# bluejay_t1_v31 requests 855246 clips 1710492 total khrs 97.846; N gpus for 1000 iters 106.906; 4 gpus for x iters 26726.438; n unique users 198438 n pro users 186068\n",
    "# bluejay_t1_v31 requests 829339 clips 1658678 total khrs 95.430; N gpus for 1000 iters 103.667; 4 gpus for x iters 25916.844; n unique users 195575 n pro users 183523\n",
    "# bluejay_t1_v31 requests 932058 clips 1864116 total khrs 107.136; N gpus for 1000 iters 116.507; 4 gpus for x iters 29126.812; n unique users 203738 n pro users 190639\n",
    "# bluejay_t1_v41 requests 1122075 clips 2244150 total khrs 128.936; N gpus for 1000 iters 140.259; 4 gpus for x iters 35064.844; n unique users 232511 n pro users 214225\n",
    "# bluejay_t1_v42 requests 1234273 clips 2468546 total khrs 140.767; N gpus for 1000 iters 154.284; 4 gpus for x iters 38571.031; n unique users 246967 n pro users 226020\n",
    "# crow_t1_v1 requests 131955 clips 263910 total khrs 12.518; N gpus for 1000 iters 16.494; 4 gpus for x iters 4123.594; n unique users 61682 n pro users 56836\n",
    "# crow_t1_v3 requests 272342 clips 544684 total khrs 26.003; N gpus for 1000 iters 34.043; 4 gpus for x iters 8510.688; n unique users 106845 n pro users 96483\n",
    "# crow_t1_v6 requests 369774 clips 739548 total khrs 35.356; N gpus for 1000 iters 46.222; 4 gpus for x iters 11555.438; n unique users 134281 n pro users 119919\n",
    "# crow_t1_v13 requests 558885 clips 1117770 total khrs 53.496; N gpus for 1000 iters 69.861; 4 gpus for x iters 17465.156; n unique users 182267 n pro users 160132\n",
    "# crow_t1_v17 requests 792420 clips 1584840 total khrs 77.740; N gpus for 1000 iters 99.052; 4 gpus for x iters 24763.125; n unique users 212279 n pro users 199173\n",
    "# crow_t1_v21 requests 1078035 clips 2156070 total khrs 106.094; N gpus for 1000 iters 134.754; 4 gpus for x iters 33688.594; n unique users 263938 n pro users 245247\n",
    "# crow_t1_v31 requests 1569448 clips 3138896 total khrs 155.285; N gpus for 1000 iters 196.181; 4 gpus for x iters 49045.250; n unique users 343002 n pro users 310320\n",
    "# crow_t1_v31 requests 1128871 clips 2257742 total khrs 113.348; N gpus for 1000 iters 141.109; 4 gpus for x iters 35277.219; n unique users 198445 n pro users 192068\n",
    "# crow_t1_vxx requests 1275697 clips 2551394 total khrs 128.287; N gpus for 1000 iters 159.462; 4 gpus for x iters 39865.531; n unique users 216731 n pro users 207174"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:16.708782Z",
     "iopub.status.busy": "2025-11-14T04:25:16.708629Z",
     "iopub.status.idle": "2025-11-14T04:25:16.728875Z",
     "shell.execute_reply": "2025-11-14T04:25:16.728422Z",
     "shell.execute_reply.started": "2025-11-14T04:25:16.708765Z"
    }
   },
   "outputs": [],
   "source": [
    "# import time\n",
    "# time.sleep(3600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:16.730077Z",
     "iopub.status.busy": "2025-11-14T04:25:16.729926Z",
     "iopub.status.idle": "2025-11-14T04:25:29.546373Z",
     "shell.execute_reply": "2025-11-14T04:25:29.545781Z",
     "shell.execute_reply.started": "2025-11-14T04:25:16.730062Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total hoot cer scores: 4600918\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [00:02<00:00, 688.32it/s]\n"
     ]
    }
   ],
   "source": [
    "# Load existing hoot CER cache\n",
    "with open(\"/home/tony/Data/Preference/crow_t1/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))\n",
    "# clip_id_to_cer = {}\n",
    "JSON_DIR = \"/app2/suno/data/dpo/crow_t1_json/\"\n",
    "\n",
    "# Extract unique s3_ids from df_slice that are not already in the cache\n",
    "clip_ids = set(df_slice[\"s3_id\"]) - set(clip_id_to_cer.keys())\n",
    "\n",
    "for clip_id in tqdm(clip_ids):\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(hoot_json_path, \"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",
    "# Update the hoot CER cache\n",
    "with open(\"/home/tony/Data/Preference/crow_t1/hoot_cer.json\", \"w\") as file:\n",
    "    json.dump(clip_id_to_cer, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:29.547222Z",
     "iopub.status.busy": "2025-11-14T04:25:29.547067Z",
     "iopub.status.idle": "2025-11-14T04:25:31.536960Z",
     "shell.execute_reply": "2025-11-14T04:25:31.536359Z",
     "shell.execute_reply.started": "2025-11-14T04:25:29.547205Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count    16488.000000\n",
      "mean         0.012069\n",
      "std          0.146108\n",
      "min         -1.000000\n",
      "10%         -0.080375\n",
      "50%          0.000000\n",
      "90%          0.107000\n",
      "95%          0.229124\n",
      "96%          0.285658\n",
      "97%          0.368460\n",
      "98%          0.500000\n",
      "99%          0.684369\n",
      "99.5%        0.800642\n",
      "99.9%        0.911256\n",
      "max          0.962025\n",
      "Name: cer_diff, dtype: float64\n",
      "crow_t1_c1c2_v1 requests 15854 clips 31708 total khrs 1.650; N gpus for 1000 iters 1.982; 4 gpus for x iters 495.438; n unique users 14617 n pro users 14201\n"
     ]
    }
   ],
   "source": [
    "# # add the cer to the df\n",
    "df_slice[\"cer\"] = df_slice[\"s3_id\"].map(clip_id_to_cer)\n",
    "df_slice = df_slice.fillna({\"cer\": 1})\n",
    "df_slice[\"cer_diff\"] = df_slice[\"cer\"].diff()\n",
    "df_slice = df_slice.fillna({\"cer_diff\": 0})\n",
    "print(df_slice[df_slice[\"preference\"]][\"cer_diff\"].describe(percentiles=[0.1, 0.5, 0.9, 0.95, 0.96, 0.97, 0.98, 0.99, 0.995, 0.999]))\n",
    "\n",
    "# just trimming off tail is fine and safe. Used to be 0.2. Multiple rounds now so probably 0.3 is still fine.\n",
    "cer_mask = (df_slice[\"preference\"]) & (df_slice[\"cer_diff\"] < 0.3)\n",
    "pos_cer_filter_requests = df_slice[cer_mask][\"request_id\"].unique()\n",
    "df_slice = df_slice[df_slice[\"request_id\"].isin(set(pos_cer_filter_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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.537711Z",
     "iopub.status.busy": "2025-11-14T04:25:31.537558Z",
     "iopub.status.idle": "2025-11-14T04:25:31.563078Z",
     "shell.execute_reply": "2025-11-14T04:25:31.562646Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.537696Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "source\n",
       "web        24868\n",
       "ios         3650\n",
       "android     3190\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_slice[\"source\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.277006Z",
     "start_time": "2024-05-16T13:59:41.252105Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.563691Z",
     "iopub.status.busy": "2025-11-14T04:25:31.563550Z",
     "iopub.status.idle": "2025-11-14T04:25:31.592417Z",
     "shell.execute_reply": "2025-11-14T04:25:31.591955Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.563676Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive in playlist (2649, 201)\n",
      "task\n",
      "                      17304\n",
      "cover                  8070\n",
      "artist_consistency     2162\n",
      "extend                 1492\n",
      "artist_cover           1150\n",
      "playlist_condition      534\n",
      "infill                  316\n",
      "artist_extend           220\n",
      "upload_extend           196\n",
      "overpainting            160\n",
      "underpainting           104\n",
      "Name: count, dtype: int64\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(df_slice[\"task\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.593050Z",
     "iopub.status.busy": "2025-11-14T04:25:31.592912Z",
     "iopub.status.idle": "2025-11-14T04:25:31.609361Z",
     "shell.execute_reply": "2025-11-14T04:25:31.608897Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.593036Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "is_public\n",
      "False    30756\n",
      "True       952\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(df_slice[\"is_public\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.609991Z",
     "iopub.status.busy": "2025-11-14T04:25:31.609853Z",
     "iopub.status.idle": "2025-11-14T04:25:31.630032Z",
     "shell.execute_reply": "2025-11-14T04:25:31.629594Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.609977Z"
    }
   },
   "outputs": [],
   "source": [
    "df_slice[\"npz_path\"] = df_slice[\"s3_id\"].map(lambda x: f\"{NPZ_DIR}/{x}.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.631980Z",
     "iopub.status.busy": "2025-11-14T04:25:31.631714Z",
     "iopub.status.idle": "2025-11-14T04:25:31.718933Z",
     "shell.execute_reply": "2025-11-14T04:25:31.718362Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.631965Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31708, 202)\n"
     ]
    }
   ],
   "source": [
    "df_total = df_slice.copy()\n",
    "print(df_total.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SOME SNOWFLAKE LYRICS SHIT YOU DON\"T WNAT TO KNOW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:31.719734Z",
     "iopub.status.busy": "2025-11-14T04:25:31.719581Z",
     "iopub.status.idle": "2025-11-14T04:25:32.369030Z",
     "shell.execute_reply": "2025-11-14T04:25:32.368530Z",
     "shell.execute_reply.started": "2025-11-14T04:25:31.719718Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-11-14 04:25:31,745 - INFO - Snowflake Connector for Python Version: 3.12.3, Python Version: 3.10.15, Platform: Linux-5.15.0-122-generic-x86_64-with-glibc2.35\n",
      "2025-11-14 04:25:31,746 - INFO - Connecting to GLOBAL Snowflake domain\n",
      "2025-11-14 04:25:31,746 - INFO - This connection is in OCSP Fail Open Mode. TLS Certificates would be checked for validity and revocation status. Any other Certificate Revocation related exceptions or OCSP Responder failures would be disregarded in favor of connectivity.\n",
      "2025-11-14 04:25:32,135 - INFO - Snowpark Session information: \n",
      "\"version\" : 1.23.0,\n",
      "\"python.version\" : 3.10.15,\n",
      "\"python.connector.version\" : 3.12.3,\n",
      "\"python.connector.session.id\" : 1572757012335726,\n",
      "\"os.name\" : Linux\n",
      "\n",
      "2025-11-14 04:25:32,136 - INFO - New root object was created for <snowflake.snowpark.session.Session object at 0x7f0f37daa380>\n",
      "2025-11-14 04:25:32,136 - INFO - performing a HTTP GET call to /api/v2/session/parameters/effective\n",
      "2025-11-14 04:25:32,366 - INFO - telemetry client created for <snowflake.connector.connection.SnowflakeConnection object at 0x7f0e3f4c9fc0>, telemetry enabled: True\n",
      "2025-11-14 04:25:32,366 - INFO - Snowflake Core version: 0.11.0, on Python 3.10.15, on platform: Linux-5.15.0-122-generic-x86_64-with-glibc2.35\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROD\n"
     ]
    }
   ],
   "source": [
    "home_dir = os.path.expanduser(\"~\")\n",
    "snow_password_path = os.path.join(home_dir, \".aws\", \"snow_pw.txt\")\n",
    "if os.path.exists(snow_password_path):\n",
    "    # !pip install snowflake\n",
    "    from snowflake.core import Root\n",
    "    from snowflake.snowpark import Session\n",
    "\n",
    "    with open(snow_password_path, \"r\") as fp:\n",
    "        fp_lines = fp.readlines()\n",
    "        snow_password = fp_lines[0].strip()\n",
    "        snow_username = fp_lines[1].strip()\n",
    "\n",
    "    CONNECTION_PARAMETERS = {\n",
    "        \"account\": \"fu90569.us-east-2.aws\",\n",
    "        \"user\": snow_username,\n",
    "        \"private_key_file\": \"/home/tony/.aws/rsa_key.p8\",\n",
    "        \"role\": \"ACCOUNTADMIN\",\n",
    "        \"database\": \"SUNO_PROD\",\n",
    "        \"warehouse\": \"SUNO_PROD_LARGE\",\n",
    "        \"schema\": \"PROD\",\n",
    "    }\n",
    "\n",
    "if not os.path.exists(snow_password_path):\n",
    "    raise Exception(\"you are not authorized to access snowflake -- please setup\")\n",
    "\n",
    "snow_session = Session.builder.configs(CONNECTION_PARAMETERS).create()\n",
    "\n",
    "snow_root = Root(snow_session)\n",
    "snow_schema = snow_root.databases[\"SUNO_PROD\"].schemas[\"PROD\"]\n",
    "print(snow_schema.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:25:32.369739Z",
     "iopub.status.busy": "2025-11-14T04:25:32.369593Z",
     "iopub.status.idle": "2025-11-14T04:26:30.312497Z",
     "shell.execute_reply": "2025-11-14T04:26:30.311691Z",
     "shell.execute_reply.started": "2025-11-14T04:25:32.369723Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No new clip IDs to query.\n",
      "Shape of df_snow_prompt:\n",
      "Rows: 4789542\n",
      "Columns: 2\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "from typing import List\n",
    "\n",
    "PROMPT_PATH = \"/home/tony/Data/Preference/crow_t1/clip_prompt.pkl\"\n",
    "\n",
    "# Load existing prompts if available\n",
    "df_existing_prompts = pd.read_pickle(PROMPT_PATH)\n",
    "df_existing_prompts[\"id\"] = df_existing_prompts[\"id\"].astype(str)\n",
    "\n",
    "df_total[\"id\"] = df_total[\"id\"].astype(str)\n",
    "# Only query for new clip ids not already in the prompt cache\n",
    "existing_ids = set(df_existing_prompts[\"id\"])\n",
    "all_clip_ids = set(df_total[\"id\"].unique())\n",
    "new_clip_ids = list(all_clip_ids - existing_ids)\n",
    "\n",
    "# df_total[\"id\"] = df_total[\"id\"].astype(str)\n",
    "# new_clip_ids = list(df_total[\"id\"].unique())\n",
    "\n",
    "snow_batch_size = 100_000\n",
    "snow_results: List[pd.DataFrame] = []\n",
    "\n",
    "if new_clip_ids:\n",
    "    for clip_ids_chunk in tqdm(\n",
    "        [new_clip_ids[i : i + snow_batch_size] for i in range(0, len(new_clip_ids), snow_batch_size)]\n",
    "    ):\n",
    "        id_query_str = \",\".join(\"'\" + x + \"'\" for x in clip_ids_chunk)\n",
    "        print(f\"Number of new clip IDs in this chunk: {len(clip_ids_chunk)}\")\n",
    "        print(f\"Length of the ID query string: {len(id_query_str)}\")\n",
    "\n",
    "        session_query = snow_session.sql(\n",
    "            f\"\"\"select ID, PROMPT_TEXT\n",
    "            from DDB_CLIP_META_HEAVY\n",
    "            where ID in ({id_query_str})\n",
    "            order by p_hour desc;\"\"\"\n",
    "        )\n",
    "        temp_df_snow = pd.DataFrame(session_query.collect())\n",
    "        snow_results.append(temp_df_snow)\n",
    "    print(f\"Number of new result batches: {len(snow_results)}\")\n",
    "    df_snow_new = pd.concat(snow_results, ignore_index=True)\n",
    "    df_snow_new = df_snow_new.rename(columns=lambda x: x.lower())\n",
    "    # Combine with existing prompts and drop duplicates (keep latest)\n",
    "    df_snow_prompt = pd.concat([df_existing_prompts, df_snow_new], ignore_index=True)\n",
    "    df_snow_prompt = df_snow_prompt.drop_duplicates(subset=[\"id\"], keep=\"last\")\n",
    "else:\n",
    "    print(\"No new clip IDs to query.\")\n",
    "    df_snow_prompt = df_existing_prompts\n",
    "\n",
    "# Save the updated prompt DataFrame every time\n",
    "df_snow_prompt.to_pickle(PROMPT_PATH)\n",
    "\n",
    "print(\"Shape of df_snow_prompt:\")\n",
    "print(f\"Rows: {df_snow_prompt.shape[0]}\")\n",
    "print(f\"Columns: {df_snow_prompt.shape[1]}\")\n",
    "\n",
    "df_total = df_total.rename(columns={'prompt_text': 'prompt_text_old'})\n",
    "df_total = df_total.merge(df_snow_prompt, on=\"id\", how=\"left\")\n",
    "# print((df_total[\"prompt_text\"] == df_total[\"prompt_text_old\"]).value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:30.313883Z",
     "iopub.status.busy": "2025-11-14T04:26:30.313347Z",
     "iopub.status.idle": "2025-11-14T04:26:31.856410Z",
     "shell.execute_reply": "2025-11-14T04:26:31.855644Z",
     "shell.execute_reply.started": "2025-11-14T04:26:30.313863Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crow_t1_c1c2_v1 requests 15850 clips 31700 total khrs 1.649; N gpus for 1000 iters 1.981; 4 gpus for x iters 495.312; n unique users 14615 n pro users 14199\n"
     ]
    }
   ],
   "source": [
    "identical_prompt_mask = (\n",
    "        df_total.groupby(\"request_id\")[\"prompt_text\"]\n",
    "        .nunique()\n",
    "        .eq(1)\n",
    "    )\n",
    "same_prompt_text_requests = set(identical_prompt_mask[identical_prompt_mask].index)\n",
    "df_slice = df_total[df_total[\"request_id\"].isin(set(same_prompt_text_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",
    "df_total = df_slice.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetch inference parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:31.857374Z",
     "iopub.status.busy": "2025-11-14T04:26:31.857203Z",
     "iopub.status.idle": "2025-11-14T04:26:39.462541Z",
     "shell.execute_reply": "2025-11-14T04:26:39.461753Z",
     "shell.execute_reply.started": "2025-11-14T04:26:31.857357Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "missing 0\n",
      "(173234, 71)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from fetch_gen_config import query_dynamodb_by_uuids_optimized\n",
    "\n",
    "DYNAMO_PATH = \"/home/tony/Data/Preference/crow_t1/clip_dynamo.pkl\"\n",
    "\n",
    "# Load existing DynamoDB records if available\n",
    "if os.path.exists(DYNAMO_PATH):\n",
    "    df_dynamo = pd.read_pickle(DYNAMO_PATH)\n",
    "    existing_ids = set(df_dynamo[\"clipId\"].tolist())\n",
    "else:\n",
    "    df_dynamo = pd.DataFrame()\n",
    "    existing_ids = set()\n",
    "\n",
    "total_s3_ids = set(df_total[\"id\"].tolist())\n",
    "missing_ids = list(total_s3_ids - existing_ids)\n",
    "print(\"missing\", len(missing_ids))\n",
    "chunk_size = 10_000\n",
    "all_records = []\n",
    "\n",
    "for i in range(0, len(missing_ids), chunk_size):\n",
    "    chunk = missing_ids[i:i + chunk_size]\n",
    "    try:\n",
    "        records = query_dynamodb_by_uuids_optimized(chunk, profile_name=\"default\")\n",
    "        all_records.extend(records)\n",
    "    except Exception as e:\n",
    "        # Log and continue with next chunk\n",
    "        print(f\"Error querying DynamoDB for chunk {i // chunk_size}: {e}\")\n",
    "\n",
    "if all_records:\n",
    "    df_new = pd.DataFrame(all_records)\n",
    "    df_dynamo = pd.concat([df_dynamo, df_new], ignore_index=True)\n",
    "\n",
    "print(df_dynamo.shape)\n",
    "df_dynamo[df_dynamo[\"type\"] == \"gpt\"].to_pickle(DYNAMO_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.463548Z",
     "iopub.status.busy": "2025-11-14T04:26:39.463368Z",
     "iopub.status.idle": "2025-11-14T04:26:39.486181Z",
     "shell.execute_reply": "2025-11-14T04:26:39.485617Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.463529Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_total.to_pickle(\"/home/tony/Data/Preference/bluejay_t1/interesting_clips_bluejay_t1_20250817_subset_total.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.486993Z",
     "iopub.status.busy": "2025-11-14T04:26:39.486837Z",
     "iopub.status.idle": "2025-11-14T04:26:39.505767Z",
     "shell.execute_reply": "2025-11-14T04:26:39.505254Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.486976Z"
    }
   },
   "outputs": [],
   "source": [
    "# OUT_DATA_DIR = \"/app2/suno/data/dpo/crow_t1_c1c2_v1\"\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/bluejay_t1_npz\"\n",
    "# N_TOKENS_AUDIO = 25 * 8 * 60\n",
    "# print(\"N_TOKENS_AUDIO\", N_TOKENS_AUDIO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.506674Z",
     "iopub.status.busy": "2025-11-14T04:26:39.506365Z",
     "iopub.status.idle": "2025-11-14T04:26:39.522409Z",
     "shell.execute_reply": "2025-11-14T04:26:39.521897Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.506658Z"
    }
   },
   "outputs": [],
   "source": [
    "# # # v3 # 09-23 -- 09-26 217218\n",
    "# # # v4 # 09-26 -- 09-28 166358\n",
    "# # # v5 # 09-28 -- 10-01 156902\n",
    "# # # \n",
    "# # # v6 # 09-23 -- 09-26 232690\n",
    "# # # v7 # 09-26 -- 09-28 184956\n",
    "# # # v8 # 09-28 -- 10-01 231470\n",
    "# # # \n",
    "# # # v9 (v6, v7, v8)\n",
    "# # # v10 # 09-23 -- 09-26 217344\n",
    "# # # v11 # 09-26 -- 09-28 167036\n",
    "# # # v12 # 09-28 -- 10-01 157570\n",
    "# # # \n",
    "# # # data 10-06\n",
    "# # # v13 # 09-23 -- 09-26 251008\n",
    "# # # v14 # 09-26 -- 09-28 205040\n",
    "# # # v15 # 09-28 -- 10-01 274238\n",
    "# # # v16 # 10-01 -- 10-04 246298\n",
    "# # # \n",
    "# # # data 10-09\n",
    "# # # v13 # 09-23 -- 09-26 260700\n",
    "# # # v14 # 09-26 -- 09-28 213860\n",
    "# # # v15 # 09-28 -- 10-01 289968\n",
    "# # # v16 # 10-01 -- 10-04 277330\n",
    "# # # \n",
    "# # # data 10-11\n",
    "# # # v17 # 09-23 -- 09-26 254828\n",
    "# # # v18 # 09-26 -- 09-29 304158\n",
    "# # # v19 # 09-29 -- 10-02 281620\n",
    "# # # v20 # 10-02 -- 10-05 274700\n",
    "# # # \n",
    "# # # data 10-16\n",
    "# # # v21 # 09-23 -- 09-26 266992\n",
    "# # # v22 # 09-26 -- 09-29 320004\n",
    "# # # v23 # 09-29 -- 10-02 299798\n",
    "# # # v24 # 10-02 -- 10-05 299228\n",
    "# # # v25 # 10-05 -- 10-08 283690\n",
    "# # # v26 # 10-08 -- 10-11 277962\n",
    "# # # v27 # 10-11 -- 10-15 320808\n",
    "# # # \n",
    "# # # data 10-25\n",
    "# # # v31 # 09-23 -- 09-26 200286\n",
    "# # # v32 # 09-26 -- 09-29 242158\n",
    "# # # v33 # 09-29 -- 10-02 230224\n",
    "# # # v34 # 10-02 -- 10-05 232714\n",
    "# # # v35 # 10-05 -- 10-11 450570\n",
    "# # # v36 # 10-11 -- 10-17 438742\n",
    "# # # v37 # 10-17 -- 10-26 446970\n",
    "# # # \n",
    "# # # data 10-25\n",
    "# # # v41 # 09-23 -- 09-26 135520\n",
    "# # # v42 # 09-26 -- 09-29 162702\n",
    "# # # v43 # 09-29 -- 10-02 152880\n",
    "# # # v44 # 10-02 -- 10-05 153712\n",
    "# # # v45 # 10-05 -- 10-11 292718\n",
    "# # # v46 # 10-11 -- 10-17 277006\n",
    "# # # v47 # 10-17 -- 10-26 253528\n",
    "# # # \n",
    "# # # data 10-30\n",
    "# # # v51 # 09-23 -- 09-26 209870\n",
    "# # # v52 # 09-26 -- 09-29 254476\n",
    "# # # v53 # 09-29 -- 10-02 242088\n",
    "# # # v54 # 10-02 -- 10-05 245902\n",
    "# # # v55 # 10-05 -- 10-12 566242\n",
    "# # # v56 # 10-12 -- 10-20 625112 \n",
    "# # # v57 # 10-20 -- 10-31 589974\n",
    "# # # data 11-04 (1st full)\n",
    "# # # v58 # 09-23 -- 10-26 462146 \n",
    "# # # data 11-08 (1st full)\n",
    "# # # v59 # 10-26 -- 11-09 370696 \n",
    "# df_total[\"created_at\"] = pd.to_datetime(df_total[\"created_at\"], utc=True)\n",
    "# start_cutoff_date = pd.to_datetime(\"2025-10-26\", utc=True)\n",
    "# end_cutoff_date = pd.to_datetime(\"2025-11-09\", utc=True)\n",
    "# date_mask = (df_total[\"created_at\"] < end_cutoff_date) & (df_total[\"created_at\"] >= start_cutoff_date)\n",
    "# print(df_total.shape, df_total[date_mask].shape)\n",
    "# df_slice = df_total[date_mask].copy()\n",
    "# print(\"after date cut\", df_slice.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.523230Z",
     "iopub.status.busy": "2025-11-14T04:26:39.523087Z",
     "iopub.status.idle": "2025-11-14T04:26:39.538043Z",
     "shell.execute_reply": "2025-11-14T04:26:39.537530Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.523216Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice = df_total.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.538881Z",
     "iopub.status.busy": "2025-11-14T04:26:39.538738Z",
     "iopub.status.idle": "2025-11-14T04:26:39.553934Z",
     "shell.execute_reply": "2025-11-14T04:26:39.553422Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.538867Z"
    }
   },
   "outputs": [],
   "source": [
    "# from tqdm import tqdm\n",
    "\n",
    "# list_of_past_data = [\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v51\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v52\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v53\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v54\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v55\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v56\",\n",
    "#     \"/app2/suno/data/dpo/crow_t1_v57\",     \n",
    "# ]\n",
    "\n",
    "# # Collect all known train IDs from previous meta_tr.jsonl files, showing progress and set growth\n",
    "# known_train_ids = set()\n",
    "# for data_dir in tqdm(list_of_past_data, desc=\"Collecting known train IDs\"):\n",
    "#     metas = read_jsonl(os.path.join(data_dir, \"meta_tr.jsonl\"))\n",
    "#     before = len(known_train_ids)\n",
    "#     known_train_ids.update(meta[\"id\"] for meta in metas)\n",
    "#     after = len(known_train_ids)\n",
    "#     tqdm.write(f\"Added {after - before} new IDs from {data_dir} (total: {after})\")\n",
    "\n",
    "# print(f\"Total known train IDs: {len(known_train_ids)}\")\n",
    "# print(f\"Original df_slice shape: {df_slice.shape}\")\n",
    "# df_slice = df_slice[~df_slice[\"id\"].isin(known_train_ids)].copy()\n",
    "# print(f\"Filtered df_slice shape: {df_slice.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.554698Z",
     "iopub.status.busy": "2025-11-14T04:26:39.554555Z",
     "iopub.status.idle": "2025-11-14T04:26:39.569439Z",
     "shell.execute_reply": "2025-11-14T04:26:39.568926Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.554684Z"
    }
   },
   "outputs": [],
   "source": [
    "# from suno_analytics.preference_data_selection import plot_clip_distribution\n",
    "# plot_clip_distribution(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.570178Z",
     "iopub.status.busy": "2025-11-14T04:26:39.570033Z",
     "iopub.status.idle": "2025-11-14T04:26:39.584968Z",
     "shell.execute_reply": "2025-11-14T04:26:39.584455Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.570164Z"
    }
   },
   "outputs": [],
   "source": [
    "# from suno_analytics.preference_data_selection import plot_clip_distribution\n",
    "# plot_clip_distribution(df_total)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.585784Z",
     "iopub.status.busy": "2025-11-14T04:26:39.585639Z",
     "iopub.status.idle": "2025-11-14T04:26:39.600847Z",
     "shell.execute_reply": "2025-11-14T04:26:39.600338Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.585770Z"
    }
   },
   "outputs": [],
   "source": [
    "# from collections import Counter, defaultdict\n",
    "# from typing import List, Dict, Set\n",
    "\n",
    "# # Read all meta files\n",
    "# meta_files: List[str] = [\n",
    "#     \"/app/suno/data/dpo/30b_t1_v23/meta_tr.jsonl\",\n",
    "#     \"/app/suno/data/dpo/30b_t6_v35/meta_tr.jsonl\",\n",
    "#     \"/app/suno/data/dpo/13b_s32_v34/meta_tr.jsonl\",\n",
    "#     \"/app/suno/data/dpo/auk_mix_t1_v6/meta_tr.jsonl\",\n",
    "#     \"/app/suno/data/dpo/auk_mix_t1_v14/meta_tr.jsonl\",\n",
    "#     \"/app2/suno/data/dpo/auk_t1_v6/meta_tr.jsonl\",\n",
    "#     \"/app2/suno/data/dpo/auk_t1_v7/meta_tr.jsonl\",\n",
    "#     \"/app2/suno/data/dpo/auk_t1_v19/meta_tr.jsonl\",\n",
    "#     \"/app2/suno/data/dpo/auk_t1_v29/meta_tr.jsonl\",\n",
    "# ]\n",
    "\n",
    "# user_id_counter: Counter[str] = Counter()\n",
    "# for meta_file in tqdm(meta_files):\n",
    "#     metas = read_jsonl(meta_file)\n",
    "#     # Count each user_id once per dataset\n",
    "#     user_ids: Set[str] = {meta[\"user_id\"] for meta in metas if \"user_id\" in meta}\n",
    "#     user_id_counter.update(user_ids)\n",
    "\n",
    "# # Map from count to set of user_ids\n",
    "# count_to_users: Dict[int, Set[str]] = defaultdict(set)\n",
    "# for user_id, count in user_id_counter.items():\n",
    "#     count_to_users[count].add(user_id)\n",
    "\n",
    "# for count in sorted(count_to_users):\n",
    "#     users = count_to_users[count]\n",
    "#     print(f\"Count {count}: {len(users)} users\")\n",
    "\n",
    "# # with open(\"/home/tony/Data/top_user_8.json\", \"w\") as fp:\n",
    "# #     json.dump(list(count_to_users[8]) , fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.601785Z",
     "iopub.status.busy": "2025-11-14T04:26:39.601408Z",
     "iopub.status.idle": "2025-11-14T04:26:39.616329Z",
     "shell.execute_reply": "2025-11-14T04:26:39.615818Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.601769Z"
    }
   },
   "outputs": [],
   "source": [
    "# with open(\"/home/tony/Data/top_user/mask_control_low.json\", \"r\") as fp:\n",
    "#     very_good_users = json.load(fp)\n",
    "# # very_good_users = count_to_users[5]\n",
    "# very_good_users_mask = df_total[\"user_id\"].isin(very_good_users)\n",
    "# print(df_total[very_good_users_mask].shape, df_total.shape)\n",
    "# print(df_total[very_good_users_mask][\"user_id\"].nunique(), df_total[\"user_id\"].nunique())\n",
    "# # df_total[very_good_users_mask][\"task\"].value_counts()\n",
    "# df_slice = df_total[very_good_users_mask].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.617276Z",
     "iopub.status.busy": "2025-11-14T04:26:39.616899Z",
     "iopub.status.idle": "2025-11-14T04:26:39.632176Z",
     "shell.execute_reply": "2025-11-14T04:26:39.631667Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.617261Z"
    }
   },
   "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": 100,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T14:00:20.866354Z",
     "start_time": "2024-05-16T14:00:12.443344Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:26:39.633055Z",
     "iopub.status.busy": "2025-11-14T04:26:39.632757Z",
     "iopub.status.idle": "2025-11-14T04:26:39.674479Z",
     "shell.execute_reply": "2025-11-14T04:26:39.671932Z",
     "shell.execute_reply.started": "2025-11-14T04:26:39.633039Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31700, 203)\n",
      "task\n",
      "                      17300\n",
      "cover                  8070\n",
      "artist_consistency     2160\n",
      "extend                 1490\n",
      "artist_cover           1150\n",
      "playlist_condition      534\n",
      "infill                  316\n",
      "artist_extend           220\n",
      "upload_extend           196\n",
      "overpainting            160\n",
      "underpainting           104\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[100], 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[1;32m      7\u001b[0m \u001b[38;5;66;03m# (269282, 199)\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'BREAK' is not defined"
     ]
    }
   ],
   "source": [
    "# df_slice.to_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250426_full_long_slice.pkl\"\n",
    "# )\n",
    "print(df_slice.shape)\n",
    "print(df_slice[\"task\"].value_counts())\n",
    "BREAK\n",
    "# (269282, 199)"
   ]
  },
  {
   "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": 101,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:44.633839Z",
     "iopub.status.busy": "2025-11-14T04:28:44.633398Z",
     "iopub.status.idle": "2025-11-14T04:28:45.905013Z",
     "shell.execute_reply": "2025-11-14T04:28:45.904231Z",
     "shell.execute_reply.started": "2025-11-14T04:28:44.633818Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(31700, 203)\n",
      "(31700, 203)\n",
      "(31700, 203)\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": 102,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932296Z",
     "start_time": "2024-05-16T13:59:41.932287Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:45.906441Z",
     "iopub.status.busy": "2025-11-14T04:28:45.906150Z",
     "iopub.status.idle": "2025-11-14T04:28:45.930677Z",
     "shell.execute_reply": "2025-11-14T04:28:45.930108Z",
     "shell.execute_reply.started": "2025-11-14T04:28:45.906422Z"
    }
   },
   "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": 103,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932966Z",
     "start_time": "2024-05-16T13:59:41.932957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:45.931565Z",
     "iopub.status.busy": "2025-11-14T04:28:45.931414Z",
     "iopub.status.idle": "2025-11-14T04:28:45.956671Z",
     "shell.execute_reply": "2025-11-14T04:28:45.956118Z",
     "shell.execute_reply.started": "2025-11-14T04:28:45.931550Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15850\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": 104,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934277Z",
     "start_time": "2024-05-16T13:59:41.934268Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:45.957546Z",
     "iopub.status.busy": "2025-11-14T04:28:45.957255Z",
     "iopub.status.idle": "2025-11-14T04:28:46.215114Z",
     "shell.execute_reply": "2025-11-14T04:28:46.214354Z",
     "shell.execute_reply.started": "2025-11-14T04:28:45.957532Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15691 159\n",
      "(31382, 203) (318, 203)\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": 105,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.935620Z",
     "start_time": "2024-05-16T13:59:41.935613Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:46.216923Z",
     "iopub.status.busy": "2025-11-14T04:28:46.216603Z",
     "iopub.status.idle": "2025-11-14T04:28:47.367638Z",
     "shell.execute_reply": "2025-11-14T04:28:47.366894Z",
     "shell.execute_reply.started": "2025-11-14T04:28:46.216904Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 31382/31382 [00:01<00:00, 27899.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1,633 hours of 31382 clips, 1.961375 nodes, 40.861979166666664 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 / 6 / 16} iters\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936268Z",
     "start_time": "2024-05-16T13:59:41.936260Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:47.369103Z",
     "iopub.status.busy": "2025-11-14T04:28:47.368924Z",
     "iopub.status.idle": "2025-11-14T04:28:55.522228Z",
     "shell.execute_reply": "2025-11-14T04:28:55.521519Z",
     "shell.execute_reply.started": "2025-11-14T04:28:47.369086Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t_data_memmap is set to: 12000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 318/318 [00:08<00:00, 39.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 318 clips, 0 different prompts, 0 different tags, 0 different negative tags\n",
      "11 hours of False\n",
      "12 hours of True\n",
      "gen: 9.2 hours\n",
      "cover: 8.1 hours\n",
      "infill: 0.1 hours\n",
      "extend: 0.6 hours\n",
      "artist_consistency: 2.3 hours\n",
      "artist_cover: 1.6 hours\n",
      "artist_extend: 0.5 hours\n",
      "playlist_condition: 0.4 hours\n",
      "\n",
      "--- Gender Distribution ---\n",
      "  female: 34 (10.7%)\n",
      "  male: 52 (16.4%)\n",
      "  unspecified: 232 (73.0%)\n",
      "\n",
      "--- Negative Tags Usage ---\n",
      "  has_neg_tags: 22 (6.9%)\n",
      "  no_neg_tags: 296 (93.1%)\n",
      "\n",
      "--- Control Slider Usage ---\n",
      "  has_control_slider: 110 (34.6% of clips)\n",
      "  no_control_slider: 208 (65.4% 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": 107,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:55.523216Z",
     "iopub.status.busy": "2025-11-14T04:28:55.522945Z",
     "iopub.status.idle": "2025-11-14T04:28:55.545245Z",
     "shell.execute_reply": "2025-11-14T04:28:55.544677Z",
     "shell.execute_reply.started": "2025-11-14T04:28:55.523198Z"
    }
   },
   "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": 108,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936964Z",
     "start_time": "2024-05-16T13:59:41.936957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:28:55.546184Z",
     "iopub.status.busy": "2025-11-14T04:28:55.545878Z",
     "iopub.status.idle": "2025-11-14T04:36:13.829147Z",
     "shell.execute_reply": "2025-11-14T04:36:13.828606Z",
     "shell.execute_reply.started": "2025-11-14T04:28:55.546168Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t_data_memmap is set to: 12000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 31382/31382 [07:18<00:00, 71.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 31382 clips, 0 different prompts, 0 different tags, 0 different negative tags\n",
      "1,076 hours of False\n",
      "1,161 hours of True\n",
      "gen: 936.3 hours\n",
      "artist_consistency: 220.4 hours\n",
      "cover: 759.7 hours\n",
      "artist_cover: 129.1 hours\n",
      "infill: 6.1 hours\n",
      "artist_extend: 20.7 hours\n",
      "extend: 97.8 hours\n",
      "playlist_condition: 45.0 hours\n",
      "underpainting: 8.9 hours\n",
      "overpainting: 12.9 hours\n",
      "\n",
      "--- Gender Distribution ---\n",
      "  female: 3,402 (10.8%)\n",
      "  male: 5,074 (16.2%)\n",
      "  unspecified: 22,906 (73.0%)\n",
      "\n",
      "--- Negative Tags Usage ---\n",
      "  has_neg_tags: 2,150 (6.9%)\n",
      "  no_neg_tags: 29,232 (93.1%)\n",
      "\n",
      "--- Control Slider Usage ---\n",
      "  has_control_slider: 10,586 (33.7% of clips)\n",
      "  no_control_slider: 20,796 (66.3% of clips)\n",
      "Done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\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": 109,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.937879Z",
     "start_time": "2024-05-16T13:59:41.937870Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:13.829894Z",
     "iopub.status.busy": "2025-11-14T04:36:13.829741Z",
     "iopub.status.idle": "2025-11-14T04:36:15.006350Z",
     "shell.execute_reply": "2025-11-14T04:36:15.005848Z",
     "shell.execute_reply.started": "2025-11-14T04:36:13.829879Z"
    }
   },
   "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": 110,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.007057Z",
     "iopub.status.busy": "2025-11-14T04:36:15.006908Z",
     "iopub.status.idle": "2025-11-14T04:36:15.024979Z",
     "shell.execute_reply": "2025-11-14T04:36:15.024525Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.007041Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({None: 164, 'cover': 92, 'artist_consistency': 22, 'artist_cover': 14, 'extend': 12, 'infill': 6, 'artist_extend': 4, 'playlist_condition': 4})\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": 111,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.938629Z",
     "start_time": "2024-05-16T13:59:41.938621Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.025612Z",
     "iopub.status.busy": "2025-11-14T04:36:15.025475Z",
     "iopub.status.idle": "2025-11-14T04:36:15.040541Z",
     "shell.execute_reply": "2025-11-14T04:36:15.040107Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.025599Z"
    }
   },
   "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": 112,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939205Z",
     "start_time": "2024-05-16T13:59:41.939198Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.041179Z",
     "iopub.status.busy": "2025-11-14T04:36:15.041036Z",
     "iopub.status.idle": "2025-11-14T04:36:15.056056Z",
     "shell.execute_reply": "2025-11-14T04:36:15.055626Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.041165Z"
    }
   },
   "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": 113,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939977Z",
     "start_time": "2024-05-16T13:59:41.939969Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.056828Z",
     "iopub.status.busy": "2025-11-14T04:36:15.056687Z",
     "iopub.status.idle": "2025-11-14T04:36:15.071317Z",
     "shell.execute_reply": "2025-11-14T04:36:15.070891Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.056814Z"
    }
   },
   "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": 114,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.940610Z",
     "start_time": "2024-05-16T13:59:41.940603Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.073477Z",
     "iopub.status.busy": "2025-11-14T04:36:15.073320Z",
     "iopub.status.idle": "2025-11-14T04:36:15.088204Z",
     "shell.execute_reply": "2025-11-14T04:36:15.087767Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.073463Z"
    }
   },
   "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": 115,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941167Z",
     "start_time": "2024-05-16T13:59:41.941159Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.088798Z",
     "iopub.status.busy": "2025-11-14T04:36:15.088662Z",
     "iopub.status.idle": "2025-11-14T04:36:15.103441Z",
     "shell.execute_reply": "2025-11-14T04:36:15.103018Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.088784Z"
    }
   },
   "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": 116,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941801Z",
     "start_time": "2024-05-16T13:59:41.941793Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.104057Z",
     "iopub.status.busy": "2025-11-14T04:36:15.103921Z",
     "iopub.status.idle": "2025-11-14T04:36:15.121044Z",
     "shell.execute_reply": "2025-11-14T04:36:15.120572Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.104044Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "159 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(\n",
    "                    input_metas[idx].get(\"id\"),\n",
    "                    input_metas[idx].get(\"tags\"),\n",
    "                    input_metas[pos_idx].get(\"id\"),\n",
    "                    input_metas[pos_idx].get(\"tags\"),\n",
    "                )\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",
    "            elif input_metas[idx].get(\"text\") != input_metas[pos_idx].get(\"text\"):\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": 117,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.942520Z",
     "start_time": "2024-05-16T13:59:41.942511Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.121684Z",
     "iopub.status.busy": "2025-11-14T04:36:15.121546Z",
     "iopub.status.idle": "2025-11-14T04:36:15.139627Z",
     "shell.execute_reply": "2025-11-14T04:36:15.139188Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.121670Z"
    }
   },
   "outputs": [],
   "source": [
    "train_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_tr.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.943072Z",
     "start_time": "2024-05-16T13:59:41.943065Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.140212Z",
     "iopub.status.busy": "2025-11-14T04:36:15.140078Z",
     "iopub.status.idle": "2025-11-14T04:36:15.156580Z",
     "shell.execute_reply": "2025-11-14T04:36:15.156137Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.140199Z"
    }
   },
   "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": 119,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.944246Z",
     "start_time": "2024-05-16T13:59:41.944237Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.157151Z",
     "iopub.status.busy": "2025-11-14T04:36:15.157015Z",
     "iopub.status.idle": "2025-11-14T04:36:15.172545Z",
     "shell.execute_reply": "2025-11-14T04:36:15.172090Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.157138Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total samples 31382 (31382, 203)\n"
     ]
    }
   ],
   "source": [
    "total_iters = len(n_neg_tr) + len(n_pos_tr)\n",
    "print(\"total samples\", total_iters, train_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.173369Z",
     "iopub.status.busy": "2025-11-14T04:36:15.173223Z",
     "iopub.status.idle": "2025-11-14T04:36:15.748018Z",
     "shell.execute_reply": "2025-11-14T04:36:15.747475Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.173354Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15691 0\n"
     ]
    }
   ],
   "source": [
    "metas_tr = read_jsonl(os.path.join(OUT_DATA_DIR, \"meta_tr.jsonl\"))\n",
    "validation_on_metas(metas_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.748714Z",
     "iopub.status.busy": "2025-11-14T04:36:15.748566Z",
     "iopub.status.idle": "2025-11-14T04:36:15.767436Z",
     "shell.execute_reply": "2025-11-14T04:36:15.766994Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.748700Z"
    }
   },
   "outputs": [],
   "source": [
    "# new_metas_tr = []\n",
    "# for index, l in enumerate(metas_tr):\n",
    "#     if index % 2 == 1:\n",
    "#         last_l = new_metas_tr[-1]\n",
    "#         if l[\"tags\"] != last_l[\"tags\"]:\n",
    "#             print(l[\"tags\"], last_l[\"tags\"])\n",
    "#             l[\"tags\"] = last_l[\"tags\"]\n",
    "#     new_metas_tr.append(l)\n",
    "# validation_on_metas(new_metas_tr)\n",
    "# write_jsonl(new_metas_tr, os.path.join(OUT_DATA_DIR, \"meta_tr.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945249Z",
     "start_time": "2024-05-16T13:59:41.945241Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:46.745657Z",
     "iopub.status.busy": "2025-11-14T04:36:46.745230Z",
     "iopub.status.idle": "2025-11-14T04:36:46.766505Z",
     "shell.execute_reply": "2025-11-14T04:36:46.765989Z",
     "shell.execute_reply.started": "2025-11-14T04:36:46.745637Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 epoch per batch 4, total 122.5859375\n",
      "1 epoch per batch 6, total 40.861979166666664\n",
      "1 epoch per batch 8, total 30.646484375\n"
     ]
    }
   ],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 8 / 8 / 4)\n",
    "print(\"1 epoch per batch 6, total\", total_iters / 16 / 8 / 6)\n",
    "print(\"1 epoch per batch 8, total\", total_iters / 16 / 8 / 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.785715Z",
     "iopub.status.busy": "2025-11-14T04:36:15.785576Z",
     "iopub.status.idle": "2025-11-14T04:36:15.800096Z",
     "shell.execute_reply": "2025-11-14T04:36:15.799662Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.785701Z"
    }
   },
   "outputs": [],
   "source": [
    "# import time\n",
    "# time.sleep(60 * 60 * 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.800715Z",
     "iopub.status.busy": "2025-11-14T04:36:15.800580Z",
     "iopub.status.idle": "2025-11-14T04:36:15.815200Z",
     "shell.execute_reply": "2025-11-14T04:36:15.814778Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.800702Z"
    }
   },
   "outputs": [],
   "source": [
    "# df[\"control_tags\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945972Z",
     "start_time": "2024-05-16T13:59:41.945964Z"
    },
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.815811Z",
     "iopub.status.busy": "2025-11-14T04:36:15.815674Z",
     "iopub.status.idle": "2025-11-14T04:36:15.830184Z",
     "shell.execute_reply": "2025-11-14T04:36:15.829754Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.815797Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm/crow && sbatch sbatch_ipo_crow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:15.830804Z",
     "iopub.status.busy": "2025-11-14T04:36:15.830667Z",
     "iopub.status.idle": "2025-11-14T04:36:16.115073Z",
     "shell.execute_reply": "2025-11-14T04:36:16.114614Z",
     "shell.execute_reply.started": "2025-11-14T04:36:15.830790Z"
    }
   },
   "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_crow_t1_c1c2.ipynb\",\n",
    "    os.path.join(OUT_DATA_DIR, \"make_dataset.ipynb\"),\n",
    ")\n",
    "print(\"Cache kept!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-14T04:36:16.115717Z",
     "iopub.status.busy": "2025-11-14T04:36:16.115572Z",
     "iopub.status.idle": "2025-11-14T04:36:16.140943Z",
     "shell.execute_reply": "2025-11-14T04:36:16.139891Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.115703Z"
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'BREAK' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[127], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mBREAK\u001b[49m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'BREAK' is not defined"
     ]
    }
   ],
   "source": [
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# some gymathtics loading prev data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.946562Z",
     "start_time": "2024-05-16T13:59:41.946555Z"
    },
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.141300Z",
     "iopub.status.idle": "2025-11-14T04:36:16.141477Z",
     "shell.execute_reply": "2025-11-14T04:36:16.141394Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.141386Z"
    }
   },
   "outputs": [],
   "source": [
    "# prev_v3_data = \"/app/suno/data/dpo/7v_v20_full/\"\n",
    "\n",
    "# test_val_metas = read_jsonl(os.path.join(prev_v3_data, f\"meta_val.jsonl\"))\n",
    "# test_tr_metas = read_jsonl(os.path.join(prev_v3_data, f\"meta_tr.jsonl\"))\n",
    "\n",
    "# all_ids = set()\n",
    "# for meta in test_val_metas:\n",
    "#     all_ids.add(meta[\"id\"])\n",
    "# for meta in test_tr_metas:\n",
    "#     all_ids.add(meta[\"id\"])\n",
    "# print(len(all_ids), len(test_val_metas) + len(test_tr_metas))\n",
    "\n",
    "# all_ids = list(all_ids)\n",
    "# with open(\"/home/tony/Data/Preference/7b_v2/7v_v20_full_recut_id.json\", \"w\") as fp:\n",
    "#     json.dump(all_ids, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.142275Z",
     "iopub.status.idle": "2025-11-14T04:36:16.142449Z",
     "shell.execute_reply": "2025-11-14T04:36:16.142366Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.142358Z"
    }
   },
   "outputs": [],
   "source": [
    "# x_data = train_df[train_df[\"preference\"]][\"similarity\"]\n",
    "# y_data = train_df[~train_df[\"preference\"]][\"similarity\"]\n",
    "# from matplotlib.colors import LogNorm\n",
    "\n",
    "# fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 10))\n",
    "\n",
    "# # 2D Histogram\n",
    "# h = ax1.hist2d(\n",
    "#     x_data,\n",
    "#     y_data,\n",
    "#     bins=(50, 50),\n",
    "#     cmap=\"coolwarm\",\n",
    "#     range=[[0, 1], [0, 1]],\n",
    "#     norm=LogNorm(),\n",
    "# )\n",
    "\n",
    "# ax1.set_xlabel(\"Semantic Distance (Preferred)\")\n",
    "# ax1.set_ylabel(\"Semantic Distance (Non-Preferred)\")\n",
    "# ax1.set_title(\n",
    "#     \"2D Histogram of Semantic Distances: Preferred vs Non-Preferred (Log Scale)\"\n",
    "# )\n",
    "\n",
    "# cbar1 = plt.colorbar(h[3], ax=ax1)\n",
    "# cbar1.set_label(\"Number of Request IDs (Log Scale)\")\n",
    "\n",
    "# # Scatter plot\n",
    "# ax2.scatter(x_data, y_data, alpha=0.1, s=1)\n",
    "# ax2.set_xlabel(\"Semantic Distance (Preferred)\")\n",
    "# ax2.set_ylabel(\"Semantic Distance (Non-Preferred)\")\n",
    "# ax2.set_title(\"Scatter Plot of Semantic Distances: Preferred vs Non-Preferred\")\n",
    "# ax2.set_xlim(0, 1)\n",
    "# ax2.set_ylim(0, 1)\n",
    "\n",
    "# plt.tight_layout()\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.142859Z",
     "iopub.status.idle": "2025-11-14T04:36:16.143024Z",
     "shell.execute_reply": "2025-11-14T04:36:16.142950Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.142942Z"
    }
   },
   "outputs": [],
   "source": [
    "# train_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"meta_tr.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.143748Z",
     "iopub.status.idle": "2025-11-14T04:36:16.143908Z",
     "shell.execute_reply": "2025-11-14T04:36:16.143833Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.143825Z"
    }
   },
   "outputs": [],
   "source": [
    "# train_info.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.144304Z",
     "iopub.status.idle": "2025-11-14T04:36:16.144459Z",
     "shell.execute_reply": "2025-11-14T04:36:16.144391Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.144383Z"
    }
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "\n",
    "# a = torch.tensor([6.2500e-04, 3.9062e-05, 2.3462e-03])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.145202Z",
     "iopub.status.idle": "2025-11-14T04:36:16.145374Z",
     "shell.execute_reply": "2025-11-14T04:36:16.145284Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.145276Z"
    }
   },
   "outputs": [],
   "source": [
    "# a.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.145799Z",
     "iopub.status.idle": "2025-11-14T04:36:16.145946Z",
     "shell.execute_reply": "2025-11-14T04:36:16.145877Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.145870Z"
    }
   },
   "outputs": [],
   "source": [
    "# import time\n",
    "# time.sleep(3600 * 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.146437Z",
     "iopub.status.idle": "2025-11-14T04:36:16.146593Z",
     "shell.execute_reply": "2025-11-14T04:36:16.146516Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.146508Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm/diffusion && sbatch run_diffusion_infill.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.147098Z",
     "iopub.status.idle": "2025-11-14T04:36:16.147265Z",
     "shell.execute_reply": "2025-11-14T04:36:16.147186Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.147178Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_control_slider(metadata):\n",
    "    if \"param_experiment\" in metadata:\n",
    "        exp = metadata.get(\"param_experiment\", \"\")\n",
    "        if exp:\n",
    "            if exp == \"mask_control_slider\":\n",
    "                if not metadata.get(\"control_sliders\", None):\n",
    "                    return False\n",
    "                else:\n",
    "                    return True\n",
    "            return None\n",
    "\n",
    "df[\"mask_control\"] = df.apply(\n",
    "    lambda row: get_control_slider(row[\"metadata\"]), axis=1\n",
    ")\n",
    "masked_request_id = df[~df[\"mask_control\"].isna()][\"request_id\"].unique()\n",
    "subset_df = df[df[\"request_id\"].isin(masked_request_id)].copy()\n",
    "print(\"check subset shape\", df.shape, subset_df.shape)\n",
    "# control masked\n",
    "user_pref_pct = (\n",
    "    subset_df[subset_df[\"mask_control\"] == True].groupby(\"user_id\")[\"preference\"]\n",
    "    .apply(lambda x: x.mean())\n",
    ")\n",
    "print(user_pref_pct.head())\n",
    "# Plot a histogram of the user preference percentages\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.hist(user_pref_pct, bins=np.linspace(0, 1, 100), edgecolor=\"black\")\n",
    "plt.xlabel(\"Preference % for mask_control\")\n",
    "plt.ylabel(\"Number of Users\")\n",
    "plt.title(\"Histogram of User Preference for 'mask_control'\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.147862Z",
     "iopub.status.idle": "2025-11-14T04:36:16.148020Z",
     "shell.execute_reply": "2025-11-14T04:36:16.147946Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.147938Z"
    }
   },
   "outputs": [],
   "source": [
    "user_counts = subset_df.groupby(\"user_id\")[\"preference\"].count()\n",
    "eligible_users = user_counts[user_counts >= 4].index\n",
    "len(eligible_users)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.148551Z",
     "iopub.status.idle": "2025-11-14T04:36:16.148703Z",
     "shell.execute_reply": "2025-11-14T04:36:16.148631Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.148623Z"
    }
   },
   "outputs": [],
   "source": [
    "# control masked\n",
    "eligible_user_pref_pct = (\n",
    "    subset_df[subset_df[\"user_id\"].isin(eligible_users) & (subset_df[\"mask_control\"] == True)].groupby(\"user_id\")[\"preference\"]\n",
    "    .apply(lambda x: x.mean())\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.149187Z",
     "iopub.status.idle": "2025-11-14T04:36:16.149335Z",
     "shell.execute_reply": "2025-11-14T04:36:16.149266Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.149259Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot a histogram of the user preference percentages\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.hist(eligible_user_pref_pct, bins=np.linspace(0, 1, 20), edgecolor=\"black\")\n",
    "plt.xlabel(\"Preference % for mask_control\")\n",
    "plt.ylabel(\"Number of Users\")\n",
    "plt.title(\"Histogram of User Preference for 'mask_control'\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.150168Z",
     "iopub.status.idle": "2025-11-14T04:36:16.150356Z",
     "shell.execute_reply": "2025-11-14T04:36:16.150262Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.150253Z"
    }
   },
   "outputs": [],
   "source": [
    "eligible_user_pref_pct[172918] # kakermix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## sliders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.150659Z",
     "iopub.status.idle": "2025-11-14T04:36:16.150809Z",
     "shell.execute_reply": "2025-11-14T04:36:16.150739Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.150732Z"
    }
   },
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "import pandas as pd\n",
    "\n",
    "def map_control_sliders_to_df(df: pd.DataFrame) -> pd.DataFrame:\n",
    "    \"\"\"Extracts 'style_weight' and 'weirdness_constraint' from the 'control_sliders' dict in the 'metadata' column\n",
    "    and adds them as new columns to the DataFrame.\n",
    "\n",
    "    Args:\n",
    "        df (pd.DataFrame): DataFrame with a 'metadata' column containing a 'control_sliders' dict.\n",
    "\n",
    "    Returns:\n",
    "        pd.DataFrame: DataFrame with added 'style_weight' and 'weirdness_constraint' columns.\n",
    "\n",
    "    Raises:\n",
    "        KeyError: If 'control_sliders', 'style_weight', or 'weirdness_constraint' are missing in any row.\n",
    "        TypeError: If the extracted values are not floats.\n",
    "\n",
    "    Example:\n",
    "        >>> import pandas as pd\n",
    "        >>> data = [{'metadata': {'control_sliders': {'style_weight': 0.89, 'weirdness_constraint': 0.8}}}]\n",
    "        >>> df = pd.DataFrame(data)\n",
    "        >>> df = map_control_sliders_to_df(df)\n",
    "        >>> df[['style_weight', 'weirdness_constraint']].iloc[0].tolist()\n",
    "        [0.89, 0.8]\n",
    "    \"\"\"\n",
    "    # Vectorized extraction for performance\n",
    "    sliders = df[\"metadata\"].map(lambda m: m.get(\"control_sliders\", {}))\n",
    "    style_weight = sliders.map(lambda s: s.get(\"style_weight\", None))\n",
    "    weirdness_constraint = sliders.map(lambda s: s.get(\"weirdness_constraint\", None))\n",
    "    audio_weight = sliders.map(lambda s: s.get(\"audio_weight\", None))\n",
    "\n",
    "    df[\"style_weight\"] = style_weight\n",
    "    df[\"weirdness_constraint\"] = weirdness_constraint\n",
    "    df[\"audio_weight\"] = audio_weight\n",
    "    return df\n",
    "\n",
    "def print_percentiles(\n",
    "    data: np.ndarray,\n",
    "    percentiles: list[float] = [5, 10, 20, 25, 50, 75, 80, 90, 95]\n",
    ") -> None:\n",
    "    \"\"\"Prints specified percentiles of the data.\n",
    "\n",
    "    Args:\n",
    "        data (np.ndarray): Array of values to compute percentiles for.\n",
    "        percentiles (list[float], optional): List of percentiles to print. Defaults to [5, 10, 20, 25, 50, 75, 80, 90, 95].\n",
    "\n",
    "    Example:\n",
    "        >>> print_percentiles(np.array([1, 2, 3, 4, 5]))\n",
    "    \"\"\"\n",
    "    data = data[~data.isna()]\n",
    "    results = np.percentile(data, percentiles)\n",
    "    print(\"Percentiles:\")\n",
    "    for p, v in zip(percentiles, results):\n",
    "        print(f\"  {p:>3}%: {v:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.151506Z",
     "iopub.status.idle": "2025-11-14T04:36:16.151679Z",
     "shell.execute_reply": "2025-11-14T04:36:16.151590Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.151583Z"
    }
   },
   "outputs": [],
   "source": [
    "df_total = map_control_sliders_to_df(df_total)\n",
    "df_total[[\"style_weight\",\"weirdness_constraint\",\"audio_weight\"]].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.152046Z",
     "iopub.status.idle": "2025-11-14T04:36:16.152192Z",
     "shell.execute_reply": "2025-11-14T04:36:16.152124Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.152117Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 5))\n",
    "df_total[\"style_weight\"].hist(bins=np.linspace(0, 1, 100), color=\"skyblue\", edgecolor=\"black\")\n",
    "print_percentiles(df_total[\"style_weight\"])\n",
    "plt.title(\"Distribution of Style Weight\")\n",
    "plt.xlabel(\"Style Weight\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.grid(True, linestyle=\"--\", alpha=0.6)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.152900Z",
     "iopub.status.idle": "2025-11-14T04:36:16.153085Z",
     "shell.execute_reply": "2025-11-14T04:36:16.152992Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.152984Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 5))\n",
    "df_total[\"weirdness_constraint\"].hist(bins=np.linspace(0, 1, 100), color=\"skyblue\", edgecolor=\"black\")\n",
    "print_percentiles(df_total[\"weirdness_constraint\"])\n",
    "plt.title(\"Distribution of Weirdness\")\n",
    "plt.xlabel(\"weirdness_constraint\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.grid(True, linestyle=\"--\", alpha=0.6)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.153409Z",
     "iopub.status.idle": "2025-11-14T04:36:16.153566Z",
     "shell.execute_reply": "2025-11-14T04:36:16.153495Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.153487Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 5))\n",
    "df_total[\"audio_weight\"].hist(bins=np.linspace(0, 1, 100), color=\"skyblue\", edgecolor=\"black\")\n",
    "print_percentiles(df_total[\"audio_weight\"])\n",
    "plt.title(\"Distribution of audio_weight\")\n",
    "plt.xlabel(\"audio_weight\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.grid(True, linestyle=\"--\", alpha=0.6)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fetch other parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.154157Z",
     "iopub.status.idle": "2025-11-14T04:36:16.154309Z",
     "shell.execute_reply": "2025-11-14T04:36:16.154236Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.154229Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_total[df_total[\"task\"] == \"stem_condition\"][[\"s3_id\", \"preference\", \"request_id\", \"control_tags\"]].head(n=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.154969Z",
     "iopub.status.idle": "2025-11-14T04:36:16.155139Z",
     "shell.execute_reply": "2025-11-14T04:36:16.155051Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.155043Z"
    }
   },
   "outputs": [],
   "source": [
    "len(metas_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.155467Z",
     "iopub.status.idle": "2025-11-14T04:36:16.155620Z",
     "shell.execute_reply": "2025-11-14T04:36:16.155549Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.155542Z"
    }
   },
   "outputs": [],
   "source": [
    "c = Counter()\n",
    "for data_meta in metas_tr:\n",
    "    c[data_meta.get(\"control_tags\", \"\") if isinstance(data_meta.get(\"control_tags\", \"\"), str) else \"\"] += 1\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.156168Z",
     "iopub.status.idle": "2025-11-14T04:36:16.156326Z",
     "shell.execute_reply": "2025-11-14T04:36:16.156251Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.156244Z"
    }
   },
   "outputs": [],
   "source": [
    "# with open(\"/app2/suno/data/dpo/crow_t1_v3/dodo_t23_2025-08-28_01-30-09_cached_loss.json\", \"r\") as fp:\n",
    "#     old_loss = json.load(fp)\n",
    "# with open(\"/app2/suno/data/dpo/crow_t1_v3/dodo_t23_2025-08-28_01-30-09_redo3_cached_loss.json\", \"r\") as fp:\n",
    "#     new_loss = json.load(fp)\n",
    "\n",
    "# sem_losses = []\n",
    "# for curr_id, o_v in old_loss[\"train\"].items():\n",
    "#     n_v = new_loss[\"train\"][curr_id]\n",
    "#     for k in o_v:\n",
    "#         assert o_v[k] == n_v[k]\n",
    "#         if k ==  \"semantic_0\":\n",
    "#             if o_v[\"orig_idx\"] % 2 == 1:\n",
    "#                 sem_losses.append(o_v[k])\n",
    "#             # if o_v[k] > 4 and o_v[\"orig_idx\"] % 2 == 1:\n",
    "#             #     print(o_v, o_v[\"orig_idx\"] % 3360, o_v[\"orig_idx\"] % 3360 // 4)\n",
    "# sem_losses.sort()\n",
    "# plt.hist(sem_losses, bins=np.linspace(0, 10, 100))\n",
    "# plt.yscale(\"log\")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.157199Z",
     "iopub.status.idle": "2025-11-14T04:36:16.157370Z",
     "shell.execute_reply": "2025-11-14T04:36:16.157280Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.157273Z"
    }
   },
   "outputs": [],
   "source": [
    "# df[df[\"is_public\"]].to_pickle(\"/home/tony/Data/Preference/crow_t1/interesting_clips_crow_t1_20251020_public_only.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.157698Z",
     "iopub.status.idle": "2025-11-14T04:36:16.157849Z",
     "shell.execute_reply": "2025-11-14T04:36:16.157775Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.157768Z"
    }
   },
   "outputs": [],
   "source": [
    "df[df[\"is_public\"] & (df[\"tags\"].str.len() < 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.158248Z",
     "iopub.status.idle": "2025-11-14T04:36:16.158397Z",
     "shell.execute_reply": "2025-11-14T04:36:16.158328Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.158321Z"
    }
   },
   "outputs": [],
   "source": [
    "from suno_analytics.preference_data_selection import (\n",
    "    plot_clip_distribution\n",
    ")\n",
    "plot_clip_distribution(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.158699Z",
     "iopub.status.idle": "2025-11-14T04:36:16.158840Z",
     "shell.execute_reply": "2025-11-14T04:36:16.158774Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.158767Z"
    }
   },
   "outputs": [],
   "source": [
    "plot_clip_distribution(df_total)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.159419Z",
     "iopub.status.idle": "2025-11-14T04:36:16.159573Z",
     "shell.execute_reply": "2025-11-14T04:36:16.159501Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.159493Z"
    }
   },
   "outputs": [],
   "source": [
    "plot_clip_distribution(df_total[df_total[\"source\"] == \"web\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-11-14T04:36:16.160309Z",
     "iopub.status.idle": "2025-11-14T04:36:16.160477Z",
     "shell.execute_reply": "2025-11-14T04:36:16.160389Z",
     "shell.execute_reply.started": "2025-11-14T04:36:16.160382Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Assuming df is your dataframe with 'created_at' column\n",
    "# df = pd.read_csv('your_data.csv')\n",
    "\n",
    "# Convert created_at to datetime if not already\n",
    "df['created_at'] = pd.to_datetime(df['created_at'])\n",
    "\n",
    "# Extract day and hour\n",
    "df['date'] = df['created_at'].dt.date\n",
    "df['hour'] = df['created_at'].dt.hour\n",
    "\n",
    "# Define your cut_frac function here\n",
    "def compute_cut_frac(group_df):\n",
    "    \"\"\"\n",
    "    Define your cut fraction logic here.\n",
    "    \n",
    "    Examples:\n",
    "    - Fraction filtered out: (group_df['filtered'] == True).mean()\n",
    "    - Fraction above threshold: (group_df['margin'] > 0.5).mean()\n",
    "    - Fraction in bottom quantile: (group_df['score'] < group_df['score'].quantile(0.2)).mean()\n",
    "    \"\"\"\n",
    "    # CUSTOMIZE THIS:\n",
    "    # Example: fraction of rows where margin > threshold\n",
    "    threshold = 0.5\n",
    "    cut_frac = (group_df[\"preference\"] & (group_df[\"norm_play_frac\"] > 5.1) \n",
    "                & (group_df[\"norm_play_frac\"] >= group_df[\"\"] / 3)\n",
    "               ).mean()\n",
    "    \n",
    "    return cut_frac\n",
    "\n",
    "# Group by date and hour\n",
    "daily_hourly_stats = df.groupby(['date', 'hour']).apply(compute_cut_frac).reset_index()\n",
    "daily_hourly_stats.columns = ['date', 'hour', 'cut_frac']\n",
    "\n",
    "# Also compute count per day-hour\n",
    "daily_hourly_counts = df.groupby(['date', 'hour']).size().reset_index(name='count')\n",
    "daily_hourly_stats = daily_hourly_stats.merge(daily_hourly_counts, on=['date', 'hour'])\n",
    "\n",
    "# Create plots\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))\n",
    "\n",
    "# Plot 1: Line plot with one line per day\n",
    "for date in sorted(df['date'].unique()):\n",
    "    date_data = daily_hourly_stats[daily_hourly_stats['date'] == date]\n",
    "    ax1.plot(date_data['hour'], date_data['cut_frac'], \n",
    "             marker='o', linewidth=2, markersize=4, label=str(date), alpha=0.7)\n",
    "\n",
    "ax1.set_xlabel('Hour of Day', fontsize=12)\n",
    "ax1.set_ylabel('Cut Fraction', fontsize=12)\n",
    "ax1.set_title('Cut Fraction by Day and Hour', fontsize=14, fontweight='bold')\n",
    "ax1.grid(True, alpha=0.3)\n",
    "ax1.set_xlim(-0.5, 23.5)\n",
    "ax1.set_xticks(range(0, 24, 2))\n",
    "ax1.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Date')\n",
    "\n",
    "# Plot 2: Heatmap\n",
    "pivot_table = daily_hourly_stats.pivot(index='date', columns='hour', values='cut_frac')\n",
    "im = ax2.imshow(pivot_table, aspect='auto', cmap='viridis', interpolation='nearest')\n",
    "\n",
    "ax2.set_xlabel('Hour of Day', fontsize=12)\n",
    "ax2.set_ylabel('Date', fontsize=12)\n",
    "ax2.set_title('Cut Fraction Heatmap', fontsize=14, fontweight='bold')\n",
    "ax2.set_xticks(range(24))\n",
    "ax2.set_xticklabels(range(24))\n",
    "ax2.set_yticks(range(len(pivot_table)))\n",
    "ax2.set_yticklabels([str(d) for d in pivot_table.index])\n",
    "\n",
    "# Add colorbar\n",
    "cbar = plt.colorbar(im, ax=ax2)\n",
    "cbar.set_label('Cut Fraction', fontsize=12)\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Print summary stats\n",
    "print(f\"\\nSummary Statistics:\")\n",
    "print(f\"Overall cut_frac: {daily_hourly_stats['cut_frac'].mean():.3f}\")\n",
    "print(f\"Min cut_frac: {daily_hourly_stats['cut_frac'].min():.3f}\")\n",
    "print(f\"Max cut_frac: {daily_hourly_stats['cut_frac'].max():.3f}\")\n",
    "print(f\"Total samples: {daily_hourly_stats['count'].sum()}\")\n",
    "print(f\"\\nPer-day averages:\")\n",
    "day_avg = daily_hourly_stats.groupby('date')['cut_frac'].mean()\n",
    "for date, avg in day_avg.items():\n",
    "    print(f\"  {date}: {avg:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.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
}
