{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04dc079",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T22:42:11.882092Z",
     "start_time": "2024-04-11T22:42:11.872532Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e338c94",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T22:42:14.273454Z",
     "start_time": "2024-04-11T22:42:11.883068Z"
    }
   },
   "outputs": [],
   "source": [
    "from preference_data_selection import *\n",
    "from preference_helper import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15b7a2e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T22:42:14.467215Z",
     "start_time": "2024-04-11T22:42:14.275017Z"
    }
   },
   "outputs": [],
   "source": [
    "engine = sqlalchemy.create_engine(\n",
    "    \"postgresql://postgres:cfAoZgBhraL0bDRd2TAn@suno-rds-prod.cluster-cnfvffydbwvc.us-east-2.rds.amazonaws.com/studio_hga1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bbe068",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T22:55:32.776731Z",
     "start_time": "2024-04-11T22:42:14.469081Z"
    }
   },
   "outputs": [],
   "source": [
    "updated_info_dict = get_info_from_engine(engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca283520",
   "metadata": {},
   "source": [
    "# load the prev one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8161dc0f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T23:33:52.284604Z",
     "start_time": "2024-04-11T22:55:32.777933Z"
    }
   },
   "outputs": [],
   "source": [
    "# Read CSV in an iterable way\n",
    "prev_clip_df = pd.read_csv('/app/suno/data/dpo/data_prep/clip_df_v3_20240410.csv', engine='python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11630295",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-11T23:33:52.686431Z",
     "start_time": "2024-04-11T23:33:52.286594Z"
    }
   },
   "outputs": [],
   "source": [
    "# # only if you need to split\n",
    "# non_empty_mask = (~clip_df[\"model_name\"].isna()) & (~clip_df[\"id\"].isna())\n",
    "# print(sum(non_empty_mask))\n",
    "# v3_mask = clip_df[\"model_name\"].str.contains(\"v3\")\n",
    "# v3_final_mask = non_empty_mask & v3_mask\n",
    "# print(sum(v3_final_mask))\n",
    "# # clip_df[v3_final_mask].to_csv('/app/suno/data/dpo/data_prep/clip_df_v3_20240410.csv', index=False)\n",
    "\n",
    "# v2_mask = clip_df[\"model_name\"].str.contains(\"v2\")\n",
    "# v2_final_mask = non_empty_mask & v2_mask\n",
    "# print(sum(v2_final_mask))\n",
    "# # clip_df[v2_final_mask].to_csv('/app/suno/data/dpo/data_prep/clip_df_v2_20240410.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce1b7f97",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:03:28.906243Z",
     "start_time": "2024-04-12T01:03:28.590910Z"
    }
   },
   "outputs": [],
   "source": [
    "# Get the most up to date clips"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "833d7bc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_clip_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9b8e42",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:07:25.999017Z",
     "start_time": "2024-04-12T01:07:22.823585Z"
    }
   },
   "outputs": [],
   "source": [
    "CLIP_QUERY = \"\"\"\n",
    "SELECT * FROM bots_generatedclip\n",
    "WHERE status='complete' AND created_at>='2024-04-10' \n",
    "\"\"\"\n",
    "total_clip_df = pd.read_sql_query(CLIP_QUERY, engine)\n",
    "print(f\"Total clips df: {total_clip_df.shape[0]:,} rows\")\n",
    "# filter on versions\n",
    "new_clip_df = total_clip_df[\n",
    "    (total_clip_df[\"model_name\"].str.contains(\"v3\"))  # or v3...\n",
    "    & (total_clip_df[\"created_at\"] >= \"2024-04-10\")\n",
    "].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "686ee6bc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:19:09.660433Z",
     "start_time": "2024-04-12T01:19:08.602512Z"
    }
   },
   "outputs": [],
   "source": [
    "# the csv conversion makes things...consistent...you will see\n",
    "new_clip_df.to_csv('/app/suno/data/dpo/data_prep/clip_df_v3_20240411.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c533820",
   "metadata": {},
   "source": [
    "# Concat new query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e473096",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:19:12.091198Z",
     "start_time": "2024-04-12T01:19:11.646335Z"
    }
   },
   "outputs": [],
   "source": [
    "new_clip_df = pd.read_csv('/app/suno/data/dpo/data_prep/clip_df_v3_20240411.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33eb0bdf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:19:30.586945Z",
     "start_time": "2024-04-12T01:19:12.971242Z"
    }
   },
   "outputs": [],
   "source": [
    "# this can still take forever\n",
    "clip_df_with_dup = pd.concat([prev_clip_df, new_clip_df], ignore_index=True)\n",
    "print(clip_df_with_dup.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4711d9a8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:21:33.320593Z",
     "start_time": "2024-04-12T01:19:30.588603Z"
    }
   },
   "outputs": [],
   "source": [
    "clip_df = clip_df_with_dup.drop_duplicates(subset=[\"id\"], keep='last')\n",
    "print(clip_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1243bd0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:26:28.578502Z",
     "start_time": "2024-04-12T01:26:28.211540Z"
    }
   },
   "outputs": [],
   "source": [
    "updated_info_dict[\"clip_df\"] = clip_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45df22db",
   "metadata": {},
   "source": [
    "# Feature Eng"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0026df7f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:26:48.931121Z",
     "start_time": "2024-04-12T01:26:48.915071Z"
    }
   },
   "outputs": [],
   "source": [
    "bots_action_df = updated_info_dict[\"bots_action_df\"]\n",
    "reaction_df = updated_info_dict[\"reaction_df\"]\n",
    "playlist_clip_df = updated_info_dict[\"playlist_clip_df\"]\n",
    "auth_user_df = updated_info_dict[\"auth_user_df\"]\n",
    "total_clip_df = updated_info_dict[\"clip_df\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8071a11e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:26:50.965922Z",
     "start_time": "2024-04-12T01:26:50.882595Z"
    }
   },
   "outputs": [],
   "source": [
    "# Get downvoted clips\n",
    "flagged_df = reaction_df[reaction_df[\"flagged\"]].copy()\n",
    "print(f\"Flagged: {flagged_df.shape[0]:,} rows\")\n",
    "flagged_ids = flagged_df[\"clip_id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06ba5589",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:28:17.724453Z",
     "start_time": "2024-04-12T01:28:17.294884Z"
    }
   },
   "outputs": [],
   "source": [
    "total_clip_df[\"metadata\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc9f3f75",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:37:53.393553Z",
     "start_time": "2024-04-12T01:37:52.198497Z"
    }
   },
   "outputs": [],
   "source": [
    "# new_clip_df['metadata'] = new_clip_df['metadata'].apply(ast.literal_eval)\n",
    "# test_df = pd.json_normalize(new_clip_df['metadata'])\n",
    "# df = df.join(test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d41b24ee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:43:14.385448Z",
     "start_time": "2024-04-12T01:42:10.911602Z"
    }
   },
   "outputs": [],
   "source": [
    "total_clip_df[\"is_in_playlist\"] = total_clip_df[\"id\"].isin(playlist_clip_df[\"clip_id\"].unique())\n",
    "# the nans are concats, we want to drop them for now\n",
    "concated_clips = total_clip_df[total_clip_df[\"request_id\"].isna()].copy()\n",
    "# need to kick them out...for now, since no pairs\n",
    "clip_df = total_clip_df[~total_clip_df[\"request_id\"].isna()]\n",
    "print(f\"concat clips frac = {concated_clips.shape[0] / total_clip_df.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24410718",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:27:27.537043Z",
     "start_time": "2024-04-12T01:27:27.537033Z"
    }
   },
   "outputs": [],
   "source": [
    "total_clip_df[\"continued_parent\"] = total_clip_df[\"metadata\"].apply(lambda x: parse_parent_id(x))\n",
    "clip_history_df = total_clip_df[~total_clip_df[\"continued_parent\"].isna()].copy()\n",
    "# these are the parent's ids\n",
    "has_continued_children_ids = clip_history_df[\"continued_parent\"]\n",
    "print(\n",
    "    \"clips that have children\",\n",
    "    len(has_continued_children_ids),\n",
    "    len(has_continued_children_ids.unique()),\n",
    "    \"\\n\",\n",
    "    \"Average continues from clip = \",\n",
    "    len(has_continued_children_ids) / len(has_continued_children_ids.unique()),\n",
    ")\n",
    "has_continued_children_ids = set(list(has_continued_children_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5de8c30",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:44:12.457918Z",
     "start_time": "2024-04-12T01:43:14.386998Z"
    }
   },
   "outputs": [],
   "source": [
    "# Filter out bad models\n",
    "print(f\"pre model filter {clip_df.shape}\")\n",
    "LIST_OF_MODELS = [\n",
    "    \"chirp-v3-engine-d\",\n",
    "    \"chirp-v3-engine-v0\",\n",
    "    \"chirp-v3-engine-i\",\n",
    "    \"chirp-v3-engine-i-d\",\n",
    "    \"chirp-v3-engine-s\",\n",
    "]\n",
    "clip_df = clip_df[clip_df[\"model_name\"].isin(LIST_OF_MODELS)].copy()\n",
    "print(f\"after model filter {clip_df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e6eb8e1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-12T01:58:51.022140Z",
     "start_time": "2024-04-12T01:56:43.771260Z"
    }
   },
   "outputs": [],
   "source": [
    "# Deal with concats\n",
    "concat_reaction_df = reaction_df[\n",
    "    reaction_df[\"clip_id\"].astype(str).isin(concated_clips[\"id\"])\n",
    "].copy()\n",
    "print(\n",
    "    \"Concat pieces reactions\",\n",
    "    concat_reaction_df.shape,\n",
    "    concat_reaction_df[\"clip_id\"].nunique(),\n",
    ")\n",
    "# bascially a grouped play counts\n",
    "concat_total_play_reaction_df_sum = concat_reaction_df.groupby(\"clip_id\")[\n",
    "    \"play_count\"\n",
    "].sum()\n",
    "concat_total_play_reaction_df_sum_df = (\n",
    "    concat_total_play_reaction_df_sum.reset_index().rename(\n",
    "        columns={\"clip_id\": \"id\", \"play_count\": \"reaction_play_count\"}\n",
    "    )\n",
    ")\n",
    "concated_clips = concated_clips.merge(\n",
    "    concat_total_play_reaction_df_sum_df, on=\"id\", how=\"left\"\n",
    ")\n",
    "print(f\"before play count filter: {concated_clips.shape}\")\n",
    "concated_clips = concated_clips[concated_clips[\"reaction_play_count\"] > 0]\n",
    "print(f\"after play count filter: {concated_clips.shape}\")\n",
    "# Because there are like with play count 0, it is better to filter them out\n",
    "concat_total_play_reaction_df_like = (\n",
    "    concat_reaction_df[concat_reaction_df[\"reaction_type\"] == \"L\"]\n",
    "    .groupby(\"clip_id\")[\"reaction_type\"]\n",
    "    .count()\n",
    ")\n",
    "concat_total_play_reaction_df_like_df = (\n",
    "    concat_total_play_reaction_df_like.reset_index().rename(\n",
    "        columns={\"clip_id\": \"id\", \"reaction_type\": \"reaction_like_count\"}\n",
    "    )\n",
    ")\n",
    "concated_clips = concated_clips.merge(\n",
    "    concat_total_play_reaction_df_like_df, on=\"id\", how=\"left\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cfd1d8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "concat_clips_ids = {}\n",
    "for _, row in tqdm.tqdm(concated_clips.iterrows()):\n",
    "    if history_ids := row[\"metadata\"][\"concat_history\"]:\n",
    "        total_duration = row[\"metadata\"][\"duration\"]\n",
    "        start_s = 0\n",
    "        for history_id in history_ids:\n",
    "            # print(history_ids, row[\"metadata\"][\"duration\"])\n",
    "            if isinstance(history_id, dict) and \"id\" in history_id:\n",
    "                # the other key is `continue_at`\n",
    "                concat_clips_ids[history_id[\"id\"]] = {\n",
    "                    \"total_start_s\": start_s,\n",
    "                    \"total_clip_s\": total_duration,\n",
    "                    \"concat_play_counts\": row[\"reaction_play_count\"],\n",
    "                    \"concat_in_playlist\": row[\"is_in_playlist\"],\n",
    "                    \"concat_likes\": row[\"upvote_count\"],\n",
    "                }\n",
    "                start_s = history_id[\"continue_at\"]\n",
    "print(\"total concat unique clips are:\", len(concat_clips_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9f522a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if a clip is generated by pro\n",
    "clip_df[\"is_pro_user\"] = clip_df[\"user_id\"].isin(auth_user_df[\"user_id\"].unique())\n",
    "print(\n",
    "    \"unique pro users generated clips\",\n",
    "    clip_df[\"user_id\"][clip_df[\"is_pro_user\"]].nunique(),\n",
    ")\n",
    "# set user number of clips generated\n",
    "clip_df[\"user_n_clips\"] = clip_df[\"user_id\"].map(clip_df[\"user_id\"].value_counts())"
   ]
  }
 ],
 "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.13"
  },
  "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": 5
}