{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:21.040680Z",
     "start_time": "2024-05-16T13:58:19.777010Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:27:53.672777Z",
     "iopub.status.busy": "2025-03-31T02:27:53.672635Z",
     "iopub.status.idle": "2025-03-31T02:27:55.968687Z",
     "shell.execute_reply": "2025-03-31T02:27:55.968006Z",
     "shell.execute_reply.started": "2025-03-31T02:27:53.672760Z"
    }
   },
   "outputs": [],
   "source": [
    "import ast\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "from collections import defaultdict\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from preference_data_preparation_4min_30b_task 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",
    "import matplotlib.pyplot as plt\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:21.082172Z",
     "start_time": "2024-05-16T13:58:21.041926Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:27:55.969664Z",
     "iopub.status.busy": "2025-03-31T02:27:55.969421Z",
     "iopub.status.idle": "2025-03-31T02:27:56.023788Z",
     "shell.execute_reply": "2025-03-31T02:27:56.023179Z",
     "shell.execute_reply.started": "2025-03-31T02:27:55.969647Z"
    }
   },
   "outputs": [],
   "source": [
    "OUT_DATA_DIR = \"/app/suno/data/dpo/13b_s32_v34/\"\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 = \"/app/suno/data/dpo/13b_s32_npz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:53.962528Z",
     "start_time": "2024-05-16T13:58:21.105919Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:27:56.024551Z",
     "iopub.status.busy": "2025-03-31T02:27:56.024400Z",
     "iopub.status.idle": "2025-03-31T02:29:13.826417Z",
     "shell.execute_reply": "2025-03-31T02:29:13.825643Z",
     "shell.execute_reply.started": "2025-03-31T02:27:56.024535Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preference data shape (2303948, 90)\n"
     ]
    }
   ],
   "source": [
    "# df = pd.read_csv(\n",
    "#     \"/home/tony/Data/Preference/13b_v0/interesting_clips_v3p5_s_8_20240813.csv\"\n",
    "# )  # , engine='python')\n",
    "df = pd.read_pickle(\n",
    "    \"/home/tony/Data/Preference/13b_v32/interesting_clips_v4_h_s_32_20250330_full_long.pkl\"\n",
    ")  # , engine='python')\n",
    "print(\"Preference data shape\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.199480Z",
     "start_time": "2024-05-16T13:58:53.963687Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:29:13.827338Z",
     "iopub.status.busy": "2025-03-31T02:29:13.827168Z",
     "iopub.status.idle": "2025-03-31T02:32:37.465816Z",
     "shell.execute_reply": "2025-03-31T02:32:37.465066Z",
     "shell.execute_reply.started": "2025-03-31T02:29:13.827320Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10966595\n",
      "10966595\n",
      "pre-downloaded df (2303948, 90)\n",
      "downloaded df (2303948, 90)\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",
    "if \"cycle\" in NPZ_DIR:\n",
    "    # hack in the cycle label\n",
    "    df[\"s3_id\"] += \"_gen_cycle\"\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": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.467253Z",
     "start_time": "2024-05-16T13:58:56.207647Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:37.467854Z",
     "iopub.status.busy": "2025-03-31T02:32:37.467555Z",
     "iopub.status.idle": "2025-03-31T02:32:38.655729Z",
     "shell.execute_reply": "2025-03-31T02:32:38.655112Z",
     "shell.execute_reply.started": "2025-03-31T02:32:37.467833Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "is_13b\n",
       "True    2303948\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"is_13b\"] = df[\"model_name\"].str.contains(\"-s-\")\n",
    "df[\"is_13b\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:38.656644Z",
     "iopub.status.busy": "2025-03-31T02:32:38.656450Z",
     "iopub.status.idle": "2025-03-31T02:32:38.680294Z",
     "shell.execute_reply": "2025-03-31T02:32:38.679730Z",
     "shell.execute_reply.started": "2025-03-31T02:32:38.656627Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "is_public\n",
      "False    2171763\n",
      "True      132185\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(df[\"is_public\"].value_counts())\n",
    "# remove public for now cause fucking users\n",
    "# df = df[~df[\"is_public\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LET's do the data prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.592883Z",
     "start_time": "2024-05-16T13:58:56.470781Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:38.681235Z",
     "iopub.status.busy": "2025-03-31T02:32:38.680920Z",
     "iopub.status.idle": "2025-03-31T02:32:42.257664Z",
     "shell.execute_reply": "2025-03-31T02:32:42.256922Z",
     "shell.execute_reply.started": "2025-03-31T02:32:38.681218Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preference  model_name     \n",
      "False       chirp-v4-h-s-32    1151974\n",
      "True        chirp-v4-h-s-32    1151974\n",
      "Name: count, dtype: int64\n",
      "(2303948, 91)\n",
      "(2303948, 91)\n"
     ]
    }
   ],
   "source": [
    "## for 13b this is easy for now\n",
    "print(df.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "print(df.shape)\n",
    "df = df[df[\"model_name\"].isin([\"chirp-v4-h-s-32\"])]\n",
    "print(df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.909539Z",
     "start_time": "2024-05-16T13:58:56.595736Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:42.258539Z",
     "iopub.status.busy": "2025-03-31T02:32:42.258375Z",
     "iopub.status.idle": "2025-03-31T02:32:45.767471Z",
     "shell.execute_reply": "2025-03-31T02:32:45.766866Z",
     "shell.execute_reply.started": "2025-03-31T02:32:42.258522Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2303948, 91)\n",
      "(2303948, 91)\n",
      "preference  model_name     \n",
      "False       chirp-v4-h-s-32    1151974\n",
      "True        chirp-v4-h-s-32    1151974\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(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(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": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:45.768220Z",
     "iopub.status.busy": "2025-03-31T02:32:45.768070Z",
     "iopub.status.idle": "2025-03-31T02:32:45.783455Z",
     "shell.execute_reply": "2025-03-31T02:32:45.783015Z",
     "shell.execute_reply.started": "2025-03-31T02:32:45.768205Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def custom_parse(x):\n",
    "    try:\n",
    "        return json.loads(x)\n",
    "    except:\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:36.043975Z",
     "start_time": "2024-05-16T13:58:56.910958Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:32:45.784112Z",
     "iopub.status.busy": "2025-03-31T02:32:45.783966Z",
     "iopub.status.idle": "2025-03-31T02:40:05.948549Z",
     "shell.execute_reply": "2025-03-31T02:40:05.947769Z",
     "shell.execute_reply.started": "2025-03-31T02:32:45.784098Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unique_requests 1151974\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_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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:36.393047Z",
     "start_time": "2024-05-16T13:59:36.048831Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:06.923799Z",
     "iopub.status.busy": "2025-03-31T02:40:06.923378Z",
     "iopub.status.idle": "2025-03-31T02:40:07.313676Z",
     "shell.execute_reply": "2025-03-31T02:40:07.312945Z",
     "shell.execute_reply.started": "2025-03-31T02:40:06.923780Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unique_requests 1151974\n"
     ]
    }
   ],
   "source": [
    "# GPT requests are also fine for now\n",
    "print(\"unique_requests\", df[\"request_id\"].nunique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:40.799375Z",
     "start_time": "2024-05-16T13:59:36.394236Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:07.314592Z",
     "iopub.status.busy": "2025-03-31T02:40:07.314413Z",
     "iopub.status.idle": "2025-03-31T02:40:34.232465Z",
     "shell.execute_reply": "2025-03-31T02:40:34.231895Z",
     "shell.execute_reply.started": "2025-03-31T02:40:07.314574Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "102393\n",
      "good_continue_at\n",
      "True     2296413\n",
      "False       7535\n",
      "Name: count, dtype: int64\n",
      "\n",
      " Check some basics... \n",
      " preference\n",
      "False    1151974\n",
      "True     1151974\n",
      "Name: count, dtype: int64 is_13b\n",
      "True    2303948\n",
      "Name: count, dtype: int64 model_name\n",
      "chirp-v4-h-s-32    2303948\n",
      "Name: count, dtype: int64 preference  model_name     \n",
      "False       chirp-v4-h-s-32    1151974\n",
      "True        chirp-v4-h-s-32    1151974\n",
      "Name: count, dtype: int64\n",
      "task\n",
      "                 2036082\n",
      "extend            167108\n",
      "upload_extend     100294\n",
      "cover                464\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "df = df.loc[:, ~df.columns.duplicated()].copy()\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",
    "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[\"s3_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[\"is_13b\"].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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:34.233199Z",
     "iopub.status.busy": "2025-03-31T02:40:34.233047Z",
     "iopub.status.idle": "2025-03-31T02:40:40.563818Z",
     "shell.execute_reply": "2025-03-31T02:40:40.563317Z",
     "shell.execute_reply.started": "2025-03-31T02:40:34.233183Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pos_diff_preference\n",
       "1.0    756448\n",
       "2.0    395526\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = df.sort_values(by=[\"request_id\", \"preference\", \"diff_preference\"])\n",
    "df[\"pos_diff_preference\"] = df[\"diff_preference\"].diff()\n",
    "# df[\"cer_diff_preference\"] = df[\"cer\"].diff()\n",
    "df[df[\"preference\"]][\"pos_diff_preference\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:40.564532Z",
     "iopub.status.busy": "2025-03-31T02:40:40.564385Z",
     "iopub.status.idle": "2025-03-31T02:40:40.579542Z",
     "shell.execute_reply": "2025-03-31T02:40:40.579100Z",
     "shell.execute_reply.started": "2025-03-31T02:40:40.564518Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 174624 duplicated prompts 87314 unique requests\n",
      "Found 44053 request_ids with duplicate prompts but not highest play counts in their group\n",
      "['eda856f3-1c7f-4306-8d52-e572c47533e8', '92a0b525-b8ac-48dd-a005-0e9c7d108af9', '7dd2a262-d173-4124-bd35-36136a495b31', '07d49149-8138-413e-9cc5-074a3720a863', '331cc6af-e788-4280-b04a-3ab80efef438', '79153b01-431a-46a6-ab4d-47f01c27c4d3', 'cb38b71b-89ab-4a88-a26f-fd23dce97811', 'e52d6833-763a-4584-ac74-7105661352ec', 'ee08ace4-86e4-43af-85fd-d1c2701352ac', '8a970e1b-f9b1-4065-a6a8-d967b59155e7']\n",
      "Before dedup user gen requests 2303948\n",
      "After dedup user gen requests 2215842\n"
     ]
    }
   ],
   "source": [
    "# Find duplicated prompts with count > 2\n",
    "duplicate_entries = df.groupby([\"user_id\", \"prompt_text\", \"tags\", \"task\"]).filter(\n",
    "    lambda x: len(x) > 2\n",
    ")\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([\"user_id\", \"prompt_text\", \"tags\", \"task\"])\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": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:40.580170Z",
     "iopub.status.busy": "2025-03-31T02:40:40.580033Z",
     "iopub.status.idle": "2025-03-31T02:40:41.752428Z",
     "shell.execute_reply": "2025-03-31T02:40:41.751844Z",
     "shell.execute_reply.started": "2025-03-31T02:40:40.580156Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive with likes (687306, 157)\n"
     ]
    }
   ],
   "source": [
    "test_mask = (df[\"preference\"] == True) & (df[\"upvote_count\"] >= 1)\n",
    "print(\"positive with likes\", df[test_mask].shape)\n",
    "# positive with likes (285955, 149) -- 0215 data\n",
    "# positive with likes (517621, 151) -- 0217 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.035167Z",
     "start_time": "2024-05-16T13:59:40.801098Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:41.753167Z",
     "iopub.status.busy": "2025-03-31T02:40:41.753021Z",
     "iopub.status.idle": "2025-03-31T02:40:45.447140Z",
     "shell.execute_reply": "2025-03-31T02:40:45.446561Z",
     "shell.execute_reply.started": "2025-03-31T02:40:41.753152Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "negative 1107789 positive 224697\n",
      "total pair requests 1107921 selected pair requests 224676 frac 0.203\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 = 2\n",
    "# this is a filter on the concated clip\n",
    "concat_total_play_count = 3\n",
    "\n",
    "neg_filter_selection_mask = (\n",
    "    (df[\"preference\"] == False)  # get basics aligned\n",
    "    & (df[\"reaction_play_count\"] >= 1)  # has to be played once\n",
    "    # & (df[\"play_count\"] <= 3)  # if it is actually bad, shouldn't be listened often\n",
    "    & (df[\"duration\"] >= 5)  # can't be too short, otherwise it is obvious\n",
    "    # & (df[\"duration\"] <= 240)  # can't be badly long\n",
    "    & (df[\"has_continue_and_start_continue_at\"].isna())  # won't have any continues\n",
    "    & (df[\"norm_play_frac\"] <= 3.1)\n",
    "    # & (df[\"dislike_count\"] >= 1) # this is kinda strict\n",
    "    #     & (\n",
    "    #         (df_slice[\"is_in_playlist\"] == False)\n",
    "    #         & (df_slice[\"concat_in_playlist\"] == False)\n",
    "    #     )  # can't be part of a playlist -- otherwise there are some like signal in it?\n",
    ")\n",
    "pos_filter_selectin_mask = (\n",
    "    (df[\"preference\"] == True)  # get basics aligned\n",
    "    & (\n",
    "        df[\"good_continue_at\"] == True\n",
    "    )  # if continue, needs to continue off a certain percentage\n",
    "    & (df[\"reaction_play_count\"] >= 1)\n",
    "    & (df[\"play_rel_diff\"] >= 0)  # this is more like quality assurance\n",
    "    & (df[\"duration\"] >= 5)  # can't be too short, otherwise it is obvious\n",
    "    # & (df[\"duration\"] <= 240)  # can't be badly long\n",
    "    & (df[\"dislike_count\"] == 0)  # can't have dislikes\n",
    "    & (df[\"flag_count\"] == 0)  # can't have issues\n",
    "    & (\n",
    "        (\n",
    "            (df[\"part_of_concat\"] == True)\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\"] == False)\n",
    "            & (df[\"reaction_play_count\"] >= normal_pos_play_count)\n",
    "        )\n",
    "    )\n",
    "    & (df[\"user_n_clips\"] >= 100)  # user needs to have genereated at least 20\n",
    "    # & (df[\"duration_rel_diff\"] < 10) # positive isn't just longer\n",
    "    # & ((df[\"upvote_count\"] >= 1) )\n",
    "    & ((df[\"norm_play_frac\"] >= 5.1) | (~df[\"continued_parent\"].isna()))\n",
    "    & (\n",
    "        df[\"norm_play_frac\"] >= df[\"reaction_play_count\"] / 2\n",
    "    )  # play duration is not low on average\n",
    "    & (\n",
    "        df[\"task\"].isin([\"\", \"extend\", \"upload_extend\"])\n",
    "    )  # play duration is not low on average\n",
    "    # & (df[\"pos_diff_preference\"] == 2)\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",
    "    \"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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.250737Z",
     "start_time": "2024-05-16T13:59:41.036434Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:45.447897Z",
     "iopub.status.busy": "2025-03-31T02:40:45.447744Z",
     "iopub.status.idle": "2025-03-31T02:40:47.317572Z",
     "shell.execute_reply": "2025-03-31T02:40:47.317009Z",
     "shell.execute_reply.started": "2025-03-31T02:40:45.447881Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " requests 224676 clips 449352 total khrs 22.374; N gpus for 1000 iters 28.084; 4 gpus for x iters 3510.562; n unique users 64856 n pro users 56564\n"
     ]
    }
   ],
   "source": [
    "df_slice = df[df[\"request_id\"].isin(set(unique_requests))].copy()\n",
    "print(\n",
    "    f\"{os.path.basename(OUT_DATA_DIR)} requests\",\n",
    "    df_slice[\"request_id\"].nunique(),\n",
    "    \"clips\",\n",
    "    df_slice.shape[0],\n",
    "    f\"total khrs {sum(df_slice['duration'] / 3600 / 1000):.3f};\",\n",
    "    f\"N gpus for 1000 iters {df_slice.shape[0] / 8 / 2 / 1000:.3f};\",\n",
    "    f\"4 gpus for x iters {df_slice.shape[0] / 8 / 4 / 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",
    "# 76171 152342 total khrs 2.880 n gpus for 1250 iters 3.809\n",
    "# 8 v9 requests 106299 clips 212598 total khrs 9.776; N gpus for 1000 iters 13.287; n unique users 27775\n",
    "# 29 v2 requests 22572 clips 45144 total khrs 2.252; N gpus for 1000 iters 2.821; 4 gpus for x iters 705.375; n unique users 18149 n pro users 10155\n",
    "# 29 v4 requests 32479 clips 64958 total khrs 3.248; N gpus for 1000 iters 4.060; 4 gpus for x iters 1014.969; n unique users 25192 n pro users 13547\n",
    "# samve for v5\n",
    "# 31 v9  requests 63922 clips 127844 total khrs 6.631; N gpus for 1000 iters 7.990; 4 gpus for x iters 1997.562; n unique users 25461 n pro users 24396\n",
    "# 32 v1  requests 42336 clips 84672 total khrs 4.311; N gpus for 1000 iters 5.292; 4 gpus for x iters 1323.000; n unique users 20440 n pro users 14949\n",
    "# 32 v2  requests 83218 clips 166436 total khrs 8.672; N gpus for 1000 iters 10.402; 4 gpus for x iters 2600.562; n unique users 36912 n pro users 24545\n",
    "# 32 v3  requests 71020 clips 142040 total khrs 7.562; N gpus for 1000 iters 8.877; 4 gpus for x iters 2219.375; n unique users 25143 n pro users 22341\n",
    "# 32 v5  requests 83743 clips 167486 total khrs 8.923; N gpus for 1000 iters 10.468; 4 gpus for x iters 2616.969; n unique users 28348 n pro users 24894\n",
    "# 32 v6  requests 63924 clips 127848 total khrs 6.725; N gpus for 1000 iters 7.990; 4 gpus for x iters 1997.625; n unique users 24675 n pro users 22071\n",
    "# 32 v7  requests 106565 clips 213130 total khrs 11.335; N gpus for 1000 iters 13.321; 4 gpus for x iters 3330.156; n unique users 34779 n pro users 29833\n",
    "# 32 v8  requests 135955 clips 271910 total khrs 14.459; N gpus for 1000 iters 16.994; 4 gpus for x iters 4248.594; n unique users 42469 n pro users 35114\n",
    "# 32 v9  requests 60046 clips 120092 total khrs 6.444; N gpus for 1000 iters 7.506; 4 gpus for x iters 1876.438; n unique users 15885 n pro users 14566\n",
    "# 32 v10  requests 79415 clips 158830 total khrs 8.497; N gpus for 1000 iters 9.927; 4 gpus for x iters 2481.719; n unique users 19892 n pro users 17860\n",
    "# 32 v13  requests 173814 clips 347628 total khrs 18.463; N gpus for 1000 iters 21.727; 4 gpus for x iters 5431.688; n unique users 43848 n pro users 39116\n",
    "# 32 v14  requests 87118 clips 174236 total khrs 8.689; N gpus for 1000 iters 10.890; 4 gpus for x iters 2722.438; n unique users 22826 n pro users 20899\n",
    "# 32 v15  requests 68523 clips 137046 total khrs 7.169; N gpus for 1000 iters 8.565; 4 gpus for x iters 1070.672; n unique users 26101 n pro users 23037\n",
    "# 32 v16  requests 88908 clips 177816 total khrs 8.876; N gpus for 1000 iters 11.114; 4 gpus for x iters 1389.188; n unique users 23225 n pro users 21243\n",
    "# 32 v17  requests 43677 clips 87354 total khrs 4.468; N gpus for 1000 iters 5.460; 4 gpus for x iters 682.453; n unique users 17554 n pro users 16917\n",
    "# using 0215 s8 v9 selection  requests 72501 clips 145002 total khrs 7.577; N gpus for 1000 iters 9.063; 4 gpus for x iters 1132.828; n unique users 19147 n pro users 17467\n",
    "# using 0217 s8 v9 selection  requests 72957 clips 145914 total khrs 7.631; N gpus for 1000 iters 9.120; 4 gpus for x iters 1139.953; n unique users 19167 n pro users 17480\n",
    "# 32 v18  requests 72957 clips 145914 total khrs 7.631; N gpus for 1000 iters 9.120; 4 gpus for x iters 1139.953; n unique users 19167 n pro users 17480\n",
    "# 32 v20  requests 52910 clips 105820 total khrs 5.572; N gpus for 1000 iters 6.614; 4 gpus for x iters 826.719; n unique users 16708 n pro users 15184\n",
    "# 32 v21  requests 54396 cqlips 108792 total khrs 5.732; N gpus for 1000 iters 6.800; 4 gpus for x iters 849.938; n unique users 17148 n pro users 15552\n",
    "# 32 v24  requests 375375 clips 750750 total khrs 35.534; N gpus for 1000 iters 46.922; 4 gpus for x iters 5865.234; n unique users 81216 n pro users 70009\n",
    "# 32 v25  requests 367757 clips 735514 total khrs 35.610; N gpus for 1000 iters 45.970; 4 gpus for x iters 5746.203; n unique users 82243 n pro users 71597\n",
    "# 32 v26  requests 130175 clips 260350 total khrs 13.203; N gpus for 1000 iters 16.272; 4 gpus for x iters 2033.984; n unique users 30091 n pro users 26580\n",
    "# 32 v28 cycle requests 365502 clips 731004 total khrs 35.369; N gpus for 1000 iters 45.688; 4 gpus for x iters 5710.969; n unique users 81990 n pro users 71365\n",
    "# 32 v29  requests 397502 clips 795004 total khrs 38.561; N gpus for 1000 iters 49.688; 4 gpus for x iters 6210.969; n unique users 87074 n pro users 74445\n",
    "# 32 v31  requests 335233 clips 670466 total khrs 33.674; N gpus for 1000 iters 41.904; 4 gpus for x iters 5238.016; n unique users 43643 n pro users 38216\n",
    "# 32 v32  equests 335233 clips 670466 total khrs 33.674; N gpus for 1000 iters 41.904; 4 gpus for x iters 5238.016; n unique users 43643 n pro users 38216\n",
    "# 32 v33  requests 335233 clips 670466 total khrs 33.674; N gpus for 1000 iters 41.904; 4 gpus for x iters 5238.016; n unique users 43643 n pro users 38216\n",
    "# 32 v34  requests 231589 clips 463178 total khrs 22.980; N gpus for 1000 iters 28.949; 4 gpus for x iters 3618.578; n unique users 65278 n pro users 56940"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.277006Z",
     "start_time": "2024-05-16T13:59:41.252105Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.320762Z",
     "iopub.status.busy": "2025-03-31T02:40:47.320510Z",
     "iopub.status.idle": "2025-03-31T02:40:47.472336Z",
     "shell.execute_reply": "2025-03-31T02:40:47.471784Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.320744Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive in playlist (67846, 157)\n"
     ]
    }
   ],
   "source": [
    "test_mask = (df_slice[\"preference\"] == True) & (\n",
    "    (df_slice[\"is_in_playlist\"] == True) | (df_slice[\"concat_in_playlist\"] == True)\n",
    ")\n",
    "print(\"positive in playlist\", df_slice[test_mask].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.473003Z",
     "iopub.status.busy": "2025-03-31T02:40:47.472854Z",
     "iopub.status.idle": "2025-03-31T02:40:47.682689Z",
     "shell.execute_reply": "2025-03-31T02:40:47.682102Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.472988Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive with likes (122892, 157)\n"
     ]
    }
   ],
   "source": [
    "test_mask = (df_slice[\"preference\"] == True) & (df_slice[\"upvote_count\"] >= 1)\n",
    "print(\"positive with likes\", df_slice[test_mask].shape)\n",
    "# positive with likes (60349, 149)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.683389Z",
     "iopub.status.busy": "2025-03-31T02:40:47.683243Z",
     "iopub.status.idle": "2025-03-31T02:40:47.697961Z",
     "shell.execute_reply": "2025-03-31T02:40:47.697520Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.683374Z"
    }
   },
   "outputs": [],
   "source": [
    "# BREAK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.323409Z",
     "start_time": "2024-05-16T13:59:41.278278Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.698583Z",
     "iopub.status.busy": "2025-03-31T02:40:47.698444Z",
     "iopub.status.idle": "2025-03-31T02:40:47.711015Z",
     "shell.execute_reply": "2025-03-31T02:40:47.710593Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.698570Z"
    }
   },
   "outputs": [],
   "source": [
    "# interesting_clips_must_be_positive_mask = (\n",
    "#     (df_slice[\"upvoted\"] == True)\n",
    "#     | (df_slice[\"has_action\"] == True)\n",
    "#     | (df_slice[\"part_of_concat\"] == True)\n",
    "# )\n",
    "# interesting_clips_must_be_not_negative_mask = (df_slice[\"downvoted\"] == False) # & (df_slice[\"dislike_count\"] < 1)\n",
    "# interesting_clips_mask = interesting_clips_must_be_positive_mask & interesting_clips_must_be_not_negative_mask\n",
    "# assert interesting_clips_mask.eq(df_slice[\"preference\"]).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.392244Z",
     "start_time": "2024-05-16T13:59:41.324472Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.711689Z",
     "iopub.status.busy": "2025-03-31T02:40:47.711557Z",
     "iopub.status.idle": "2025-03-31T02:40:47.722338Z",
     "shell.execute_reply": "2025-03-31T02:40:47.721920Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.711676Z"
    }
   },
   "outputs": [],
   "source": [
    "# save positive ids\n",
    "# positive_preference_ids = df_slice[df_slice[\"preference\"] == False][\"s3_id\"].to_json(orient='values')\n",
    "# with open('/home/tony/Data/Preference/7b_v2/7v_v20_full_recut_id_negative.json', 'w') as file:\n",
    "#     file.write(positive_preference_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.722917Z",
     "iopub.status.busy": "2025-03-31T02:40:47.722787Z",
     "iopub.status.idle": "2025-03-31T02:40:47.733659Z",
     "shell.execute_reply": "2025-03-31T02:40:47.733237Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.722904Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[df_slice[\"task\"] == \"upload_extend\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.734208Z",
     "iopub.status.busy": "2025-03-31T02:40:47.734077Z",
     "iopub.status.idle": "2025-03-31T02:40:47.744990Z",
     "shell.execute_reply": "2025-03-31T02:40:47.744565Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.734194Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[\"promotion\"].value_counts()\n",
    "# df_slice[\"source\"].value_counts()\n",
    "# df_slice[\"gpt_description_prompt\"].isna().value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.745583Z",
     "iopub.status.busy": "2025-03-31T02:40:47.745448Z",
     "iopub.status.idle": "2025-03-31T02:40:47.815843Z",
     "shell.execute_reply": "2025-03-31T02:40:47.815418Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.745569Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "person_type\n",
       "1    321812\n",
       "0    127540\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_slice.groupby(\"user_id\")[\"gpt_description_prompt\"].transform(\"any\").value_counts()\n",
    "df_slice[\"person_type\"] = df_slice[\"gpt_description_prompt\"].isna().astype(int)\n",
    "df_slice[\"person_type\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.816461Z",
     "iopub.status.busy": "2025-03-31T02:40:47.816321Z",
     "iopub.status.idle": "2025-03-31T02:40:47.963907Z",
     "shell.execute_reply": "2025-03-31T02:40:47.963352Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.816447Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "449352\n"
     ]
    }
   ],
   "source": [
    "result_dict = dict(zip(df_slice[\"id\"], df_slice[\"person_type\"]))\n",
    "print(len(result_dict))\n",
    "# with open(os.path.join(OUT_DATA_DIR, \"person_info.json\"), \"w\") as fp:\n",
    "#     json.dump(result_dict, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:47.964598Z",
     "iopub.status.busy": "2025-03-31T02:40:47.964450Z",
     "iopub.status.idle": "2025-03-31T02:40:48.003303Z",
     "shell.execute_reply": "2025-03-31T02:40:48.002867Z",
     "shell.execute_reply.started": "2025-03-31T02:40:47.964582Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(150613, 64856, 411, 78538800)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    df[\"user_id\"].nunique(),\n",
    "    df_slice[\"user_id\"].nunique(),\n",
    "    df_slice[\"user_id\"].min(),\n",
    "    df_slice[\"user_id\"].max(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:48.003955Z",
     "iopub.status.busy": "2025-03-31T02:40:48.003816Z",
     "iopub.status.idle": "2025-03-31T02:40:48.017124Z",
     "shell.execute_reply": "2025-03-31T02:40:48.016692Z",
     "shell.execute_reply.started": "2025-03-31T02:40:48.003941Z"
    }
   },
   "outputs": [],
   "source": [
    "# with open(f\"/app/suno/data/dpo/13b_s32_v29/quality/full_pair_quality.json\", \"r\") as f:\n",
    "#    result = json.load(f)\n",
    "\n",
    "# # need to take care of the tail only\n",
    "# def fast_slope(y):\n",
    "#     \"\"\"\n",
    "#     Calculate the slope of a linear regression line extremely quickly.\n",
    "\n",
    "#     Parameters:\n",
    "#     y (list or array): List of float values\n",
    "\n",
    "#     Returns:\n",
    "#     float: The slope of the linear regression line\n",
    "#     \"\"\"\n",
    "#     # detect if x is long enough\n",
    "#     if len(y) < 3 * 60 / 5:\n",
    "#         return 0\n",
    "#     y = np.asarray(y, dtype=np.float64)\n",
    "#     n = len(y)\n",
    "\n",
    "#     # Fast calculation using vectorized operations\n",
    "#     x = np.arange(n)\n",
    "#     x_mean = (n - 1) / 2  # Analytical mean of range(n)\n",
    "#     y_mean = np.mean(y)\n",
    "\n",
    "#     # Optimized computation of slope using vectorized operations\n",
    "#     # Formula: slope = sum((x_i - x_mean) * (y_i - y_mean)) / sum((x_i - x_mean)^2)\n",
    "#     numerator = np.sum(y * x) - n * x_mean * y_mean\n",
    "#     denominator = np.sum(x * x) - n * x_mean * x_mean\n",
    "\n",
    "#     return numerator / (denominator + 0.001)\n",
    "\n",
    "# df_slice[\"ave_ear\"] = df_slice[\"s3_id\"].map(lambda x: np.mean(result[x]) if x in result and result[x] else 0)\n",
    "# df_slice[\"decay_ear\"] = df_slice[\"s3_id\"].map(lambda x: fast_slope(result[x]) if x in result and result[x] else 0)\n",
    "\n",
    "# df_slice[\"ave_ear_diff\"] = df_slice[\"ave_ear\"].diff()\n",
    "# df_slice[\"decay_ear_diff\"] = df_slice[\"decay_ear\"].diff()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:48.017735Z",
     "iopub.status.busy": "2025-03-31T02:40:48.017597Z",
     "iopub.status.idle": "2025-03-31T02:40:48.029815Z",
     "shell.execute_reply": "2025-03-31T02:40:48.029398Z",
     "shell.execute_reply.started": "2025-03-31T02:40:48.017721Z"
    }
   },
   "outputs": [],
   "source": [
    "# _ = df_slice[df_slice[\"preference\"]][\"ave_ear_diff\"].hist(bins=np.linspace(-10, 10, 100))\n",
    "# print(df_slice[df_slice[\"preference\"]][\"ave_ear_diff\"].quantile(0.05))\n",
    "# plt.show()\n",
    "# _ = df_slice[df_slice[\"preference\"]][\"decay_ear\"].hist(bins=np.linspace(-1, 1, 100))\n",
    "# print(df_slice[df_slice[\"preference\"]][\"decay_ear\"].quantile(0.05))\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:48.030517Z",
     "iopub.status.busy": "2025-03-31T02:40:48.030384Z",
     "iopub.status.idle": "2025-03-31T02:40:48.041454Z",
     "shell.execute_reply": "2025-03-31T02:40:48.041034Z",
     "shell.execute_reply.started": "2025-03-31T02:40:48.030503Z"
    }
   },
   "outputs": [],
   "source": [
    "# extra cut here to filter out the tail of the qualities\n",
    "# unique_non_decay_positive_requets = df_slice[(df_slice[\"preference\"]) &(df_slice[\"ave_ear_diff\"] > -3.5) & (df_slice[\"decay_ear\"] > -0.21)][\"request_id\"]\n",
    "\n",
    "# df_slice[df_slice[\"request_id\"].isin(set(unique_non_decay_positive_requets))].shape[0]/df_slice.shape[0]\n",
    "\n",
    "# df_slice = df_slice[df_slice[\"request_id\"].isin(set(unique_non_decay_positive_requets))].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:48.052049Z",
     "iopub.status.busy": "2025-03-31T02:40:48.051911Z",
     "iopub.status.idle": "2025-03-31T02:40:48.065972Z",
     "shell.execute_reply": "2025-03-31T02:40:48.065565Z",
     "shell.execute_reply.started": "2025-03-31T02:40:48.052035Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "person_type\n",
       "1    321812\n",
       "0    127540\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_slice[\"person_type\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T02:40:48.066676Z",
     "iopub.status.busy": "2025-03-31T02:40:48.066545Z",
     "iopub.status.idle": "2025-03-31T02:40:48.360322Z",
     "shell.execute_reply": "2025-03-31T02:40:48.359565Z",
     "shell.execute_reply.started": "2025-03-31T02:40:48.066662Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(449352, 158)\n",
      "task\n",
      "                 368064\n",
      "extend            53196\n",
      "upload_extend     28092\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[33], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(df_slice\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m      3\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----> 4\u001b[0m \u001b[43mBREAK\u001b[49m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'BREAK' is not defined"
     ]
    }
   ],
   "source": [
    "# df_slice.to_pickle(\"/home/tony/Data/Preference/13b_v32/interesting_clips_v4_h_s_32_20250330_full_long_bluejay_r2.pkl\")\n",
    "print(df_slice.shape)\n",
    "print(df_slice[\"task\"].value_counts())\n",
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 30b loss filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:32.329733Z",
     "iopub.status.busy": "2025-03-31T03:51:32.329367Z",
     "iopub.status.idle": "2025-03-31T03:51:33.153418Z",
     "shell.execute_reply": "2025-03-31T03:51:33.152935Z",
     "shell.execute_reply.started": "2025-03-31T03:51:32.329715Z"
    }
   },
   "outputs": [],
   "source": [
    "# info_val = read_jsonl(os.path.join(\"/app/suno/data/dpo/13b_s32_v31\", f\"meta_val.jsonl\"))\n",
    "# info_tr = read_jsonl(os.path.join(\"/app/suno/data/dpo/13b_s32_v31\", f\"meta_tr.jsonl\"))\n",
    "# with open(\"/app/suno/data/dpo/13b_s32_v31/30b_t6_bt16_cached_loss.json\", \"r\") as fp:\n",
    "#     loss_30b_lookup = json.load(fp)\n",
    "# with open(\"/app/suno/data/dpo/13b_s32_v31/13b_dpo_s32_cached_loss.json\", \"r\") as fp:\n",
    "#     loss_13b_lookup = json.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.154484Z",
     "iopub.status.busy": "2025-03-31T03:51:33.154144Z",
     "iopub.status.idle": "2025-03-31T03:51:33.166479Z",
     "shell.execute_reply": "2025-03-31T03:51:33.166055Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.154469Z"
    }
   },
   "outputs": [],
   "source": [
    "# id_to_30b_loss = {}\n",
    "# for clip_id, loss in loss_30b_lookup[\"val\"].items():\n",
    "#     id_to_30b_loss[info_val[int(clip_id)][\"id\"]] = loss['semantic_0']\n",
    "# for clip_id, loss in loss_30b_lookup[\"train\"].items():\n",
    "#     id_to_30b_loss[info_tr[int(clip_id)][\"id\"]] = loss['semantic_0']\n",
    "\n",
    "# id_to_13b_loss = {}\n",
    "# for clip_id, loss in loss_13b_lookup[\"val\"].items():\n",
    "#     id_to_13b_loss[info_val[int(clip_id)][\"id\"]] = loss['semantic_0']\n",
    "# for clip_id, loss in loss_13b_lookup[\"train\"].items():\n",
    "#     id_to_13b_loss[info_tr[int(clip_id)][\"id\"]] = loss['semantic_0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.167078Z",
     "iopub.status.busy": "2025-03-31T03:51:33.166942Z",
     "iopub.status.idle": "2025-03-31T03:51:33.177864Z",
     "shell.execute_reply": "2025-03-31T03:51:33.177446Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.167064Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[\"13b_loss\"] = df_slice[\"id\"].map(id_to_13b_loss)\n",
    "# df_slice[\"30b_loss\"] = df_slice[\"id\"].map(id_to_30b_loss)\n",
    "# df_slice[\"13b_loss_diff\"] = df_slice[\"13b_loss\"].diff()\n",
    "# df_slice[\"30b_loss_diff\"] = df_slice[\"30b_loss\"].diff()\n",
    "# df_slice[df_slice[\"preference\"]][[\"13b_loss\", \"30b_loss\", \"13b_loss_diff\", \"30b_loss_diff\" ]].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.179035Z",
     "iopub.status.busy": "2025-03-31T03:51:33.178777Z",
     "iopub.status.idle": "2025-03-31T03:51:33.189634Z",
     "shell.execute_reply": "2025-03-31T03:51:33.189213Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.179020Z"
    }
   },
   "outputs": [],
   "source": [
    "# plt.hist(df_slice[df_slice[\"preference\"]][\"30b_loss_diff\"], bins=np.linspace(-2, 2, 100))\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.190311Z",
     "iopub.status.busy": "2025-03-31T03:51:33.190181Z",
     "iopub.status.idle": "2025-03-31T03:51:33.200921Z",
     "shell.execute_reply": "2025-03-31T03:51:33.200490Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.190298Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[df_slice[\"preference\"]][\"30b_loss_diff\"].describe()\n",
    "# df_slice[df_slice[\"preference\"]][\"30b_loss_diff\"].quantile([0.01, 0.1, 0.9, 0.99])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.201515Z",
     "iopub.status.busy": "2025-03-31T03:51:33.201380Z",
     "iopub.status.idle": "2025-03-31T03:51:33.212205Z",
     "shell.execute_reply": "2025-03-31T03:51:33.211773Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.201502Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[df_slice[\"preference\"]][\"30b_loss\"].quantile([0.01, 0.99])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.212813Z",
     "iopub.status.busy": "2025-03-31T03:51:33.212668Z",
     "iopub.status.idle": "2025-03-31T03:51:33.223626Z",
     "shell.execute_reply": "2025-03-31T03:51:33.223203Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.212799Z"
    }
   },
   "outputs": [],
   "source": [
    "# potential_good_requests = df_slice[\n",
    "#     (df_slice[\"preference\"]) & (df_slice[\"30b_loss_diff\"] > -0.35) & (df_slice[\"30b_loss_diff\"] < 0.75)\n",
    "#     & (df_slice[\"30b_loss\"] > 1.0) & (df_slice[\"30b_loss\"] < 2.7)\n",
    "# ][\"request_id\"].unique()\n",
    "# print(df_slice.shape)\n",
    "# df_slice_2 = df_slice[df_slice[\"request_id\"].isin(potential_good_requests)].copy()\n",
    "# print(df_slice_2.shape)\n",
    "# # df_slice[(df_slice[\"preference\"]) & (df_slice[\"30b_loss_diff\"] > 1.0)].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.224304Z",
     "iopub.status.busy": "2025-03-31T03:51:33.224170Z",
     "iopub.status.idle": "2025-03-31T03:51:33.234839Z",
     "shell.execute_reply": "2025-03-31T03:51:33.234420Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.224290Z"
    }
   },
   "outputs": [],
   "source": [
    "# plt.hist(df_slice_2[df_slice_2[\"preference\"]][\"13b_loss_diff\"], bins=np.linspace(-2, 2, 100), alpha=0.5)\n",
    "# plt.hist(df_slice[df_slice[\"preference\"]][\"13b_loss_diff\"], bins=np.linspace(-2, 2, 100), alpha=0.5)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.235575Z",
     "iopub.status.busy": "2025-03-31T03:51:33.235442Z",
     "iopub.status.idle": "2025-03-31T03:51:33.246109Z",
     "shell.execute_reply": "2025-03-31T03:51:33.245689Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.235562Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice = df_slice_2.copy()"
   ]
  },
  {
   "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": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932296Z",
     "start_time": "2024-05-16T13:59:41.932287Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.247826Z",
     "iopub.status.busy": "2025-03-31T03:51:33.247566Z",
     "iopub.status.idle": "2025-03-31T03:51:33.258382Z",
     "shell.execute_reply": "2025-03-31T03:51:33.257958Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.247811Z"
    }
   },
   "outputs": [],
   "source": [
    "# don't have continue at\n",
    "# df_slice[df_slice[\"continue_at\"].isna()][\"request_id\"].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932966Z",
     "start_time": "2024-05-16T13:59:41.932957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.258978Z",
     "iopub.status.busy": "2025-03-31T03:51:33.258847Z",
     "iopub.status.idle": "2025-03-31T03:51:33.340498Z",
     "shell.execute_reply": "2025-03-31T03:51:33.340018Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.258965Z"
    }
   },
   "outputs": [],
   "source": [
    "final_filtered_requests = df_slice[\"request_id\"].unique()\n",
    "print(len(final_filtered_requests))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.933558Z",
     "start_time": "2024-05-16T13:59:41.933550Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.341125Z",
     "iopub.status.busy": "2025-03-31T03:51:33.340981Z",
     "iopub.status.idle": "2025-03-31T03:51:33.352261Z",
     "shell.execute_reply": "2025-03-31T03:51:33.351837Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.341111Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice.to_csv(\"/home/tony/Data/Preference/7b_v2/7b_before_recode_20240412\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934277Z",
     "start_time": "2024-05-16T13:59:41.934268Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:33.352840Z",
     "iopub.status.busy": "2025-03-31T03:51:33.352708Z",
     "iopub.status.idle": "2025-03-31T03:51:34.845239Z",
     "shell.execute_reply": "2025-03-31T03:51:34.844648Z",
     "shell.execute_reply.started": "2025-03-31T03:51:33.352826Z"
    }
   },
   "outputs": [],
   "source": [
    "train_requests, val_requests = train_test_split(\n",
    "    sorted(list(final_filtered_requests)), test_size=0.01, random_state=42\n",
    ")\n",
    "print(len(train_requests), len(val_requests))\n",
    "\n",
    "train_df = df_slice[df_slice[\"request_id\"].isin(set(train_requests))].copy()\n",
    "val_df = df_slice[df_slice[\"request_id\"].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",
    "\n",
    "print(train_df.shape, val_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:34.845950Z",
     "iopub.status.busy": "2025-03-31T03:51:34.845802Z",
     "iopub.status.idle": "2025-03-31T03:51:34.860826Z",
     "shell.execute_reply": "2025-03-31T03:51:34.860382Z",
     "shell.execute_reply.started": "2025-03-31T03:51:34.845934Z"
    }
   },
   "outputs": [],
   "source": [
    "# BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actually make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934954Z",
     "start_time": "2024-05-16T13:59:41.934946Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:34.861549Z",
     "iopub.status.busy": "2025-03-31T03:51:34.861409Z",
     "iopub.status.idle": "2025-03-31T03:51:34.873882Z",
     "shell.execute_reply": "2025-03-31T03:51:34.873454Z",
     "shell.execute_reply.started": "2025-03-31T03:51:34.861535Z"
    }
   },
   "outputs": [],
   "source": [
    "# val_df[[\"request_id\", \"metadata\", \"updated_at\", \"user_id\", \"preference\"]].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.935620Z",
     "start_time": "2024-05-16T13:59:41.935613Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:34.874472Z",
     "iopub.status.busy": "2025-03-31T03:51:34.874341Z",
     "iopub.status.idle": "2025-03-31T03:51:50.837936Z",
     "shell.execute_reply": "2025-03-31T03:51:50.837417Z",
     "shell.execute_reply.started": "2025-03-31T03:51:34.874458Z"
    }
   },
   "outputs": [],
   "source": [
    "total_duration = 0\n",
    "for i, row in tqdm(train_df.iterrows(), total=len(train_df)):\n",
    "    # we need to alternate between preference: neg, pos\n",
    "    # print(i, row)\n",
    "    try:\n",
    "        assert row[\"preference\"] == (i % 2 == 1)\n",
    "    except:\n",
    "        print(i, row)\n",
    "    total_duration += row[\"duration\"]\n",
    "print(\n",
    "    f\"{round(total_duration / 60 / 60):,} hours of {train_df.shape[0]} clips, {train_df.shape[0] / 8 / 4 / 1000} nodes, {train_df.shape[0] / 8 / 8 / 4} steps\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936268Z",
     "start_time": "2024-05-16T13:59:41.936260Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:51:50.838634Z",
     "iopub.status.busy": "2025-03-31T03:51:50.838483Z",
     "iopub.status.idle": "2025-03-31T03:53:04.637015Z",
     "shell.execute_reply": "2025-03-31T03:53:04.636291Z",
     "shell.execute_reply.started": "2025-03-31T03:51:50.838619Z"
    }
   },
   "outputs": [],
   "source": [
    "make_dataset(val_df, OUT_DATA_DIR, is_val=True, npz_dir=NPZ_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.936964Z",
     "start_time": "2024-05-16T13:59:41.936957Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-31T03:53:04.637912Z",
     "iopub.status.busy": "2025-03-31T03:53:04.637738Z"
    }
   },
   "outputs": [],
   "source": [
    "make_dataset(train_df, OUT_DATA_DIR, is_val=False, npz_dir=NPZ_DIR)"
   ]
  },
  {
   "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": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.937879Z",
     "start_time": "2024-05-16T13:59:41.937870Z"
    }
   },
   "outputs": [],
   "source": [
    "# verify\n",
    "mm = np.memmap(os.path.join(OUT_DATA_DIR, f\"data_val.bin\"), dtype=np.uint16, mode=\"r\")\n",
    "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"meta_val.jsonl\"))\n",
    "test_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_val.json\"))\n",
    "mm = mm.reshape(-1, 6016, 13)\n",
    "assert len(mm) == len(test_metas)\n",
    "assert mm[:100, :, 0].min() >= 0\n",
    "assert mm[:100, :, 0].max() <= 4000\n",
    "assert mm[:100, :, 1:].min() >= 0\n",
    "assert mm[:100, :, 1:].max() <= 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.938629Z",
     "start_time": "2024-05-16T13:59:41.938621Z"
    }
   },
   "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": 53,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939205Z",
     "start_time": "2024-05-16T13:59:41.939198Z"
    }
   },
   "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": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.939977Z",
     "start_time": "2024-05-16T13:59:41.939969Z"
    }
   },
   "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": 55,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.940610Z",
     "start_time": "2024-05-16T13:59:41.940603Z"
    }
   },
   "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": 56,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941167Z",
     "start_time": "2024-05-16T13:59:41.941159Z"
    }
   },
   "outputs": [],
   "source": [
    "# original_npz_path = f\"/app/suno/data/dpo/7b_npz/{test_metas[idx]['id']}.npz\"\n",
    "# original_npz_path = \"/app/suno/data/dpo/7b_npz/729c3011-f672-4ccd-8d82-1cbf2b52ff69.npz\"\n",
    "# original_arr = np.load(original_npz_path)[\"v2_raw\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941801Z",
     "start_time": "2024-05-16T13:59:41.941793Z"
    }
   },
   "outputs": [],
   "source": [
    "def validation_on_metas(input_metas):\n",
    "    total_bad = 0\n",
    "    total_good = 0\n",
    "    for idx in range(len(input_metas)):\n",
    "        if idx % 2 == 0:\n",
    "            pos_idx = idx + 1\n",
    "            if input_metas[idx].get(\"tags\") != input_metas[pos_idx].get(\"tags\"):\n",
    "                # print(test_metas[idx].get(\"text\") == test_metas[pos_idx].get(\"text\"), test_metas[idx].get(\"tags\"), test_metas[pos_idx].get(\"tags\"))\n",
    "                total_bad += 1\n",
    "            else:\n",
    "                total_good += 1\n",
    "    print(total_good, total_bad)\n",
    "    return\n",
    "\n",
    "\n",
    "validation_on_metas(test_metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.942520Z",
     "start_time": "2024-05-16T13:59:41.942511Z"
    }
   },
   "outputs": [],
   "source": [
    "train_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_tr.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.943072Z",
     "start_time": "2024-05-16T13:59:41.943065Z"
    }
   },
   "outputs": [],
   "source": [
    "n_neg_tr = train_info[\"perference_0\"][\"idx_list\"]\n",
    "n_pos_tr = train_info[\"perference_1\"][\"idx_list\"]\n",
    "assert len(n_pos_tr) == len(n_neg_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.944246Z",
     "start_time": "2024-05-16T13:59:41.944237Z"
    }
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945249Z",
     "start_time": "2024-05-16T13:59:41.945241Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 8 / 4 / 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945972Z",
     "start_time": "2024-05-16T13:59:41.945964Z"
    }
   },
   "outputs": [],
   "source": [
    "# !cd /home/tony/Work/tony/slurm/13b_dpo && sbatch sbatch_ipo_13b_s32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "# Basic file copy\n",
    "shutil.copy(\n",
    "    \"/home/tony/Work/tony/Preference/make_dataset_auk_13b_v34.ipynb\",\n",
    "    os.path.join(OUT_DATA_DIR, \"make_dataset.ipynb\"),\n",
    ")\n",
    "print(\"Cache kept!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# some gymathtics loading prev data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.946562Z",
     "start_time": "2024-05-16T13:59:41.946555Z"
    }
   },
   "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": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_df[\"lang\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "# # with open(f\"/home/tony/Data/Preference/13b_s32_v29/full_pair_quality.json\", \"r\") as f:\n",
    "# #    result = json.load(f)\n",
    "# result = {}\n",
    "# print(len(result))\n",
    "# for job_idx in range(8):\n",
    "#     with open(f\"/app/suno/data/dpo/13b_s32_v29/quality/full_pair_quality_{job_idx}.json\", \"r\") as fp:\n",
    "#         current_result = json.load(fp)\n",
    "#         result.update(current_result)\n",
    "# print(len(result))\n",
    "# with open(f\"/app/suno/data/dpo/13b_s32_v29/quality/full_pair_quality.json\", \"w\") as f:\n",
    "#     json.dump(result, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env_dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
