{
 "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-18T03:49:07.254713Z",
     "iopub.status.busy": "2025-03-18T03:49:07.254523Z",
     "iopub.status.idle": "2025-03-18T03:49:09.419225Z",
     "shell.execute_reply": "2025-03-18T03:49:09.418589Z",
     "shell.execute_reply.started": "2025-03-18T03:49:07.254698Z"
    }
   },
   "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_diff 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",
    "from suno_utils.audio import Audio\n",
    "from suno_analytics.preference_helper import get_preference_counts\n",
    "\n",
    "pd.set_option(\"display.max_rows\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:21.082172Z",
     "start_time": "2024-05-16T13:58:21.041926Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:09.421644Z",
     "iopub.status.busy": "2025-03-18T03:49:09.421286Z",
     "iopub.status.idle": "2025-03-18T03:49:09.434470Z",
     "shell.execute_reply": "2025-03-18T03:49:09.434025Z",
     "shell.execute_reply.started": "2025-03-18T03:49:09.421627Z"
    }
   },
   "outputs": [],
   "source": [
    "OUT_DATA_DIR = \"/app/suno/data/dpo/diff_v6_t2_comb/\"\n",
    "os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "NPZ_DIR = \"/app/suno/data/dpo/diff_v6\""
   ]
  },
  {
   "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-18T03:49:09.436210Z",
     "iopub.status.busy": "2025-03-18T03:49:09.436094Z",
     "iopub.status.idle": "2025-03-18T03:49:10.504144Z",
     "shell.execute_reply": "2025-03-18T03:49:10.503454Z",
     "shell.execute_reply.started": "2025-03-18T03:49:09.436199Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preference data shape (81394, 89)\n",
      "unique users 22476\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_pickle(\n",
    "    \"/home/tony/Data/Preference/up_v6/interesting_clips_up_u_6_20250316_full.pkl\"\n",
    ")  # , engine='python')\n",
    "print(\"Preference data shape\", df.shape)\n",
    "print(\"unique users\", df[\"user_id\"].nunique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:10.506386Z",
     "iopub.status.busy": "2025-03-18T03:49:10.506263Z",
     "iopub.status.idle": "2025-03-18T03:49:10.509946Z",
     "shell.execute_reply": "2025-03-18T03:49:10.509478Z",
     "shell.execute_reply.started": "2025-03-18T03:49:10.506372Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "is_public\n",
      "False    74842\n",
      "True      6552\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\"]]\n",
    "# print(df[\"is_public\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:10.511554Z",
     "iopub.status.busy": "2025-03-18T03:49:10.511445Z",
     "iopub.status.idle": "2025-03-18T03:49:10.562489Z",
     "shell.execute_reply": "2025-03-18T03:49:10.561906Z",
     "shell.execute_reply.started": "2025-03-18T03:49:10.511543Z"
    }
   },
   "outputs": [],
   "source": [
    "df[\"upsample_clip_id\"] = df[\"metadata\"].apply(lambda x: x.get(\"upsample_clip_id\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.199480Z",
     "start_time": "2024-05-16T13:58:53.963687Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:10.563188Z",
     "iopub.status.busy": "2025-03-18T03:49:10.563060Z",
     "iopub.status.idle": "2025-03-18T03:49:17.025017Z",
     "shell.execute_reply": "2025-03-18T03:49:17.024338Z",
     "shell.execute_reply.started": "2025-03-18T03:49:10.563175Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "97099\n",
      "32971\n",
      "64128\n",
      "pre-downloaded df (81394, 90)\n",
      "downloaded df (64596, 90)\n",
      "vae downloaded df (61118, 90)\n"
     ]
    }
   ],
   "source": [
    "all_converted_paths = os.listdir(NPZ_DIR)\n",
    "print(len(all_converted_paths))\n",
    "\n",
    "converted_paths = set(\n",
    "    [f.replace(\".npz\", \"\") for f in all_converted_paths if \"vae\" not in f]\n",
    ")\n",
    "print(len(converted_paths))\n",
    "vae_converted_paths = set(\n",
    "    [f.replace(\"_vae.npz\", \"\") for f in all_converted_paths if \"vae\" in f]\n",
    ")\n",
    "print(len(vae_converted_paths))\n",
    "\n",
    "print(\"pre-downloaded df\", df.shape)\n",
    "df[df[\"upsample_clip_id\"].isin(converted_paths)].shape\n",
    "df = df[df[\"upsample_clip_id\"].isin(converted_paths)].copy()\n",
    "print(\"downloaded df\", df.shape)\n",
    "df[df[\"s3_id\"].isin(vae_converted_paths)].shape\n",
    "df = df[df[\"s3_id\"].isin(vae_converted_paths)].copy()\n",
    "print(\"vae downloaded df\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.467253Z",
     "start_time": "2024-05-16T13:58:56.207647Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:17.026582Z",
     "iopub.status.busy": "2025-03-18T03:49:17.026320Z",
     "iopub.status.idle": "2025-03-18T03:49:17.041827Z",
     "shell.execute_reply": "2025-03-18T03:49:17.041372Z",
     "shell.execute_reply.started": "2025-03-18T03:49:17.026564Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "is_up\n",
       "True    61118\n",
       "Name: count, dtype: int64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"is_up\"] = df[\"model_name\"].str.contains(\"up\")\n",
    "df[\"is_up\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LET's do the data prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.592883Z",
     "start_time": "2024-05-16T13:58:56.470781Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:17.042697Z",
     "iopub.status.busy": "2025-03-18T03:49:17.042351Z",
     "iopub.status.idle": "2025-03-18T03:49:17.060831Z",
     "shell.execute_reply": "2025-03-18T03:49:17.060304Z",
     "shell.execute_reply.started": "2025-03-18T03:49:17.042682Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preference  model_name     \n",
      "False       chirp-v4-up-u-6    30559\n",
      "True        chirp-v4-up-u-6    30559\n",
      "Name: count, dtype: int64\n",
      "(61118, 91)\n",
      "(61118, 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-up-u-6\"])]\n",
    "print(df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:58:56.909539Z",
     "start_time": "2024-05-16T13:58:56.595736Z"
    },
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:17.061548Z",
     "iopub.status.busy": "2025-03-18T03:49:17.061415Z",
     "iopub.status.idle": "2025-03-18T03:49:17.139354Z",
     "shell.execute_reply": "2025-03-18T03:49:17.138693Z",
     "shell.execute_reply.started": "2025-03-18T03:49:17.061535Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(61118, 91)\n",
      "(61118, 91)\n",
      "preference  model_name     \n",
      "False       chirp-v4-up-u-6    30559\n",
      "True        chirp-v4-up-u-6    30559\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": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:17.140123Z",
     "iopub.status.busy": "2025-03-18T03:49:17.139978Z",
     "iopub.status.idle": "2025-03-18T03:49:22.162780Z",
     "shell.execute_reply": "2025-03-18T03:49:22.162064Z",
     "shell.execute_reply.started": "2025-03-18T03:49:17.140108Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total pair quality scores: 41822\n",
      "Total unpacked pair quality scores: 484088\n"
     ]
    }
   ],
   "source": [
    "with open(\"/home/tony/Data/Preference/up_v6/full_pair_quality.json\", \"r\") as file:\n",
    "    full_pair_quality = json.load(file)\n",
    "print(\"Total pair quality scores:\", len(full_pair_quality))\n",
    "\n",
    "unpacked_pair_quality = {}\n",
    "for request_id, pairs_of_qualities in full_pair_quality.items():\n",
    "    for clip_id, pair_quality in pairs_of_qualities.items():\n",
    "        unpacked_pair_quality[clip_id] = pair_quality\n",
    "print(\"Total unpacked pair quality scores:\", len(unpacked_pair_quality))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-18T03:49:22.163695Z",
     "iopub.status.busy": "2025-03-18T03:49:22.163433Z",
     "iopub.status.idle": "2025-03-18T03:49:22.776630Z",
     "shell.execute_reply": "2025-03-18T03:49:22.774829Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.163678Z"
    }
   },
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'pref'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[11], line 34\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m11\u001b[39m)]\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [\n\u001b[1;32m      6\u001b[0m         \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpref\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m      7\u001b[0m         \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshimmer_score\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m         \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclips_per_second\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     17\u001b[0m     ]\n\u001b[1;32m     20\u001b[0m df[\n\u001b[1;32m     21\u001b[0m     [\n\u001b[1;32m     22\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpair_quality\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     23\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtotal_shimmer_score\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     24\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mloudness_factor\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     25\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mspectral_character\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     26\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mspectral_centroid\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     27\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbass_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     28\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmid_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     29\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhigh_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     30\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstereo_width\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     31\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtotal_clips\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     32\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclips_per_second\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     33\u001b[0m     ]\n\u001b[0;32m---> 34\u001b[0m ] \u001b[38;5;241m=\u001b[39m pd\u001b[38;5;241m.\u001b[39mDataFrame(\u001b[43mdf\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43ms3_id\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mget_audio_quality_measures\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mtolist(), index\u001b[38;5;241m=\u001b[39mdf\u001b[38;5;241m.\u001b[39mindex)\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/series.py:4924\u001b[0m, in \u001b[0;36mSeries.apply\u001b[0;34m(self, func, convert_dtype, args, by_row, **kwargs)\u001b[0m\n\u001b[1;32m   4789\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mapply\u001b[39m(\n\u001b[1;32m   4790\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m   4791\u001b[0m     func: AggFuncType,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   4796\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs,\n\u001b[1;32m   4797\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DataFrame \u001b[38;5;241m|\u001b[39m Series:\n\u001b[1;32m   4798\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   4799\u001b[0m \u001b[38;5;124;03m    Invoke function on values of Series.\u001b[39;00m\n\u001b[1;32m   4800\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   4915\u001b[0m \u001b[38;5;124;03m    dtype: float64\u001b[39;00m\n\u001b[1;32m   4916\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m   4917\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mSeriesApply\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   4918\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4919\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4920\u001b[0m \u001b[43m        \u001b[49m\u001b[43mconvert_dtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconvert_dtype\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4921\u001b[0m \u001b[43m        \u001b[49m\u001b[43mby_row\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mby_row\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4922\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4923\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m-> 4924\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/apply.py:1427\u001b[0m, in \u001b[0;36mSeriesApply.apply\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1424\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mapply_compat()\n\u001b[1;32m   1426\u001b[0m \u001b[38;5;66;03m# self.func is Callable\u001b[39;00m\n\u001b[0;32m-> 1427\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply_standard\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/apply.py:1507\u001b[0m, in \u001b[0;36mSeriesApply.apply_standard\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1501\u001b[0m \u001b[38;5;66;03m# row-wise access\u001b[39;00m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# apply doesn't have a `na_action` keyword and for backward compat reasons\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m \u001b[38;5;66;03m# we need to give `na_action=\"ignore\"` for categorical data.\u001b[39;00m\n\u001b[1;32m   1504\u001b[0m \u001b[38;5;66;03m# TODO: remove the `na_action=\"ignore\"` when that default has been changed in\u001b[39;00m\n\u001b[1;32m   1505\u001b[0m \u001b[38;5;66;03m#  Categorical (GH51645).\u001b[39;00m\n\u001b[1;32m   1506\u001b[0m action \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mignore\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(obj\u001b[38;5;241m.\u001b[39mdtype, CategoricalDtype) \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 1507\u001b[0m mapped \u001b[38;5;241m=\u001b[39m \u001b[43mobj\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_map_values\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1508\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmapper\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcurried\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mna_action\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43maction\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconvert\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconvert_dtype\u001b[49m\n\u001b[1;32m   1509\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1511\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(mapped) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(mapped[\u001b[38;5;241m0\u001b[39m], ABCSeries):\n\u001b[1;32m   1512\u001b[0m     \u001b[38;5;66;03m# GH#43986 Need to do list(mapped) in order to get treated as nested\u001b[39;00m\n\u001b[1;32m   1513\u001b[0m     \u001b[38;5;66;03m#  See also GH#25959 regarding EA support\u001b[39;00m\n\u001b[1;32m   1514\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m obj\u001b[38;5;241m.\u001b[39m_constructor_expanddim(\u001b[38;5;28mlist\u001b[39m(mapped), index\u001b[38;5;241m=\u001b[39mobj\u001b[38;5;241m.\u001b[39mindex)\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/base.py:921\u001b[0m, in \u001b[0;36mIndexOpsMixin._map_values\u001b[0;34m(self, mapper, na_action, convert)\u001b[0m\n\u001b[1;32m    918\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(arr, ExtensionArray):\n\u001b[1;32m    919\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m arr\u001b[38;5;241m.\u001b[39mmap(mapper, na_action\u001b[38;5;241m=\u001b[39mna_action)\n\u001b[0;32m--> 921\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43malgorithms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap_array\u001b[49m\u001b[43m(\u001b[49m\u001b[43marr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmapper\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mna_action\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mna_action\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconvert\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconvert\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/pandas/core/algorithms.py:1743\u001b[0m, in \u001b[0;36mmap_array\u001b[0;34m(arr, mapper, na_action, convert)\u001b[0m\n\u001b[1;32m   1741\u001b[0m values \u001b[38;5;241m=\u001b[39m arr\u001b[38;5;241m.\u001b[39mastype(\u001b[38;5;28mobject\u001b[39m, copy\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m   1742\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m na_action \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 1743\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mlib\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap_infer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalues\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmapper\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconvert\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconvert\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1744\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1745\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m lib\u001b[38;5;241m.\u001b[39mmap_infer_mask(\n\u001b[1;32m   1746\u001b[0m         values, mapper, mask\u001b[38;5;241m=\u001b[39misna(values)\u001b[38;5;241m.\u001b[39mview(np\u001b[38;5;241m.\u001b[39muint8), convert\u001b[38;5;241m=\u001b[39mconvert\n\u001b[1;32m   1747\u001b[0m     )\n",
      "File \u001b[0;32mlib.pyx:2972\u001b[0m, in \u001b[0;36mpandas._libs.lib.map_infer\u001b[0;34m()\u001b[0m\n",
      "Cell \u001b[0;32mIn[11], line 6\u001b[0m, in \u001b[0;36mget_audio_quality_measures\u001b[0;34m(s3_id)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m audio_quality:\n\u001b[1;32m      4\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m11\u001b[39m)]\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\n\u001b[0;32m----> 6\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(\u001b[43maudio_quality\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mpref\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m),\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshimmer_score\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mloudness_factor\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m      9\u001b[0m     audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mspectral_character\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mspectral_centroid\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     11\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbass_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmid_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhigh_ratio\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstereo_width\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     15\u001b[0m     \u001b[38;5;28mint\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtotal_clips\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28mfloat\u001b[39m(audio_quality[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclips_per_second\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m     17\u001b[0m ]\n",
      "\u001b[0;31mKeyError\u001b[0m: 'pref'"
     ]
    }
   ],
   "source": [
    "def get_audio_quality_measures(s3_id):\n",
    "    audio_quality = unpacked_pair_quality.get(s3_id, [])\n",
    "    if not audio_quality:\n",
    "        return [None for _ in range(11)]\n",
    "    return [\n",
    "        float(audio_quality[\"pref\"]),\n",
    "        float(audio_quality[\"shimmer_score\"]),\n",
    "        float(audio_quality[\"loudness_factor\"]),\n",
    "        audio_quality[\"spectral_character\"],\n",
    "        float(audio_quality[\"spectral_centroid\"]),\n",
    "        float(audio_quality[\"bass_ratio\"]),\n",
    "        float(audio_quality[\"mid_ratio\"]),\n",
    "        float(audio_quality[\"high_ratio\"]),\n",
    "        float(audio_quality[\"stereo_width\"]),\n",
    "        int(audio_quality[\"total_clips\"]),\n",
    "        float(audio_quality[\"clips_per_second\"]),\n",
    "    ]\n",
    "\n",
    "\n",
    "df[\n",
    "    [\n",
    "        \"pair_quality\",\n",
    "        \"total_shimmer_score\",\n",
    "        \"loudness_factor\",\n",
    "        \"spectral_character\",\n",
    "        \"spectral_centroid\",\n",
    "        \"bass_ratio\",\n",
    "        \"mid_ratio\",\n",
    "        \"high_ratio\",\n",
    "        \"stereo_width\",\n",
    "        \"total_clips\",\n",
    "        \"clips_per_second\",\n",
    "    ]\n",
    "] = pd.DataFrame(df[\"s3_id\"].apply(get_audio_quality_measures).tolist(), index=df.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.776956Z",
     "iopub.status.idle": "2025-03-18T03:49:22.777133Z",
     "shell.execute_reply": "2025-03-18T03:49:22.777067Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.777060Z"
    }
   },
   "outputs": [],
   "source": [
    "print(df.shape)\n",
    "df = df.dropna(subset=[\n",
    "        \"pair_quality\",\n",
    "        \"total_shimmer_score\",\n",
    "        \"loudness_factor\",\n",
    "        \"spectral_character\",\n",
    "        \"spectral_centroid\",\n",
    "        \"bass_ratio\",\n",
    "        \"mid_ratio\",\n",
    "        \"high_ratio\",\n",
    "        \"stereo_width\",\n",
    "        \"total_clips\",\n",
    "        \"clips_per_second\",\n",
    "    ])\n",
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:36.043975Z",
     "start_time": "2024-05-16T13:58:56.910958Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.777593Z",
     "iopub.status.idle": "2025-03-18T03:49:22.777720Z",
     "shell.execute_reply": "2025-03-18T03:49:22.777662Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.777657Z"
    }
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:40.799375Z",
     "start_time": "2024-05-16T13:59:36.394236Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.778019Z",
     "iopub.status.idle": "2025-03-18T03:49:22.778236Z",
     "shell.execute_reply": "2025-03-18T03:49:22.778177Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.778171Z"
    }
   },
   "outputs": [],
   "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[\"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": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.778492Z",
     "iopub.status.idle": "2025-03-18T03:49:22.778662Z",
     "shell.execute_reply": "2025-03-18T03:49:22.778607Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.778602Z"
    }
   },
   "outputs": [],
   "source": [
    "df = df.sort_values(by=[\"request_id\", \"preference\", \"diff_preference\"])\n",
    "df[\"pos_diff_preference\"] = df[\"diff_preference\"].diff()\n",
    "df[df[\"preference\"]][\"pos_diff_preference\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.779029Z",
     "iopub.status.idle": "2025-03-18T03:49:22.779237Z",
     "shell.execute_reply": "2025-03-18T03:49:22.779180Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.779175Z"
    }
   },
   "outputs": [],
   "source": [
    "lookup_percentiles = [5, 10, 20, 50, 80, 90, 95]\n",
    "df[\"shimmer_score_diff\"] = df[\"total_shimmer_score\"].diff()\n",
    "percentiles = np.percentile(\n",
    "    df[df[\"preference\"]][\"shimmer_score_diff\"].dropna(), lookup_percentiles\n",
    ")\n",
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"shimmer_score_diff\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['shimmer_score_diff']):.2f}\",\n",
    "    bins=np.linspace(-1, 1, 100),\n",
    "    alpha=0.5,\n",
    ")\n",
    "textstr = \"\\n\".join(\n",
    "    [\n",
    "        f\"{lookup_percentiles[i]}th: {percentile:.2f}\"\n",
    "        for i, percentile in enumerate(percentiles)\n",
    "    ]\n",
    ")\n",
    "plt.gcf().text(\n",
    "    0.15,\n",
    "    0.98,\n",
    "    textstr,\n",
    "    fontsize=10,\n",
    "    verticalalignment=\"top\",\n",
    "    horizontalalignment=\"left\",\n",
    "    bbox=dict(facecolor=\"white\", alpha=0.5),\n",
    ")\n",
    "for percentile in percentiles:\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "plt.title(\n",
    "    f\"Shimmer score difference --> {lookup_percentiles[-1]}th, {percentiles[-1]:.2f}\"\n",
    ")\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.779496Z",
     "iopub.status.idle": "2025-03-18T03:49:22.779662Z",
     "shell.execute_reply": "2025-03-18T03:49:22.779607Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.779601Z"
    }
   },
   "outputs": [],
   "source": [
    "lookup_percentiles = [5, 10, 20, 50, 80, 90, 95]\n",
    "percentiles = np.percentile(\n",
    "    df[df[\"preference\"]][\"pair_quality\"].dropna(), lookup_percentiles\n",
    ")\n",
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"pair_quality\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['pair_quality']):.2f}\",\n",
    "    bins=np.linspace(0, 1, 100),\n",
    "    alpha=0.5,\n",
    ")\n",
    "textstr = \"\\n\".join(\n",
    "    [\n",
    "        f\"{lookup_percentiles[i]}th: {percentile:.2f}\"\n",
    "        for i, percentile in enumerate(percentiles)\n",
    "    ]\n",
    ")\n",
    "plt.gcf().text(\n",
    "    0.15,\n",
    "    0.98,\n",
    "    textstr,\n",
    "    fontsize=10,\n",
    "    verticalalignment=\"top\",\n",
    "    horizontalalignment=\"left\",\n",
    "    bbox=dict(facecolor=\"white\", alpha=0.5),\n",
    ")\n",
    "for percentile in percentiles:\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "plt.title(f\"Pair quality --> {lookup_percentiles[0]}th, {percentiles[0]:.2f}\")\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.779895Z",
     "iopub.status.idle": "2025-03-18T03:49:22.780009Z",
     "shell.execute_reply": "2025-03-18T03:49:22.779956Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.779951Z"
    }
   },
   "outputs": [],
   "source": [
    "# CS said -1 means stero to mono; 1 means mono to stereo\n",
    "# only cut off the left side\n",
    "df[\"stereo_width_diff\"] = df[\"stereo_width\"].diff()\n",
    "lookup_percentiles = [5, 10, 20, 50, 80, 90, 95]\n",
    "percentiles = np.percentile(\n",
    "    df[df[\"preference\"]][\"stereo_width_diff\"].dropna(), lookup_percentiles\n",
    ")\n",
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"stereo_width_diff\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['stereo_width_diff']):.2f}\",\n",
    "    bins=np.linspace(-1, 1, 400),\n",
    "    alpha=0.5,\n",
    ")\n",
    "textstr = \"\\n\".join(\n",
    "    [\n",
    "        f\"{lookup_percentiles[i]}th: {percentile:.2f}\"\n",
    "        for i, percentile in enumerate(percentiles)\n",
    "    ]\n",
    ")\n",
    "plt.gcf().text(\n",
    "    0.15,\n",
    "    0.98,\n",
    "    textstr,\n",
    "    fontsize=10,\n",
    "    verticalalignment=\"top\",\n",
    "    horizontalalignment=\"left\",\n",
    "    bbox=dict(facecolor=\"white\", alpha=0.5),\n",
    ")\n",
    "for percentile in percentiles:\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "plt.title(\n",
    "    f\"Stereo Width Difference --> {lookup_percentiles[0]}th, {percentiles[0]:.2f}\"\n",
    ")\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.780310Z",
     "iopub.status.idle": "2025-03-18T03:49:22.780554Z",
     "shell.execute_reply": "2025-03-18T03:49:22.780486Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.780479Z"
    }
   },
   "outputs": [],
   "source": [
    "# CS said no tails is good\n",
    "lookup_percentiles = [5, 10, 20, 50, 80, 90, 95]\n",
    "percentiles = np.percentile(\n",
    "    df[df[\"preference\"]][\"spectral_centroid\"].dropna(), lookup_percentiles\n",
    ")\n",
    "percentiles = np.percentile(\n",
    "    df[~df[\"preference\"]][\"spectral_centroid\"].dropna(), lookup_percentiles\n",
    ")\n",
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"spectral_centroid\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['spectral_centroid']):.2f}\",\n",
    "    bins=np.linspace(0, 8000, 400),\n",
    "    alpha=0.5,\n",
    ")\n",
    "plt.hist(\n",
    "    df[~df[\"preference\"]][\"spectral_centroid\"],\n",
    "    label=f\"neg, mean: {np.mean(df[~df['preference']]['spectral_centroid']):.2f}\",\n",
    "    bins=np.linspace(0, 8000, 400),\n",
    "    alpha=0.5,\n",
    ")\n",
    "textstr = \"\\n\".join(\n",
    "    [\n",
    "        f\"{lookup_percentiles[i]}th: {percentile:.2f}\"\n",
    "        for i, percentile in enumerate(percentiles)\n",
    "    ]\n",
    ")\n",
    "plt.gcf().text(\n",
    "    0.15,\n",
    "    0.98,\n",
    "    textstr,\n",
    "    fontsize=10,\n",
    "    verticalalignment=\"top\",\n",
    "    horizontalalignment=\"left\",\n",
    "    bbox=dict(facecolor=\"white\", alpha=0.5),\n",
    ")\n",
    "for percentile in percentiles:\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "plt.title(\"Spectral Centroid\")\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.780934Z",
     "iopub.status.idle": "2025-03-18T03:49:22.781063Z",
     "shell.execute_reply": "2025-03-18T03:49:22.781004Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.780998Z"
    }
   },
   "outputs": [],
   "source": [
    "# CS said no tails is good\n",
    "# take the relative centroid diff\n",
    "df[\"spectral_centroid_diff\"] = df[\"spectral_centroid\"].diff() / df[\"spectral_centroid\"]\n",
    "lookup_percentiles = [5, 10, 20, 50, 80, 90, 95]\n",
    "percentiles = np.percentile(\n",
    "    df[df[\"preference\"]][\"spectral_centroid_diff\"].dropna(), lookup_percentiles\n",
    ")\n",
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"spectral_centroid_diff\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['spectral_centroid_diff']):.2f}\",\n",
    "    bins=np.linspace(-1, 1, 400),\n",
    "    alpha=0.5,\n",
    ")\n",
    "textstr = \"\\n\".join(\n",
    "    [\n",
    "        f\"{lookup_percentiles[i]}th: {percentile:.2f}\"\n",
    "        for i, percentile in enumerate(percentiles)\n",
    "    ]\n",
    ")\n",
    "plt.gcf().text(\n",
    "    0.15,\n",
    "    0.98,\n",
    "    textstr,\n",
    "    fontsize=10,\n",
    "    verticalalignment=\"top\",\n",
    "    horizontalalignment=\"left\",\n",
    "    bbox=dict(facecolor=\"white\", alpha=0.5),\n",
    ")\n",
    "for percentile in percentiles:\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "plt.title(\n",
    "    f\"Spectral Centroid Difference ratio --> {lookup_percentiles[-1]}th, {percentiles[-1]:.2f}\"\n",
    ")\n",
    "# plt.yscale(\"log\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.781514Z",
     "iopub.status.idle": "2025-03-18T03:49:22.781715Z",
     "shell.execute_reply": "2025-03-18T03:49:22.781650Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.781643Z"
    }
   },
   "outputs": [],
   "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[\"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": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.782048Z",
     "iopub.status.idle": "2025-03-18T03:49:22.782292Z",
     "shell.execute_reply": "2025-03-18T03:49:22.782225Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.782218Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.hist(\n",
    "    df[df[\"preference\"]][\"total_shimmer_score\"],\n",
    "    label=f\"pos, mean: {np.mean(df[df['preference']]['total_shimmer_score']):.2f}\",\n",
    "    bins=np.linspace(0, 10, 100),\n",
    "    alpha=0.5,\n",
    ")\n",
    "plt.hist(\n",
    "    df[~df[\"preference\"]][\"total_shimmer_score\"],\n",
    "    label=f\"neg, mean: {np.mean(df[~df['preference']]['total_shimmer_score']):.2f}\",\n",
    "    bins=np.linspace(0, 10, 100),\n",
    "    alpha=0.5,\n",
    ")\n",
    "percentiles = np.percentile(df[df[\"preference\"]][\"total_shimmer_score\"].dropna(), [50, 75, 90])\n",
    "for percentile in percentiles:\n",
    "    # print(percentile)\n",
    "    plt.axvline(x=percentile, color=\"r\", linestyle=\"dashed\", linewidth=1)\n",
    "# plt.yscale(\"log\")\n",
    "plt.title(f\"Shimmer score\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.035167Z",
     "start_time": "2024-05-16T13:59:40.801098Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.782665Z",
     "iopub.status.idle": "2025-03-18T03:49:22.782797Z",
     "shell.execute_reply": "2025-03-18T03:49:22.782740Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.782734Z"
    }
   },
   "outputs": [],
   "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\"])  # 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\"] >= 30)  # can't be too short, otherwise it is obvious\n",
    "    # & (df[\"duration\"] <= 60)  # can't be badly long\n",
    "    & (df[\"has_continue_and_start_continue_at\"].isna())  # won't have any continues\n",
    "    & (df[\"norm_play_frac\"] <= 2.1)\n",
    "    # & (df[\"sum_total_play_duration_5\"] >= 31)\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\"])  # get basics aligned\n",
    "    & (\n",
    "        df[\"good_continue_at\"]\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\"] >= 30)  # can't be too short, otherwise it is obvious\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\"])\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\"])\n",
    "            & (df[\"reaction_play_count\"] >= normal_pos_play_count)\n",
    "            & (df[\"norm_play_frac\"] >= 2.1)  # this is a bit of a luxury cut...\n",
    "            & (df[\"sum_total_play_duration_5\"] >= 31)\n",
    "        )\n",
    "    )\n",
    "    # & (df[\"norm_play_frac\"] >= 1.9)\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[\"task\"] == \"\") | (df[\"task\"] == \"extend\"))\n",
    "    & (\n",
    "        (df[\"upvote_count\"] >= 1)\n",
    "        | (df[\"reaction_play_count\"] >= 5)\n",
    "        | (df[\"concat_play_counts\"] >= 5)\n",
    "    )\n",
    "    # & (df[\"pos_diff_preference\"] == 2)\n",
    "    # & ((0 < df[\"similarity\"]) &  (df[\"similarity\"] <= 0.99))\n",
    "    # & (\n",
    "    #     (df[\"cer_diff_preference\"] < 0.25) & (df[\"cer\"] < 0.8)\n",
    "    # )  # cut on hoot cer difference and abs cer\n",
    "    # & (df[\"pair_quality\"] > 0.31)  # bottom 5%\n",
    "    # & ((df[\"total_shimmer_score\"] < 1) | (df[\"shimmer_score_diff\"] < 0.4))\n",
    "    # & (df[\"stereo_width_diff\"] > -0.2)  # cut off bottom 5%\n",
    "    # & (df[\"spectral_centroid_diff\"] < 0.25)  # crop off the top 5%\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": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.250737Z",
     "start_time": "2024-05-16T13:59:41.036434Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.783162Z",
     "iopub.status.idle": "2025-03-18T03:49:22.783282Z",
     "shell.execute_reply": "2025-03-18T03:49:22.783227Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.783222Z"
    }
   },
   "outputs": [],
   "source": [
    "df_slice = df[df[\"request_id\"].isin(set(unique_requests))].copy()\n",
    "print(\n",
    "    f\"{os.path.basename(OUT_DATA_DIR)} requests\",\n",
    "    df_slice[\"request_id\"].nunique(),\n",
    "    \"clips\",\n",
    "    df_slice.shape[0],\n",
    "    f\"total khrs {sum(df_slice['duration'] / 3600 / 1000):.3f};\",\n",
    "    f\"N gpus for 1000 iters {df_slice.shape[0] / 8 / 2 / 1000:.3f};\",\n",
    "    f\"4 gpus for x iters {df_slice.shape[0] / 8 / 2 / 4:.3f};\",\n",
    "    f\"n unique users {df_slice['user_id'].nunique()}\",\n",
    "    f\"n pro users {df_slice[df_slice['is_pro_user']]['user_id'].nunique()}\",\n",
    ")\n",
    "# up t7 requests 37735 clips 75470 total khrs 3.964; N gpus for 1000 iters 4.717; 4 gpus for x iters 1179.219; n unique users 19093 n pro users 16217\n",
    "# up t17 requests 49262 clips 98524 total khrs 5.187; N gpus for 1000 iters 6.158; 4 gpus for x iters 1539.438; n unique users 24093 n pro users 20326\n",
    "# up v2 t1 requests 6008 clips 12016 total khrs 0.642; N gpus for 1000 iters 0.751; 4 gpus for x iters 187.750; n unique users 3846 n pro users 3645\n",
    "# up v2 t2 requests 10772 clips 21544 total khrs 1.154; N gpus for 1000 iters 1.347; 4 gpus for x iters 336.625; n unique users 6527 n pro users 6027\n",
    "# up v3 t10 requests 27102 clips 54204 total khrs 2.938; N gpus for 1000 iters 3.388; 4 gpus for x iters 846.938; n unique users 13932 n pro users 12187\n",
    "# up v4 t1  requests 3201 clips 6402 total khrs 0.343; N gpus for 1000 iters 0.400; 4 gpus for x iters 100.031; n unique users 2363 n pro users 2321\n",
    "# up v4 t2  requests 12818 clips 25636 total khrs 1.371; N gpus for 1000 iters 1.602; 4 gpus for x iters 400.562; n unique users 7753 n pro users 7524\n",
    "# up v4 t3  requests 15332 clips 30664 total khrs 1.637; N gpus for 1000 iters 1.917; 4 gpus for x iters 479.125; n unique users 8960 n pro users 8656\n",
    "# up v4 t4  requests 18368 clips 36736 total khrs 1.962; N gpus for 1000 iters 2.296; 4 gpus for x iters 574.000; n unique users 10427 n pro users 10049\n",
    "# up v4 t5  requests 22037 clips 44074 total khrs 2.347; N gpus for 1000 iters 2.755; 4 gpus for x iters 688.656; n unique users 12121 n pro users 11551\n",
    "# up v4 t6  requests 27374 clips 54748 total khrs 2.915; N gpus for 1000 iters 3.422; 4 gpus for x iters 855.438; n unique users 14463 n pro users 13678\n",
    "# up v4 t7  requests 31797 clips 63594 total khrs 3.384; N gpus for 1000 iters 3.975; 4 gpus for x iters 993.656; n unique users 16398 n pro users 15325\n",
    "# up v5 t1   requests 17035 clips 34070 total khrs 1.788; N gpus for 1000 iters 2.129; 4 gpus for x iters 532.344; n unique users 10858 n pro users 8065\n",
    "# up v5 t2  requests 20030 clips 40060 total khrs 2.108; N gpus for 1000 iters 2.504; 4 gpus for x iters 625.938; n unique users 12395 n pro users 9156"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.277006Z",
     "start_time": "2024-05-16T13:59:41.252105Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.783563Z",
     "iopub.status.idle": "2025-03-18T03:49:22.783777Z",
     "shell.execute_reply": "2025-03-18T03:49:22.783715Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.783709Z"
    }
   },
   "outputs": [],
   "source": [
    "test_mask = (df_slice[\"preference\"]) & (\n",
    "    (df_slice[\"is_in_playlist\"]) | (df_slice[\"concat_in_playlist\"])\n",
    ")\n",
    "print(\"positive in playlist\", df_slice[test_mask].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.392244Z",
     "start_time": "2024-05-16T13:59:41.324472Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.784024Z",
     "iopub.status.idle": "2025-03-18T03:49:22.784188Z",
     "shell.execute_reply": "2025-03-18T03:49:22.784132Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.784126Z"
    }
   },
   "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": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.784498Z",
     "iopub.status.idle": "2025-03-18T03:49:22.784620Z",
     "shell.execute_reply": "2025-03-18T03:49:22.784564Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.784558Z"
    }
   },
   "outputs": [],
   "source": [
    "# def modify_model_name(model_name, metadata):\n",
    "#     if (\n",
    "#         model_name.startswith(\"chirp-v3p5-engine-t\")\n",
    "#         or model_name.startswith(\"chirp-v3p5-engine-s\")\n",
    "#         or model_name.startswith(\"chirp-v4\")\n",
    "#         or model_name.startswith(\"chirp-v3p5-h-s-31\")\n",
    "#     ):\n",
    "#         if \"param_experiment\" in metadata:\n",
    "#             exp = metadata.get(\"param_experiment\", \"\")\n",
    "#             if exp:\n",
    "#                 return f\"{model_name}_{exp}\"\n",
    "#     return model_name\n",
    "\n",
    "# metrics_check_df_slice = df_slice.copy()\n",
    "# metrics_check_df_slice[\"model_name\"] = metrics_check_df_slice.apply(\n",
    "#     lambda row: modify_model_name(row[\"model_name\"], row[\"metadata\"]), axis=1\n",
    "# )\n",
    "# get_preference_counts(metrics_check_df_slice)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.784911Z",
     "iopub.status.idle": "2025-03-18T03:49:22.785029Z",
     "shell.execute_reply": "2025-03-18T03:49:22.784975Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.784970Z"
    }
   },
   "outputs": [],
   "source": [
    "print(df_slice[\"source\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T14:00:20.866354Z",
     "start_time": "2024-05-16T14:00:12.443344Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.785410Z",
     "iopub.status.idle": "2025-03-18T03:49:22.785643Z",
     "shell.execute_reply": "2025-03-18T03:49:22.785581Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.785574Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice.to_csv(\"/home/tony/Data/Preference/13b_v0/interesting_clips_v3p5_s_8_20240828_slice.csv\")\n",
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Need to kick out the ones has gpt prompt -- these are pairs with different text inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.932966Z",
     "start_time": "2024-05-16T13:59:41.932957Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.785919Z",
     "iopub.status.idle": "2025-03-18T03:49:22.786140Z",
     "shell.execute_reply": "2025-03-18T03:49:22.786075Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.786068Z"
    }
   },
   "outputs": [],
   "source": [
    "final_filtered_requests = df_slice[\"request_id\"].unique()\n",
    "print(len(final_filtered_requests))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934277Z",
     "start_time": "2024-05-16T13:59:41.934268Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.786408Z",
     "iopub.status.idle": "2025-03-18T03:49:22.786605Z",
     "shell.execute_reply": "2025-03-18T03:49:22.786547Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.786540Z"
    }
   },
   "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": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.786847Z",
     "iopub.status.idle": "2025-03-18T03:49:22.786964Z",
     "shell.execute_reply": "2025-03-18T03:49:22.786911Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.786905Z"
    }
   },
   "outputs": [],
   "source": [
    "# BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actually make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.934954Z",
     "start_time": "2024-05-16T13:59:41.934946Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.787366Z",
     "iopub.status.idle": "2025-03-18T03:49:22.787599Z",
     "shell.execute_reply": "2025-03-18T03:49:22.787530Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.787523Z"
    }
   },
   "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.status.busy": "2025-03-18T03:49:22.787868Z",
     "iopub.status.idle": "2025-03-18T03:49:22.788039Z",
     "shell.execute_reply": "2025-03-18T03:49:22.787980Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.787973Z"
    }
   },
   "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 / 2 / 1000} nodes, {train_df.shape[0] / 8 / 2 / 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.status.busy": "2025-03-18T03:49:22.788302Z",
     "iopub.status.idle": "2025-03-18T03:49:22.788426Z",
     "shell.execute_reply": "2025-03-18T03:49:22.788368Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.788363Z"
    }
   },
   "outputs": [],
   "source": [
    "make_dataset(\n",
    "    val_df,\n",
    "    OUT_DATA_DIR,\n",
    "    is_val=True,\n",
    "    npz_dir=NPZ_DIR,\n",
    "    do_extend_chunks=True,\n",
    "    clip_id_to_quality_scores=unpacked_pair_quality,\n",
    ")"
   ]
  },
  {
   "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.status.busy": "2025-03-18T03:49:22.788805Z",
     "iopub.status.idle": "2025-03-18T03:49:22.788925Z",
     "shell.execute_reply": "2025-03-18T03:49:22.788870Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.788865Z"
    }
   },
   "outputs": [],
   "source": [
    "make_dataset(\n",
    "    train_df,\n",
    "    OUT_DATA_DIR,\n",
    "    is_val=False,\n",
    "    npz_dir=NPZ_DIR,\n",
    "    do_extend_chunks=True,\n",
    "    clip_id_to_quality_scores=unpacked_pair_quality,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-29T19:46:47.549860Z",
     "start_time": "2024-01-29T19:46:47.548015Z"
    }
   },
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.937879Z",
     "start_time": "2024-05-16T13:59:41.937870Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.789258Z",
     "iopub.status.idle": "2025-03-18T03:49:22.789439Z",
     "shell.execute_reply": "2025-03-18T03:49:22.789380Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.789374Z"
    }
   },
   "outputs": [],
   "source": [
    "# verify\n",
    "metas_val = read_jsonl(os.path.join(OUT_DATA_DIR, \"metas_val.jsonl\"))\n",
    "print(len(metas_val) / 2)\n",
    "mm_semantic_val = np.memmap(\n",
    "    os.path.join(OUT_DATA_DIR, \"data_semantic_val.bin\"), dtype=np.uint16, mode=\"r\"\n",
    ")\n",
    "mm_vae_val = np.memmap(\n",
    "    os.path.join(OUT_DATA_DIR, \"data_vae_val.bin\"), dtype=np.float16, mode=\"r\"\n",
    ")\n",
    "\n",
    "\n",
    "mm_vae_val = mm_vae_val.reshape(-1, VAE_MEMMAP_SIZE, VAE_DIM)\n",
    "print(mm_vae_val.shape)\n",
    "\n",
    "mm_semantic_val = mm_semantic_val.reshape(-1, SEMANTIC_MEMMAP_SIZE)\n",
    "print(mm_semantic_val.shape)\n",
    "\n",
    "assert len(metas_val) == mm_vae_val.shape[0] == mm_semantic_val.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.789715Z",
     "iopub.status.idle": "2025-03-18T03:49:22.789915Z",
     "shell.execute_reply": "2025-03-18T03:49:22.789858Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.789852Z"
    }
   },
   "outputs": [],
   "source": [
    "# # load codec for decoding\n",
    "# from suno_utils.tasks.dac_vae_100hz_peaq import (  # NOTE: works for 25hz as well\n",
    "#     preload_models as preload_codec_models,\n",
    "#     decode as codec_decode,\n",
    "#     encode as codec_encode,\n",
    "#     get_embedding_rate,\n",
    "#     load_model as load_codec_model,\n",
    "# )\n",
    "\n",
    "# CODEC_FILEPATH = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "# preload_codec_models(CODEC_FILEPATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.790171Z",
     "iopub.status.idle": "2025-03-18T03:49:22.790335Z",
     "shell.execute_reply": "2025-03-18T03:49:22.790278Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.790273Z"
    }
   },
   "outputs": [],
   "source": [
    "# # decode some audio\n",
    "idx = 108\n",
    "# # ensure even index\n",
    "assert idx % 2 == 0\n",
    "# print(metas_val[idx])\n",
    "# print(\"negative\")\n",
    "# audio = codec_decode(mm_vae_val[idx])\n",
    "# audio.normalize_volume().play()\n",
    "\n",
    "# print(metas_val[idx + 1])\n",
    "# print(\"positive\")\n",
    "# audio = codec_decode(mm_vae_val[idx + 1])\n",
    "# audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.790554Z",
     "iopub.status.idle": "2025-03-18T03:49:22.790667Z",
     "shell.execute_reply": "2025-03-18T03:49:22.790615Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.790610Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.791021Z",
     "iopub.status.idle": "2025-03-18T03:49:22.791133Z",
     "shell.execute_reply": "2025-03-18T03:49:22.791080Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.791075Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.equal(torch.tensor(mm_semantic_val[idx]), torch.tensor(mm_semantic_val[idx + 1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941167Z",
     "start_time": "2024-05-16T13:59:41.941159Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.791471Z",
     "iopub.status.idle": "2025-03-18T03:49:22.791585Z",
     "shell.execute_reply": "2025-03-18T03:49:22.791532Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.791527Z"
    }
   },
   "outputs": [],
   "source": [
    "# original_npz_path = f\"/app/suno/data/dpo/7b_npz/{test_metas[idx]['id']}.npz\"\n",
    "# original_npz_path = \"/app/suno/data/dpo/7b_npz/729c3011-f672-4ccd-8d82-1cbf2b52ff69.npz\"\n",
    "# original_arr = np.load(original_npz_path)[\"v2_raw\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.941801Z",
     "start_time": "2024-05-16T13:59:41.941793Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.791883Z",
     "iopub.status.idle": "2025-03-18T03:49:22.791996Z",
     "shell.execute_reply": "2025-03-18T03:49:22.791943Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.791938Z"
    }
   },
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.792277Z",
     "iopub.status.idle": "2025-03-18T03:49:22.792388Z",
     "shell.execute_reply": "2025-03-18T03:49:22.792337Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.792332Z"
    }
   },
   "outputs": [],
   "source": [
    "metas_tr = read_jsonl(os.path.join(OUT_DATA_DIR, \"metas_tr.jsonl\"))\n",
    "validation_on_metas(metas_tr)\n",
    "print(len(metas_tr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.792713Z",
     "iopub.status.idle": "2025-03-18T03:49:22.792825Z",
     "shell.execute_reply": "2025-03-18T03:49:22.792773Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.792768Z"
    }
   },
   "outputs": [],
   "source": [
    "sum(len(meta[\"tags\"][0]) == 0 for meta in metas_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-16T13:59:41.945972Z",
     "start_time": "2024-05-16T13:59:41.945964Z"
    },
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.793204Z",
     "iopub.status.idle": "2025-03-18T03:49:22.793363Z",
     "shell.execute_reply": "2025-03-18T03:49:22.793306Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.793301Z"
    }
   },
   "outputs": [],
   "source": [
    "!cd /home/tony/Work/tony/slurm/diffusion && sbatch run_diffusion.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.793578Z",
     "iopub.status.idle": "2025-03-18T03:49:22.793687Z",
     "shell.execute_reply": "2025-03-18T03:49:22.793636Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.793631Z"
    }
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "# Basic file copy\n",
    "shutil.copy(\n",
    "    \"/home/tony/Work/tony/Preference/make_dataset_diff_upsample_v1_r4_comb.ipynb\",\n",
    "    os.path.join(OUT_DATA_DIR, \"make_dataset.ipynb\"),\n",
    ")\n",
    "print(\"Cache kept!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspections "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.794020Z",
     "iopub.status.idle": "2025-03-18T03:49:22.794133Z",
     "shell.execute_reply": "2025-03-18T03:49:22.794080Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.794075Z"
    }
   },
   "outputs": [],
   "source": [
    "# df[df[\"preference\"] & (df[\"shimmer_score_diff\"] > 3)][\n",
    "#     [\n",
    "#         \"index\",\n",
    "#         \"s3_id\",\n",
    "#         \"total_shimmer_score\",\n",
    "#         \"shimmer_score_diff\",\n",
    "#         \"request_id\",\n",
    "#         \"preference\",\n",
    "#     ]\n",
    "# ].tail()\n",
    "\n",
    "# df[df[\"preference\"] & (df[\"pair_quality\"] < 0.1)][\n",
    "#     [\n",
    "#         \"index\",\n",
    "#         \"s3_id\",\n",
    "#         \"total_shimmer_score\",\n",
    "#         \"shimmer_score_diff\",\n",
    "#         \"pair_quality\",\n",
    "#         \"request_id\",\n",
    "#         \"preference\",\n",
    "#     ]\n",
    "# ].tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.794413Z",
     "iopub.status.idle": "2025-03-18T03:49:22.794524Z",
     "shell.execute_reply": "2025-03-18T03:49:22.794472Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.794467Z"
    }
   },
   "outputs": [],
   "source": [
    "# test_pair_df = df[df[\"request_id\"] == \"621c8b02-a905-48f1-a2d5-a4e8423d1505\"]\n",
    "# print(\n",
    "#     test_pair_df[\n",
    "#         [\n",
    "#             \"s3_id\",\n",
    "#             \"total_shimmer_score\",\n",
    "#             \"pair_quality\",\n",
    "#             \"request_id\",\n",
    "#             \"preference\",\n",
    "#             \"prompt_text\",\n",
    "#         ]\n",
    "#     ]\n",
    "# )\n",
    "# negative_audio = Audio.from_s3(\n",
    "#     f\"s3://suno-data-uploads/studio/uploads/{test_pair_df['s3_id'].values[0]}.mp3\",\n",
    "#     n_channels=2,\n",
    "# )\n",
    "# print(\"negative\")\n",
    "# negative_audio.get_segment(0, 30).play()\n",
    "# positive_audio = Audio.from_s3(\n",
    "#     f\"s3://suno-data-uploads/studio/uploads/{test_pair_df['s3_id'].values[1]}.mp3\",\n",
    "#     n_channels=2,\n",
    "# )\n",
    "# print(\"positive\")\n",
    "# positive_audio.get_segment(0, 30).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.794782Z",
     "iopub.status.idle": "2025-03-18T03:49:22.794891Z",
     "shell.execute_reply": "2025-03-18T03:49:22.794839Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.794834Z"
    }
   },
   "outputs": [],
   "source": [
    "# total_dict = {}\n",
    "# total_dict.update(pair_quality_dict)\n",
    "# total_dict.update(pair_quality_1_dict)\n",
    "# total_dict.update(pair_quality_2_dict)\n",
    "# total_dict.update(pair_quality_3_dict)\n",
    "# len(total_dict)\n",
    "# with open(\n",
    "#     os.path.join(\"/home/tony/Data/Preference/up_v1\", \"pair_quality.json\"), \"w\"\n",
    "# ) as fp:\n",
    "#     json.dump(total_dict, fp, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.795321Z",
     "iopub.status.idle": "2025-03-18T03:49:22.795431Z",
     "shell.execute_reply": "2025-03-18T03:49:22.795380Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.795375Z"
    }
   },
   "outputs": [],
   "source": [
    "# import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.795684Z",
     "iopub.status.idle": "2025-03-18T03:49:22.795806Z",
     "shell.execute_reply": "2025-03-18T03:49:22.795753Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.795738Z"
    }
   },
   "outputs": [],
   "source": [
    "# test_arr = np.load(\"/home/tony/Data/test_npz/diffusion_input_tensor([ 18, 182]).npy\")\n",
    "# test_arr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.796016Z",
     "iopub.status.idle": "2025-03-18T03:49:22.796126Z",
     "shell.execute_reply": "2025-03-18T03:49:22.796075Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.796070Z"
    }
   },
   "outputs": [],
   "source": [
    "# mm_vae_val = np.memmap(\n",
    "#     os.path.join(OUT_DATA_DIR, \"data_vae_val.bin\"), dtype=np.float16, mode=\"r\"\n",
    "# )\n",
    "\n",
    "# mm_vae_val = mm_vae_val.reshape(-1, VAE_MEMMAP_SIZE, VAE_DIM)\n",
    "# print(mm_vae_val.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.796423Z",
     "iopub.status.idle": "2025-03-18T03:49:22.796533Z",
     "shell.execute_reply": "2025-03-18T03:49:22.796482Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.796477Z"
    }
   },
   "outputs": [],
   "source": [
    "# print(\"negative\")\n",
    "# audio = codec_decode(test_arr[0].T / 2.5)\n",
    "# audio.normalize_volume().play()\n",
    "\n",
    "# print(\"positive\")\n",
    "# audio = codec_decode(test_arr[1].T / 2.5)\n",
    "# audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.796743Z",
     "iopub.status.idle": "2025-03-18T03:49:22.796852Z",
     "shell.execute_reply": "2025-03-18T03:49:22.796801Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.796796Z"
    }
   },
   "outputs": [],
   "source": [
    "# print(\"negative\")\n",
    "# audio = codec_decode(mm_vae_val[18])\n",
    "# audio.normalize_volume().play()\n",
    "\n",
    "# print(\"positive\")\n",
    "# audio = codec_decode(mm_vae_val[19])\n",
    "# audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.797138Z",
     "iopub.status.idle": "2025-03-18T03:49:22.797249Z",
     "shell.execute_reply": "2025-03-18T03:49:22.797197Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.797193Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "rng = torch.quasirandom.SobolEngine(1, scramble=True, seed=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.797503Z",
     "iopub.status.idle": "2025-03-18T03:49:22.797612Z",
     "shell.execute_reply": "2025-03-18T03:49:22.797560Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.797555Z"
    }
   },
   "outputs": [],
   "source": [
    "t = rng.draw(4)[:, 0].to(torch.bfloat16)\n",
    "print(t)\n",
    "# Replace 1% of t with ones to ensure training on terminal SNR\n",
    "t = torch.where(torch.rand_like(t) < 0.5, torch.ones_like(t), t)\n",
    "print(t)\n",
    "t = torch.repeat_interleave(t, repeats=2, dim=0)\n",
    "print(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.797909Z",
     "iopub.status.idle": "2025-03-18T03:49:22.798015Z",
     "shell.execute_reply": "2025-03-18T03:49:22.797965Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.797961Z"
    }
   },
   "outputs": [],
   "source": [
    "(t * 32).to(int) / 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.798309Z",
     "iopub.status.idle": "2025-03-18T03:49:22.798419Z",
     "shell.execute_reply": "2025-03-18T03:49:22.798366Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.798362Z"
    }
   },
   "outputs": [],
   "source": [
    "t[0] = 0.99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.798695Z",
     "iopub.status.idle": "2025-03-18T03:49:22.798804Z",
     "shell.execute_reply": "2025-03-18T03:49:22.798752Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.798747Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.round(t * 32) / 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.799117Z",
     "iopub.status.idle": "2025-03-18T03:49:22.799229Z",
     "shell.execute_reply": "2025-03-18T03:49:22.799178Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.799173Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "with open(f\"/home/tony/Data/Preference/up_v4/full_pair_quality.json\", \"r\") as f:\n",
    "   result = json.load(f)\n",
    "# result = {}\n",
    "print(len(result))\n",
    "for job_idx in range(4):\n",
    "    with open(f\"/home/tony/Data/Preference/up_v4/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\"/home/tony/Data/Preference/up_v4/full_pair_quality.json\", \"w\") as f:\n",
    "#     json.dump(result, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.799516Z",
     "iopub.status.idle": "2025-03-18T03:49:22.799626Z",
     "shell.execute_reply": "2025-03-18T03:49:22.799575Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.799570Z"
    }
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# semantic_codes_chunk = torch.ones((1, 100))\n",
    "# semantic_skip_phase = 0\n",
    "# semantic_skip_factor = 4\n",
    "# mask = torch.ones_like(semantic_codes_chunk, dtype=torch.bool)\n",
    "# indices = (\n",
    "#     torch.arange(semantic_codes_chunk.size(1)) + semantic_skip_phase\n",
    "# ) % semantic_skip_factor == 0\n",
    "# mask[:, indices] = False\n",
    "# mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.799904Z",
     "iopub.status.idle": "2025-03-18T03:49:22.800011Z",
     "shell.execute_reply": "2025-03-18T03:49:22.799961Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.799956Z"
    }
   },
   "outputs": [],
   "source": [
    "df_slice.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.800587Z",
     "iopub.status.idle": "2025-03-18T03:49:22.800743Z",
     "shell.execute_reply": "2025-03-18T03:49:22.800688Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.800682Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[(df_slice[\"preference\"]) & (df_slice[\"total_shimmer_score\"] < 0.5)][[\"s3_id\", \"total_shimmer_score\"]].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-03-18T03:49:22.801010Z",
     "iopub.status.idle": "2025-03-18T03:49:22.801169Z",
     "shell.execute_reply": "2025-03-18T03:49:22.801110Z",
     "shell.execute_reply.started": "2025-03-18T03:49:22.801104Z"
    }
   },
   "outputs": [],
   "source": [
    "# df_slice[(df_slice[\"preference\"]) & (df_slice[\"total_shimmer_score\"] > 2)][[\"s3_id\", \"total_shimmer_score\"]].head()"
   ]
  },
  {
   "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
}
