{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "from collections import defaultdict\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from preference_data_preparation_4min_30b_task import *\n",
    "from preference_helper import *\n",
    "from sklearn.model_selection import train_test_split\n",
    "from suno_utils.utils.s3 import download_s3_files\n",
    "from suno_utils.utils.text import read_json, read_jsonl, write_json, write_jsonl\n",
    "from tqdm import tqdm\n",
    "\n",
    "pd.set_option(\"display.max_rows\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)\n",
    "\n",
    "# setup autoload\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "\n",
    "def custom_parse(x):\n",
    "    try:\n",
    "        return json.loads(x)\n",
    "    except:\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUT_DATA_DIR = \"/app/suno/data/dpo/30b_t6_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/30b_npz\"\n",
    "# NOTE FOR 30b we increase this from 6016 up\n",
    "N_TOKENS_AUDIO = 6016"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_pickle(\n",
    "    # \"/home/tony/Data/Preference/30b_v3/interesting_clips_v4_t_4_20240923_full_with_sem_distance_and_similarity.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_20240919_full_l10_with_cer.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v3/interesting_clips_v4_t_4_20240925_full_l10_with_cer.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v5/interesting_clips_v4_t_5_20241029_full_with_cer.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241117_full_with_cer.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241118_full.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250222_full.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250301_full.pkl\"\n",
    "    \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250227_full_long.pkl\"\n",
    "    # \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250107_full.pkl\"\n",
    ")  # , engine='python')\n",
    "# df = pd.read_csv(\n",
    "#     \"/home/tony/Data/Preference/30b_v0/interesting_clips_v4_t_1_20240808.csv\"\n",
    "# )  # , engine='python')\n",
    "print(\"Preference data shape\", df.shape)\n",
    "\n",
    "# df_cover = pd.read_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241109_full_with_sem_distance_and_similarity.pkl\"\n",
    "# )\n",
    "# print(df_cover.shape)\n",
    "\n",
    "# df_artist = pd.read_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241109_full_with_sem_distance_and_similarity_artist.pkl\"\n",
    "# )\n",
    "# print(df_cover.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_cover[df_cover[\"preference\"]][\"similarity\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# df_cover[\"similarity_diff\"] = df_cover[\"similarity\"].diff()\n",
    "# df_cover[df_cover[\"preference\"]][\"similarity_diff\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# df_cover[\"continued_parent\"] = None\n",
    "# df_cover[\"continue_at\"] = -1\n",
    "# df_cover[df_cover[\"preference\"]][\"similarity\"].describe()\n",
    "# df_cover_drops_id = df_cover[\n",
    "#     (df_cover[\"preference\"])\n",
    "#     & (~((0.1 < df_cover[\"similarity\"]) & (df_cover[\"similarity\"] <= 0.99)))\n",
    "# ][\"s3_id\"].unique()\n",
    "# print(len(df_cover_drops_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_artist[df_artist[\"preference\"]][\"similarity\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# df_artist[\"similarity_diff\"] = df_artist[\"similarity\"].diff()\n",
    "# df_artist[df_artist[\"preference\"]][\"similarity_diff\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# df_artist[\"continued_parent\"] = None\n",
    "# df_artist[\"continue_at\"] = -1\n",
    "# df_artist[df_artist[\"preference\"]][\"similarity\"].describe()\n",
    "# df_artist_drops_id = df_artist[\n",
    "#     (df_artist[\"preference\"])\n",
    "#     & ((~((0 < df_artist[\"similarity\"]) & (df_artist[\"similarity\"] <= 0.99)))\n",
    "#     | (~(df_artist[\"similarity_diff\"] > -0.2)))\n",
    "# ][\"s3_id\"].unique()\n",
    "# print(len(df_artist_drops_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"before filter on cover similarity\", df.shape)\n",
    "# df = df[~df[\"s3_id\"].isin(df_cover_drops_id)].copy()\n",
    "# print(\"after filter on cover similarity\", df.shape)\n",
    "# df = df[~df[\"s3_id\"].isin(df_artist_drops_id)].copy()\n",
    "# print(\"after filter on artist similarity\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.dropna(axis=1, how=\"all\")\n",
    "print(\"after dropna\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_paths = os.listdir(NPZ_DIR)\n",
    "print(len(converted_paths))\n",
    "\n",
    "converted_paths = set([f.replace(\".npz\", \"\") for f in converted_paths])\n",
    "print(len(converted_paths))\n",
    "\n",
    "print(\"pre-downloaded df\", df.shape)\n",
    "df[df[\"s3_id\"].isin(converted_paths)].shape\n",
    "df = df[df[\"s3_id\"].isin(converted_paths)].copy()\n",
    "print(\"downloaded df\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"is_30b\"] = df[\"model_name\"].str.contains(\"-t\")\n",
    "print(df[\"is_30b\"].value_counts())\n",
    "df[\"task\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LET's do the data prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# drop extend for now\n",
    "# reason is -- cause they are likely caused by extend from 13b\n",
    "# we don't want contamination\n",
    "# df = df[df[\"task\"] != \"extend\"].copy()\n",
    "print(df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## for 13b this is easy for now\n",
    "print(df.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "print(\"before filter on model name\", df.shape)\n",
    "df = df[df[\"model_name\"].isin([\"chirp-v4-h-t-6\", \"chirp-v3p5-engine-t-6\"])]\n",
    "# df = df[df[\"model_name\"].isin([\"chirp-v3p5-engine-t-6\"])]\n",
    "print(\"after filter on model name\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df[\"is_public\"].value_counts())\n",
    "# remove public for now cause fucking users\n",
    "# df = df[~df[\"is_public\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"before filter on request id pairs\", df.shape)\n",
    "df = df[\n",
    "    df[\"request_id\"].isin(\n",
    "        df[\"request_id\"].value_counts().index[df[\"request_id\"].value_counts() == 2]\n",
    "    )\n",
    "]\n",
    "print(\"after filter on request id pairs\", df.shape)\n",
    "print(df.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "assert df.shape[0] == df[\"request_id\"].nunique() * 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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 = df[\"metadata\"] #.apply(lambda x: custom_parse(x))\n",
    "test_slice_series = test_slice.apply(pd.Series)\n",
    "df = pd.concat([df, test_slice_series], axis=1, join=\"inner\")\n",
    "print(\"unique_requests\", df[\"request_id\"].nunique())\n",
    "# remove the duplicates\n",
    "print(\"before removing duplicates\", df.shape)\n",
    "df = df.loc[:, ~df.columns.duplicated()].copy()\n",
    "print(\"after removing duplicates\", df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"task\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(\"positive\", df[df[\"preference\"]][\"param_experiment\"].value_counts())\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df[df[\"preference\"]][\"cer_diff_preference\"].hist(bins=50)\n",
    "# print(df[df[\"preference\"]][\"cer_diff_preference\"].quantile(0.95))\n",
    "# plt.show()\n",
    "# print(df[df[\"preference\"]][\"cer\"].hist(bins=50))\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"id\"] = df[\"str_id\"]\n",
    "# get the original duration of the clips, if they are concacted\n",
    "df[\"original_duration_s\"] = df[\"total_start_s\"] + df[\"duration\"]\n",
    "# classify the continue at behavoirs by the duration choice\n",
    "audio_prompt_id_to_continue_at = {}\n",
    "\n",
    "for _, row in df[~df[\"continued_parent\"].isna()].iterrows():\n",
    "    audio_prompt_id = row[\"continued_parent\"]\n",
    "    if audio_prompt_id not in audio_prompt_id_to_continue_at:\n",
    "        audio_prompt_id_to_continue_at[audio_prompt_id] = row[\"continue_at\"]\n",
    "    else:\n",
    "        # pick the max\n",
    "        audio_prompt_id = max(\n",
    "            audio_prompt_id_to_continue_at[audio_prompt_id], row[\"continue_at\"]\n",
    "        )\n",
    "print(len(audio_prompt_id_to_continue_at))\n",
    "df[\"has_continue_and_start_continue_at\"] = df[\"id\"].apply(\n",
    "    lambda x: audio_prompt_id_to_continue_at.get(x)\n",
    ")\n",
    "# we want continue at to be at most of the clip...\n",
    "df[\"good_continue_at\"] = (\n",
    "    (df[\"has_continue_and_start_continue_at\"] / df[\"duration\"]) > 0.9\n",
    ") | df[\"has_continue_and_start_continue_at\"].isna()\n",
    "print(df[\"good_continue_at\"].value_counts())\n",
    "\n",
    "\n",
    "print(\n",
    "    \"\\n Check some basics... \\n\",\n",
    "    df[\"preference\"].value_counts(),\n",
    "    df[\"is_30b\"].value_counts(),\n",
    "    df[\"model_name\"].value_counts(),\n",
    "    df.groupby([\"preference\"])[\"model_name\"].value_counts(),\n",
    ")\n",
    "\n",
    "df = df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "df[\"duration_rel_diff\"] = df[\"duration\"].diff()\n",
    "df[\"play_rel_diff\"] = df[\"reaction_play_count\"].diff()\n",
    "print(df[\"task\"].value_counts())\n",
    "\n",
    "df[\"post_infill_duration\"] = (df[\"duration\"] \n",
    "                              + df[\"infill_context_end_s\"] \n",
    "                              - df[\"infill_context_start_s\"] \n",
    "                              - df[\"include_future_s\"] \n",
    "                              - df[\"include_history_s\"] \n",
    "                              - df[\"infill_dur_s\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"post_infill_duration\"].hist(bins=np.linspace(-5, 360, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df = df[df[\"task\"].isin([\"extend\", \"\"])].copy()\n",
    "# print(df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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 = 1\n",
    "# this is a filter on the concated clip\n",
    "concat_total_play_count = 3\n",
    "\n",
    "all_fitlers = (df[\"duration\"] >= 10) & (df[\"duration\"] <= 240)\n",
    "print(\"after duration\", all_fitlers.sum() /  df.shape[0])\n",
    "infill_duration_filter =  (~df[\"task\"].isin(\n",
    "        [\n",
    "            \"infill\",\n",
    "            \"infill_intro\",\n",
    "            \"infill_outro\",\n",
    "        ]\n",
    "    )) | (df[\"post_infill_duration\"] <= 239)\n",
    "print(\"after infill duration\", infill_duration_filter.sum() / df.shape[0])\n",
    "# negative fitlers\n",
    "total_negative = df[~df[\"preference\"]].shape[0]\n",
    "neg_filter_reaction_play_count = (~df[\"preference\"]) & (df[\"reaction_play_count\"] >= 1)\n",
    "print(\n",
    "    \"neg_filter_reaction_play_count\",\n",
    "    round(neg_filter_reaction_play_count.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_upvote_count = (~df[\"preference\"]) & (df[\"upvote_count\"] == 0)\n",
    "print(\n",
    "    \"neg_filter_upvote_count\",\n",
    "    round(neg_filter_upvote_count.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_norm_play_frac = (~df[\"preference\"]) & (df[\"norm_play_frac\"] <= 3.1)\n",
    "print(\n",
    "    \"neg_filter_norm_play_frac\",\n",
    "    round(neg_filter_norm_play_frac.sum() / total_negative, 4),\n",
    ")\n",
    "neg_filter_continues = (~df[\"preference\"]) & (\n",
    "    df[\"has_continue_and_start_continue_at\"].isna()\n",
    ")\n",
    "print(\n",
    "    \"neg_filter_continues\",\n",
    "    round(neg_filter_continues.sum() / total_negative, 4),\n",
    ")\n",
    "\n",
    "neg_filter_selection_mask = (\n",
    "    all_fitlers\n",
    "    & infill_duration_filter\n",
    "    & neg_filter_reaction_play_count\n",
    "    & neg_filter_upvote_count\n",
    "    & neg_filter_norm_play_frac\n",
    "    & neg_filter_continues\n",
    ")\n",
    "\n",
    "print(\"----------------\")\n",
    "total_positive = df[df[\"preference\"]].shape[0]\n",
    "assert total_positive == total_negative\n",
    "pos_filter_continues = (df[\"preference\"]) & (df[\"good_continue_at\"])\n",
    "print(\"pos_filter_continues\", round(pos_filter_continues.sum() / total_positive, 4))\n",
    "pos_filter_reaction_play_count = (df[\"preference\"]) & (df[\"reaction_play_count\"] >= 1)\n",
    "print(\n",
    "    \"pos_filter_reaction_play_count\",\n",
    "    round(pos_filter_reaction_play_count.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_relative_play_count = (df[\"preference\"]) & (df[\"play_rel_diff\"] >= 0)\n",
    "print(\n",
    "    \"pos_filter_relative_play_count\",\n",
    "    round(pos_filter_relative_play_count.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_cer_diff_preference = (\n",
    "    (df[\"preference\"]) # & (df[\"pos_diff_preference\"] == 2) # & (df[\"cer_diff_preference\"] < 0.5) & (df[\"cer\"] < 0.99)\n",
    ")\n",
    "print(\n",
    "    \"pos_filter_cer_diff_preference\",\n",
    "    round(pos_filter_cer_diff_preference.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_bad_flags = (\n",
    "    (df[\"preference\"]) & (df[\"flag_count\"] == 0) & (df[\"dislike_count\"] == 0)\n",
    ")\n",
    "print(\n",
    "    \"pos_filter_bad_flags\",\n",
    "    round(pos_filter_bad_flags.sum() / total_positive, 4),\n",
    ")\n",
    "pos_filter_play_counts = (df[\"preference\"]) & (\n",
    "    (\n",
    "        (df[\"part_of_concat\"])\n",
    "        & (df[\"reaction_play_count\"] >= concat_pos_play_count)\n",
    "        & (df[\"concat_play_counts\"] >= concat_total_play_count)\n",
    "    )\n",
    "    | (\n",
    "        (~df[\"part_of_concat\"]) & (df[\"reaction_play_count\"] >= normal_pos_play_count)\n",
    "        # & (df[\"norm_play_frac\"] >= 2.1)  # this is a bit of a luxury cut...\n",
    "    )\n",
    ")\n",
    "print(\n",
    "    \"after filter on play counts\",\n",
    "    round(pos_filter_play_counts.sum() / total_positive, 4),\n",
    ")\n",
    "high_quality_tasks_filter = (\n",
    "    (\n",
    "        df[\"task\"].isin(\n",
    "            [\n",
    "                \"cover\",\n",
    "                \"extend\",\n",
    "                \"artist_consistency\",\n",
    "                \"\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    & (\n",
    "        (df[\"upvote_count\"] >= 1) # (df[\"upvote_count\"] >= 1)\n",
    "        | (df[\"reaction_play_count\"] >= 5)\n",
    "        | (df[\"concat_play_counts\"] >= 5)\n",
    "    )\n",
    "    & (\n",
    "        (df[\"part_of_concat\"])\n",
    "        | (\n",
    "            (~df[\"part_of_concat\"])\n",
    "            & (df[\"norm_play_frac\"] >= 5.1)  # this is a bit of a luxury cut...\n",
    "            & (df[\"norm_play_frac\"] >= df[\"reaction_play_count\"] / 3) # play duration is not low on average\n",
    "        )\n",
    "    )\n",
    ")\n",
    "medium_quality_tasks_filter = (\n",
    "    df[\"task\"].isin(\n",
    "        [\n",
    "            \"infill\",\n",
    "            \"infill_intro\",\n",
    "            \"infill_outro\",\n",
    "        ]\n",
    "    )\n",
    ") & (\n",
    "    (df[\"upvote_count\"] >= 1) # (df[\"upvote_count\"] >= 1)  (df[\"pos_diff_preference\"] == 2)\n",
    "    | (df[\"reaction_play_count\"] >= 5)\n",
    "    | (df[\"concat_play_counts\"] >= 5)\n",
    ")\n",
    "pos_filter_higher_quality = (df[\"preference\"]) & (\n",
    "    high_quality_tasks_filter | medium_quality_tasks_filter\n",
    ")\n",
    "print(\n",
    "    \"after filter on higher quality\",\n",
    "    round(pos_filter_higher_quality.sum() / total_positive, 4),\n",
    ")\n",
    "\n",
    "user_gen_filter = (df[\"user_n_clips\"] >= 100)  # user needs to have genereated at least 100 over the time period\n",
    "\n",
    "print(\"----------------\")\n",
    "pos_filter_selectin_mask = (\n",
    "    (df[\"preference\"])  # get basics aligned\n",
    "    & all_fitlers\n",
    "    & infill_duration_filter\n",
    "    & pos_filter_continues\n",
    "    & pos_filter_reaction_play_count\n",
    "    & pos_filter_relative_play_count\n",
    "    & pos_filter_cer_diff_preference\n",
    "    & pos_filter_bad_flags\n",
    "    & pos_filter_play_counts\n",
    "    & pos_filter_higher_quality\n",
    "    & user_gen_filter\n",
    ")\n",
    "print(\n",
    "    \"negative\",\n",
    "    sum(neg_filter_selection_mask),\n",
    "    \"positive\",\n",
    "    sum(pos_filter_selectin_mask),\n",
    ")\n",
    "\n",
    "neg_filter_requests = df[neg_filter_selection_mask][\"request_id\"].unique()\n",
    "pos_filter_requests = df[pos_filter_selectin_mask][\"request_id\"].unique()\n",
    "# looking for very strong signal here:\n",
    "# listen to the positive/negative more than once\n",
    "# disliked one of the clips\n",
    "unique_requests = set(pos_filter_requests).intersection(neg_filter_requests)\n",
    "print(\"----------------\")\n",
    "print(\n",
    "    \"total pair requests\",\n",
    "    df[\"request_id\"].nunique(),\n",
    "    \" --> selected pair requests\",\n",
    "    len(unique_requests),\n",
    "    f\"frac {len(unique_requests) / df['request_id'].nunique():.3f}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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",
    "# 76171 152342 total khrs 2.880 n gpus for 1250 iters 3.809\n",
    "# v10 has 78866\n",
    "# v14 has 110402\n",
    "# 30b_t4_v12 requests 47755 clips 95510 total khrs 4.555; N gpus for 1000 iters 5.969; 4 gpus for x iters 1492.344; n unique users 13790 n pro users 13526\n",
    "# 30b_t4_v20 requests 18306 clips 36612 total khrs 1.754; N gpus for 1000 iters 2.288; 4 gpus for x iters 572.062; n unique users 6943 n pro users 6714\n",
    "# 30b_t5_v5 requests 21318 clips 42636 total khrs 1.938; N gpus for 1000 iters 2.665; 4 gpus for x iters 666.188; n unique users 11661 n pro users 9347\n",
    "# 30b_t5_v6 requests 22833 clips 45666 total khrs 2.072; N gpus for 1000 iters 2.854; 4 gpus for x iters 713.531; n unique users 12576 n pro users 9910\n",
    "# 30b_t5_v7 requests 33197 clips 66394 total khrs 2.846; N gpus for 1000 iters 4.150; 4 gpus for x iters 1037.406; n unique users 16665 n pro users 13050\n",
    "# 30b_t5_v8 requests 20907 clips 41814 total khrs 1.629; N gpus for 1000 iters 2.613; 4 gpus for x iters 653.344; n unique users 11102 n pro users 9052\n",
    "# 30b_t5_v9 requests 24645 clips 49290 total khrs 1.918; N gpus for 1000 iters 3.081; 4 gpus for x iters 770.156; n unique users 12661 n pro users 10229\n",
    "# 30b_t5_v11 requests 21975 clips 43950 total khrs 1.508; N gpus for 1000 iters 2.747; 4 gpus for x iters 686.719; n unique users 11173 n pro users 9270\n",
    "# 30b_t5_v12 requests 21436 clips 42872 total khrs 1.582; N gpus for 1000 iters 2.679; 4 gpus for x iters 669.875; n unique users 11553 n pro users 9374\n",
    "# 30b_t5_v13 requests 25912 clips 51824 total khrs 1.922; N gpus for 1000 iters 3.239; 4 gpus for x iters 809.750; n unique users 13429 n pro users 10737\n",
    "# 30b_t5_v14 requests 27687 clips 55374 total khrs 2.036; N gpus for 1000 iters 3.461; 4 gpus for x iters 865.219; n unique users 14196 n pro users 11227\n",
    "# 30b_t5_v15 requests 29563 clips 59126 total khrs 2.225; N gpus for 1000 iters 3.695; 4 gpus for x iters 923.844; n unique users 14684 n pro users 11642\n",
    "# 30b_t6_v1 requests 24678 clips 49356 total khrs 1.934; N gpus for 1000 iters 3.085; 4 gpus for x iters 771.188; n unique users 11350 n pro users 10065\n",
    "# 30b_t6_v2 requests 49341 clips 98682 total khrs 4.332; N gpus for 1000 iters 6.168; 4 gpus for x iters 1541.906; n unique users 19992 n pro users 16972\n",
    "# 30b_t6_v3 requests 54967 clips 109934 total khrs 4.826; N gpus for 1000 iters 6.871; 4 gpus for x iters 1717.719; n unique users 21653 n pro users 18266\n",
    "# 30b_t6_v5 requests 29459 clips 58918 total khrs 2.657; N gpus for 1000 iters 3.682; 4 gpus for x iters 920.594; n unique users 14091 n pro users 11884\n",
    "# 30b_t6_v7 requests 40056 clips 80112 total khrs 3.354; N gpus for 1000 iters 5.007; 4 gpus for x iters 1251.750; n unique users 17440 n pro users 14476\n",
    "# 30b_t6_v8 requests 31681 clips 63362 total khrs 2.681; N gpus for 1000 iters 3.960; 4 gpus for x iters 990.031; n unique users 15180 n pro users 15078\n",
    "# 30b_t6_v9 requests 55785 clips 111570 total khrs 4.727; N gpus for 1000 iters 6.973; 4 gpus for x iters 1743.281; n unique users 23114 n pro users 22871\n",
    "# 30b_t6_v10 requests 52209 clips 104418 total khrs 4.515; N gpus for 1000 iters 6.526; 4 gpus for x iters 1631.531; n unique users 22216 n pro users 21893\n",
    "# 30b_t6_v11 requests 31338 clips 62676 total khrs 2.493; N gpus for 1000 iters 3.917; 4 gpus for x iters 979.312; n unique users 15190 n pro users 14973\n",
    "# 30b_t6_v13 requests 50856 clips 101712 total khrs 3.987; N gpus for 1000 iters 6.357; 4 gpus for x iters 1589.250; n unique users 23547 n pro users 23097\n",
    "# 30b_t6_v15 requests 65213 clips 130426 total khrs 5.217; N gpus for 1000 iters 8.152; 4 gpus for x iters 2037.906; n unique users 28763 n pro users 27964\n",
    "# 30b_t6_v16 requests 37907 clips 75814 total khrs 2.956; N gpus for 1000 iters 4.738; 4 gpus for x iters 1184.594; n unique users 18448 n pro users 17212\n",
    "# 30b_t6_v17 requests 51543 clips 103086 total khrs 4.154; N gpus for 1000 iters 6.443; 4 gpus for x iters 1610.719; n unique users 24223 n pro users 21829\n",
    "# 30b_t6_v18 requests 94965 clips 189930 total khrs 7.483; N gpus for 1000 iters 11.871; 4 gpus for x iters 2967.656; n unique users 23852 n pro users 21767\n",
    "# 30b_t6_v19 requests 34289 clips 68578 total khrs 2.846; N gpus for 1000 iters 4.286; 4 gpus for x iters 1071.531; n unique users 11802 n pro users 10830\n",
    "# 30b_t6_v20 requests 53258 clips 106516 total khrs 4.270; N gpus for 1000 iters 6.657; 4 gpus for x iters 1664.312; n unique users 22096 n pro users 21171\n",
    "# 30b_t6_v22 requests 29971 clips 59942 total khrs 1.992; N gpus for 1000 iters 3.746; 4 gpus for x iters 936.594; n unique users 13611 n pro users 13036\n",
    "# 30b_t6_v24 requests 45333 clips 90666 total khrs 3.343; N gpus for 1000 iters 5.667; 4 gpus for x iters 1416.656; n unique users 18853 n pro users 18049\n",
    "# 30b_t6_v25 requests 28154 clips 56308 total khrs 2.401; N gpus for 1000 iters 3.519; 4 gpus for x iters 879.812; n unique users 12946 n pro users 12361\n",
    "# 30b_t6_v26 requests 30428 clips 60856 total khrs 2.599; N gpus for 1000 iters 3.804; 4 gpus for x iters 950.875; n unique users 13822 n pro users 13163\n",
    "# 30b_t6_v27 requests 43731 clips 87462 total khrs 3.862; N gpus for 1000 iters 5.466; 4 gpus for x iters 1366.594; n unique users 17898 n pro users 16865\n",
    "# 30b_t6_v28 requests 56869 clips 113738 total khrs 5.064; N gpus for 1000 iters 7.109; 4 gpus for x iters 1777.156; n unique users 22004 n pro users 20244\n",
    "# 30b_t6_v29 requests 32729 clips 65458 total khrs 2.604; N gpus for 1000 iters 4.091; 4 gpus for x iters 1022.781; n unique users 13739 n pro users 13503\n",
    "# 30b_t6_v30 requests 32553 clips 65106 total khrs 2.592; N gpus for 1000 iters 4.069; 4 gpus for x iters 1017.281; n unique users 13712 n pro users 13478\n",
    "# 30b_t6_v31 requests 41523 clips 83046 total khrs 3.164; N gpus for 1000 iters 5.190; 4 gpus for x iters 1297.594; n unique users 16266 n pro users 15967\n",
    "# 30b_t6_v32 requests 55942 clips 111884 total khrs 4.361; N gpus for 1000 iters 6.993; 4 gpus for x iters 1748.188; n unique users 20808 n pro users 20271\n",
    "# 30b_t6_v33 requests 182096 clips 364192 total khrs 13.745; N gpus for 1000 iters 22.762; 4 gpus for x iters 5690.500; n unique users 40437 n pro users 34524"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mask = (df_slice[\"preference\"]) & (\n",
    "    (df_slice[\"is_in_playlist\"]) | (df_slice[\"concat_in_playlist\"])\n",
    ")\n",
    "print(\"positive in playlist\", df_slice[test_mask].shape)\n",
    "print(df_slice[\"task\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_slice[\"post_infill_duration\"].hist(bins=np.linspace(-5, 300, 100))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": null,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "# df_slice_2 = pd.read_pickle(\"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241118_full_slice.pkl\")\n",
    "# df_total = pd.concat([df_slice, df_slice_2])\n",
    "# print(df_total.shape)\n",
    "# df_slice = df_total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice_prev = pd.read_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_v13_20240902_slice.pkl\"\n",
    "# )\n",
    "# df_slice_prev = df_slice_prev[\n",
    "#     (\n",
    "#         (df_slice_prev[\"task\"] != \"infill\")\n",
    "#         & (df_slice_prev[\"task\"] != \"cover\")\n",
    "#         & (df_slice_prev[\"task\"] != \"artist_consistency\")\n",
    "#     )\n",
    "# ].copy()\n",
    "# df_slice = pd.concat([df_slice, df_slice_prev])\n",
    "# print(df_slice.shape)\n",
    "# print(df_slice[\"task\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice[(df_slice[\"preference\"]) & ((df_slice[\"task\"].str.strip() == \"\") | (df_slice[\"task\"].str.strip() == \"cover\"))].to_pickle(\n",
    "#     \"/home/tony/Data/Preference/30b_v5/interesting_clips_v4_t_5_20241018_full_with_cer_pos_gen.pkl\"\n",
    "# )\n",
    "# df_slice[(df_slice[\"preference\"]) & ((df_slice[\"task\"].str.strip() == \"\") | (df_slice[\"task\"].str.strip() == \"cover\"))].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice.to_pickle(\"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_h_t_6_20250108_full_slice.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"/home/tony/Data/Preference/30b_v6/similarities_30b_t6.json\", \"r\") as fp:\n",
    "#     all_similarities = json.load(fp)\n",
    "    \n",
    "# df_slice[\"similarity\"] = df_slice[\"id\"].map(all_similarities)\n",
    "# df_cover = df_slice[df_slice[\"task\"] == \"cover\"].copy()\n",
    "# df_artist = df_slice[df_slice[\"task\"] == \"artist_consistency\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_cover[df_cover[\"preference\"]][\"similarity\"].hist(bins=200, label=\"positive\", alpha=0.5)\n",
    "# df_cover[~df_cover[\"preference\"]][\"similarity\"].hist(bins=200, label=\"negative\", alpha=0.5)\n",
    "# plt.legend()\n",
    "# plt.show()\n",
    "# df_cover[\"similarity_diff\"] = df_cover[\"similarity\"].diff()\n",
    "# df_cover[df_cover[\"preference\"]][\"similarity_diff\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# df_cover[\"continued_parent\"] = None\n",
    "# df_cover[\"continue_at\"] = -1\n",
    "# too_similar = (df_cover[\"similarity\"] >= 0.99)\n",
    "# too_different = ( df_cover[\"similarity\"] < 0.25)\n",
    "# too_big_difference = (df_cover[\"similarity_diff\"] < -0.2)\n",
    "# df_cover_drops_id = df_cover[\n",
    "#     (df_cover[\"preference\"])\n",
    "#     & (too_similar | too_different | too_big_difference)\n",
    "# ][\"s3_id\"].unique()\n",
    "# print(\"total rows\", df_cover.shape[0], \"total drops\", len(df_cover_drops_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_artist[df_artist[\"preference\"]][\"similarity\"].hist(bins=200, label=\"positive\", alpha=0.5)\n",
    "# df_artist[~df_artist[\"preference\"]][\"similarity\"].hist(bins=200, label=\"negative\", alpha=0.5)\n",
    "# plt.legend()\n",
    "# plt.show()\n",
    "# df_artist[\"similarity_diff\"] = df_artist[\"similarity\"].diff()\n",
    "# df_artist[df_artist[\"preference\"]][\"similarity_diff\"].hist(bins=200)\n",
    "# plt.show()\n",
    "# too_similar = (df_artist[\"similarity\"] >= 0.95)\n",
    "# too_different = ( df_artist[\"similarity\"] < 0.4)\n",
    "# too_big_difference = (df_artist[\"similarity_diff\"] < -0.2)\n",
    "# df_artist_drops_id = df_artist[\n",
    "#     (df_artist[\"preference\"])\n",
    "#     & (too_similar | too_different | too_big_difference)\n",
    "# ][\"s3_id\"].unique()\n",
    "# print(\"total rows\", df_artist.shape[0], \"total drops\", len(df_artist_drops_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"before filter on cover similarity\", df_slice.shape)\n",
    "# df_slice = df_slice[~df_slice[\"s3_id\"].isin(df_cover_drops_id)].copy()\n",
    "# print(\"after filter on cover similarity\", df_slice.shape)\n",
    "# df_slice = df_slice[~df_slice[\"s3_id\"].isin(df_artist_drops_id)].copy()\n",
    "# print(\"after filter on artist similarity\", df_slice.shape)\n",
    "# df_slice = df_slice[\n",
    "#     df_slice[\"request_id\"].isin(\n",
    "#         df_slice[\"request_id\"].value_counts().index[df_slice[\"request_id\"].value_counts() == 2]\n",
    "#     )\n",
    "# ]\n",
    "# print(\"after filter on request id pairs\", df_slice.shape)\n",
    "# print(df_slice.groupby([\"preference\"])[\"model_name\"].value_counts())\n",
    "# assert df_slice.shape[0] == df_slice[\"request_id\"].nunique() * 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df_slice[\"is_public\"].value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice.to_pickle(\"/home/tony/Data/Preference/30b_v6/interesting_clips_v4_t_6_20241118_full_slice.pkl\")\n",
    "print(df_slice.shape)\n",
    "print(df_slice[\"task\"].value_counts())\n",
    "# df_slice = df_slice[df_slice[\"task\"].isin([\"extend\", \"\"])].copy()\n",
    "print(df_slice.shape)\n",
    "print(df_slice[\"task\"].value_counts())\n",
    "BREAK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Need to kick out the ones has gpt prompt -- these are pairs with different text inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# don't have continue at\n",
    "df_slice[\"request_id\"] = df_slice[\"request_id\"].astype(str)\n",
    "# df_slice[df_slice[\"continue_at\"].isna()][\"request_id\"].nunique(), df_slice[\"request_id\"].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df_slice.shape)\n",
    "df_slice = df_slice[df_slice[\"request_id\"].apply(lambda x: len(x) > 3)]\n",
    "print(df_slice.shape)\n",
    "# df_slice = df_slice[df_slice[\"is_pro_user\"]].copy()\n",
    "print(df_slice.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_filtered_requests = df_slice[\"request_id\"].astype(str).unique()\n",
    "# final_filtered_requests = df_slice[df_slice[\"is_pro_user\"]][\"request_id\"].astype(str).unique()\n",
    "print(len(final_filtered_requests))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice.to_csv(\"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_20240902_slice.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_slice[\"continue_at\"] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_requests, val_requests = train_test_split(\n",
    "    sorted(list(final_filtered_requests)), test_size=0.01, random_state=42\n",
    ")\n",
    "print(len(train_requests), len(val_requests))\n",
    "\n",
    "train_df = df_slice[df_slice[\"request_id\"].astype(str).isin(set(train_requests))].copy()\n",
    "val_df = df_slice[df_slice[\"request_id\"].astype(str).isin(set(val_requests))].copy()\n",
    "train_df = train_df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "train_df = train_df  # .reset_index()\n",
    "val_df = val_df.sort_values(by=[\"request_id\", \"preference\"])\n",
    "val_df = val_df  # .reset_index()\n",
    "\n",
    "print(train_df.shape, val_df.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actually make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# val_df[[\"request_id\", \"metadata\", \"updated_at\", \"user_id\", \"preference\"]].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_duration = 0\n",
    "for i, row in tqdm(train_df.iterrows(), total=len(train_df)):\n",
    "    # we need to alternate between preference: neg, pos\n",
    "    # print(i, row)\n",
    "    try:\n",
    "        assert row[\"preference\"] == (i % 2 == 1)\n",
    "        total_duration += row[\"duration\"]\n",
    "    except Exception as E:\n",
    "        print(i, row)\n",
    "        print(E)\n",
    "        raise ValueError()\n",
    "\n",
    "print(\n",
    "    f\"{round(total_duration / 60 / 60):,} hours of {train_df.shape[0]} clips, {train_df.shape[0] / 8 / 2 / 1000} nodes, {train_df.shape[0] / 8 / 2 / 6} iters\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_dataset(val_df, OUT_DATA_DIR, is_val=True, npz_dir=NPZ_DIR, t_data_memmap=N_TOKENS_AUDIO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_npz = np.load(\"/app/suno/data/dpo/30b_npz/26d19085-18da-4701-af43-122684543891.npz\")\n",
    "# for k in test_npz.keys():\n",
    "#     print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_dataset(train_df, OUT_DATA_DIR, is_val=False, npz_dir=NPZ_DIR, t_data_memmap=N_TOKENS_AUDIO)"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "mm = np.memmap(os.path.join(OUT_DATA_DIR, f\"data_val.bin\"), dtype=np.uint16, mode=\"r\")\n",
    "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"meta_val.jsonl\"))\n",
    "test_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_val.json\"))\n",
    "mm = mm.reshape(-1, N_TOKENS_AUDIO, 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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_counts = Counter()\n",
    "for test_meta in test_metas:\n",
    "    task_counts[test_meta.get(\"task\")] += 1\n",
    "print(task_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # randomly listen to some stuff\n",
    "# from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models\n",
    "# from suno_utils.tasks.dac_2c_12cb import (\n",
    "#     encode as codec_encode,\n",
    "#     decode_stream_to_full_audio as codec_decode,\n",
    "#     EMBEDDING_RATE as CODEC_EMBEDDING_RATE,\n",
    "#     decode as decode\n",
    "# )\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "# _ = preload_codec_models(\"/app/suno/data/dpo/models/dac_2c_25x12.pt\", device=\"cuda\")\n",
    "# assert len(test_metas) == len(mm)\n",
    "# idx_list = list(range(len(test_metas)))\n",
    "# # random.shuffle(idx_list)\n",
    "# # idx_list = [idx for idx in idx_list if \"text\" in test_metas[idx]]\n",
    "# print(len(mm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import random\n",
    "# idx = random.choice(test_info[\"perference_0\"][\"idx_list\"])\n",
    "# assert \"original_duration_s\" in test_metas[idx]\n",
    "# # positive index should be shifted by 1\n",
    "# pos_idx = idx + 1\n",
    "# print(\n",
    "#     \"tags:\",\n",
    "#     test_metas[idx].get(\"tags\") == test_metas[pos_idx].get(\"tags\"),\n",
    "#     test_metas[idx].get(\"tags\"),\n",
    "# )\n",
    "# arr = mm[idx, 1:].copy().astype(np.int16)[:, 1:]\n",
    "# pos_arr = mm[pos_idx, 1:].copy().astype(np.int16)[:, 1:]\n",
    "# pad_idx_arr = np.where(arr == COARSE_PAD_TOKEN)[0]\n",
    "# if len(pad_idx_arr) > 0:\n",
    "#     arr = arr[: pad_idx_arr[0], :]\n",
    "# pos_pad_idx_arr = np.where(pos_arr == COARSE_PAD_TOKEN)[0]\n",
    "# if len(pos_pad_idx_arr) > 0:\n",
    "#     pos_arr = pos_arr[: pos_pad_idx_arr[0], :]\n",
    "# a = decode(arr)\n",
    "# print(\"\\n negative example \\n\", test_metas[idx])\n",
    "# a.play(compress=False)\n",
    "# pos_a = decode(pos_arr)\n",
    "# print(\"\\n positive example \\n\", test_metas[pos_idx])\n",
    "# pos_a.play(compress=False)\n",
    "# print(\n",
    "#     \"text:\",\n",
    "#     test_metas[idx].get(\"text\") == test_metas[pos_idx].get(\"text\"),\n",
    "#     test_metas[idx].get(\"text\"),\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# val_df[val_df[\"tags\"] == 'a vibrant blend of experimental jazz fusion, drum-and-bass and swagger fuzzed-out guitars']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from collections import Counter\n",
    "# c = Counter()\n",
    "# for _, row in df_slice.iterrows():\n",
    "#     # print(row[\"metadata\"])\n",
    "#     for k in ast.literal_eval(row[\"metadata\"]).keys():\n",
    "#         c[k] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "def validation_on_metas(input_metas):\n",
    "    total_bad = 0\n",
    "    total_good = 0\n",
    "    for idx in range(len(input_metas)):\n",
    "        if idx % 2 == 0:\n",
    "            pos_idx = idx + 1\n",
    "            if input_metas[idx].get(\"tags\") != input_metas[pos_idx].get(\"tags\"):\n",
    "                # print(test_metas[idx].get(\"text\") == test_metas[pos_idx].get(\"text\"), test_metas[idx].get(\"tags\"), test_metas[pos_idx].get(\"tags\"))\n",
    "                total_bad += 1\n",
    "            else:\n",
    "                total_good += 1\n",
    "    print(total_good, total_bad)\n",
    "    return\n",
    "\n",
    "\n",
    "validation_on_metas(test_metas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_info = read_json(os.path.join(OUT_DATA_DIR, f\"info_tr.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "outputs": [],
   "source": [
    "print(\"1 epoch per batch 4, total\", total_iters / 8 / 2 / 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd /home/tony/Work/tony/slurm/30b_dpo && sbatch sbatch_ipo_30b"
   ]
  },
  {
   "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_13b_v4_t6.ipynb\",\n",
    "    os.path.join(OUT_DATA_DIR, \"make_dataset.ipynb\"),\n",
    ")\n",
    "print(\"Cache kept!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# some gymathtics loading prev data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prev_v3_data = \"/app/suno/data/dpo/7v_v20_full/\"\n",
    "\n",
    "# test_val_metas = read_jsonl(os.path.join(prev_v3_data, f\"meta_val.jsonl\"))\n",
    "# test_tr_metas = read_jsonl(os.path.join(prev_v3_data, f\"meta_tr.jsonl\"))\n",
    "\n",
    "# all_ids = set()\n",
    "# for meta in test_val_metas:\n",
    "#     all_ids.add(meta[\"id\"])\n",
    "# for meta in test_tr_metas:\n",
    "#     all_ids.add(meta[\"id\"])\n",
    "# print(len(all_ids), len(test_val_metas) + len(test_tr_metas))\n",
    "\n",
    "# all_ids = list(all_ids)\n",
    "# with open(\"/home/tony/Data/Preference/7b_v2/7v_v20_full_recut_id.json\", \"w\") as fp:\n",
    "#     json.dump(all_ids, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x_data = train_df[train_df[\"preference\"]][\"similarity\"]\n",
    "# y_data = train_df[~train_df[\"preference\"]][\"similarity\"]\n",
    "# from matplotlib.colors import LogNorm\n",
    "\n",
    "# fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 10))\n",
    "\n",
    "# # 2D Histogram\n",
    "# h = ax1.hist2d(\n",
    "#     x_data,\n",
    "#     y_data,\n",
    "#     bins=(50, 50),\n",
    "#     cmap=\"coolwarm\",\n",
    "#     range=[[0, 1], [0, 1]],\n",
    "#     norm=LogNorm(),\n",
    "# )\n",
    "\n",
    "# ax1.set_xlabel(\"Semantic Distance (Preferred)\")\n",
    "# ax1.set_ylabel(\"Semantic Distance (Non-Preferred)\")\n",
    "# ax1.set_title(\n",
    "#     \"2D Histogram of Semantic Distances: Preferred vs Non-Preferred (Log Scale)\"\n",
    "# )\n",
    "\n",
    "# cbar1 = plt.colorbar(h[3], ax=ax1)\n",
    "# cbar1.set_label(\"Number of Request IDs (Log Scale)\")\n",
    "\n",
    "# # Scatter plot\n",
    "# ax2.scatter(x_data, y_data, alpha=0.1, s=1)\n",
    "# ax2.set_xlabel(\"Semantic Distance (Preferred)\")\n",
    "# ax2.set_ylabel(\"Semantic Distance (Non-Preferred)\")\n",
    "# ax2.set_title(\"Scatter Plot of Semantic Distances: Preferred vs Non-Preferred\")\n",
    "# ax2.set_xlim(0, 1)\n",
    "# ax2.set_ylim(0, 1)\n",
    "\n",
    "# plt.tight_layout()\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"meta_tr.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_info.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
