yacup2023.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7d64683a-82dc-43cd-a222-21f72e9ca896",
   "metadata": {},
   "source": [
    "## Import Libraries and Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "340aceb6-4360-49ac-8381-bb42055d6146",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from glob import glob\n",
    "from datetime import datetime\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "32801790-6619-4141-b909-a70de5e00071",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GLOBAL VARIABLES\n",
    "CUDA_DEV = 0\n",
    "NUM_TAGS = 256\n",
    "\n",
    "# EDA Preprocessing\n",
    "MAX_GENRES = 10\n",
    "MAX_SECONDS = 138\n",
    "\n",
    "# Train Constants\n",
    "N_FOLDS = 10\n",
    "N_EPOCH = 14\n",
    "BATCH_SIZE = 64\n",
    "OPTIM_LR = 3e-4\n",
    "\n",
    "# Set True if you want to use W&B library\n",
    "IS_WANDB = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "29f35090-eb3d-464b-b46f-de30c4d42d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv('train_extended.csv')\n",
    "df_test = pd.read_csv('test_extended.csv')\n",
    "\n",
    "# Keep tracks with no more than 10 genres (by default)\n",
    "df_train = df_train[df_train.tags.apply(lambda x: len(x.split(',')) <= MAX_GENRES)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b46b7d4b-785e-4187-b2bc-6c620c7b8cf1",
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████| 76714/76714 [00:52<00:00, 1455.00it/s]\n"
     ]
    }
   ],
   "source": [
    "# Upload the track-embeddings (source - https://www.kaggle.com/datasets/sharthz23/yandex-cup-2023-recsys/data?select=track_embeddings)\n",
    "track_idx2embeds = {}\n",
    "for fn in tqdm(glob('track_embeddings/*')):\n",
    "    filename = os.path.basename(fn)\n",
    "    track_idx = int(os.path.splitext(filename)[0])\n",
    "    embeds = np.load(fn)\n",
    "    track_idx2embeds[track_idx] = embeds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8fdbb3e-368b-4c0b-9c26-b54eef54e5e9",
   "metadata": {},
   "source": [
    "### Create Track DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8d1b5cae-60ae-4584-a6bb-4f6b833929aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TaggingDataset(Dataset):\n",
    "    def __init__(self, df, testing=False):\n",
    "        self.df = df\n",
    "        self.testing = testing\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.df.shape[0]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        track_idx = row.track\n",
    "        embeds = track_idx2embeds[track_idx]\n",
    "        if self.testing:\n",
    "            return track_idx, embeds\n",
    "        tags = [int(x) for x in row.tags.split(',')]\n",
    "        target = np.zeros(NUM_TAGS)\n",
    "        target[tags] = 1\n",
    "        return embeds, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e4decde5-005a-4820-804c-cf4d110c799e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(b):\n",
    "    embeds = [torch.from_numpy(x[0]) for x in b]\n",
    "    targets = np.vstack([x[1] for x in b])\n",
    "    targets = torch.from_numpy(targets)\n",
    "    return embeds, targets\n",
    "\n",
    "def collate_fn_test(b):\n",
    "    track_idxs = torch.from_numpy(np.vstack([x[0] for x in b]))\n",
    "    embeds = [torch.from_numpy(x[1]) for x in b]\n",
    "    return track_idxs, embeds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb0bd4f0-e304-401b-9220-2f738775e53e",
   "metadata": {},
   "source": [
    "## EDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6d94f7fb-802e-47cb-aa51-115266e35daf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Confidence interval: 58.13 ± 3 * 26.62, so interval is [0; 138]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>length</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>50614.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>58.128502</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>26.615302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>10.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>43.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>53.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>68.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>404.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             length\n",
       "count  50614.000000\n",
       "mean      58.128502\n",
       "std       26.615302\n",
       "min       10.000000\n",
       "25%       43.000000\n",
       "50%       53.000000\n",
       "75%       68.000000\n",
       "max      404.000000"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAX0AAAD4CAYAAAAAczaOAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAVBklEQVR4nO3df6xc5X3n8fe3QABxsxhCeoVsa0221q4o3lK4C1SJqmtQwUBVU4lGrFBjIlaWdomUqlSL2SpLmoDWWYWyjdSm6w1enKTtDaWNsEiyWa/xVZQ/+BEXgg2UchOcLVesrcbG7U0ou06/+8c8l0xv7o+547lnjv28X9JoznnOMzOfeez7nTPPnDkTmYkkqQ4/NewAkqTmWPQlqSIWfUmqiEVfkipi0Zekipw57ACLueiii3LdunULbv/BD37Aeeed11ygHrU1F5itH23NBWbrR1tzweCy7d+//28y873zbszM1l6uvPLKXMy+ffsW3T4sbc2VabZ+tDVXptn60dZcmYPLBnwrF6irTu9IUkV6KvoRcSgiDkTE8xHxrdJ2YUTsiYhXy/UFpT0i4jMRMRURL0TEFV33s6X0fzUitqzMU5IkLWQ5e/obM/PyzBwr69uAvZm5Hthb1gFuBNaXy1bgs9B5kQDuA64GrgLum32hkCQ142SmdzYDu8ryLuCWrvbPl6mlp4BVEXExcAOwJzOPZuYxYA+w6SQeX5K0TJE9nHsnIl4DjgEJ/NfM3BERb2bmqrI9gGOZuSoingC2Z+Y3y7a9wD3AOHBOZt5f2j8GvJWZn57zWFvpvENgdHT0yomJiQVzzczMMDIysrxn3IC25gKz9aOtucBs/WhrLhhcto0bN+7vmpX5R3o9ZPMDmTkdET8N7ImIv+zemJkZEQM5c1tm7gB2AIyNjeX4+PiCfScnJ1ls+7C0NReYrR9tzQVm60dbc0Ez2Xqa3snM6XJ9BPgynTn5w2XahnJ9pHSfBtZ23XxNaVuoXZLUkCWLfkScFxHvnl0GrgcOAruB2SNwtgCPl+XdwIfKUTzXAMcz8w3g68D1EXFB+QD3+tImSWpIL9M7o8CXO9P2nAn8cWb+j4h4Fng0Iu4Evgd8sPT/KnATMAX8EPgwQGYejYhPAs+Wfp/IzKMDeyaSpCUtWfQz87vAz83T/n3gunnaE7hrgfvaCexcfsz2W7ftK+8sP7KpnV/xliS/kbsCDkwfZ922r/yjFwJJagOLviRVxKIvSRWx6EtSRSz6klQRi74kVcSiL0kVsehLUkUs+pJUEYu+JFXEoi9JFbHoS1JFLPqSVBGLviRVxKIvSRWx6EtSRSz6klQRi74kVcSiL0kVsehLUkWW/GF0LayX38Dt7nNo+80rGUeSluSeviRVxKIvSRWx6EtSRSz6klQRi74kVcSiL0kVsehLUkUs+pJUEYu+JFXEoi9JFbHoS1JFLPqSVJGei35EnBERz0XEE2X9koh4OiKmIuJLEfGu0n52WZ8q29d13ce9pf2ViLhh4M9GkrSo5ezpfxR4uWv9U8BDmfkzwDHgztJ+J3CstD9U+hERlwK3AT8LbAL+ICLOOLn4kqTl6KnoR8Qa4Gbgc2U9gGuBx0qXXcAtZXlzWadsv6703wxMZObbmfkaMAVcNYDnIEnqUWTm0p0iHgP+E/Bu4LeAO4Cnyt48EbEW+FpmXhYRB4FNmfl62fYd4Grg4+U2XyztD5fbPDbnsbYCWwFGR0evnJiYWDDXzMwMIyMjy3m+A3Vg+vi87aPnwuG3frJ9w+rzVzjR0oY9Zotpa7a25gKz9aOtuWBw2TZu3Lg/M8fm27bkj6hExC8DRzJzf0SMn3SaJWTmDmAHwNjYWI6PL/yQk5OTLLZ9pd2xwI+o3L3hBA8e+MmhPXT7+AonWtqwx2wxbc3W1lxgtn60NRc0k62XX856P/ArEXETcA7wT4DfA1ZFxJmZeQJYA0yX/tPAWuD1iDgTOB/4flf7rO7bSJIasOScfmbem5lrMnMdnQ9in8zM24F9wK2l2xbg8bK8u6xTtj+ZnTmk3cBt5eieS4D1wDMDeyaSpCWdzG/k3gNMRMT9wHPAw6X9YeALETEFHKXzQkFmvhgRjwIvASeAuzLzRyfx+JKkZVpW0c/MSWCyLH+XeY6+ycy/B35tgds/ADyw3JCSpMHwG7mSVBGLviRVxKIvSRWx6EtSRSz6klQRi74kVcSiL0kVsehLUkUs+pJUEYu+JFXkZM69o2Va13Uq5kPbbx5iEkm1ck9fkipi0Zekilj0JakiFn1JqohFX5IqYtGXpIpY9CWpIhZ9SaqIRV+SKuI3cpep+1u1knSqcU9fkipi0Zekilj0JakiFn1JqohFX5IqYtGXpIpY9CWpIhZ9SaqIRV+SKmLRl6SKWPQlqSIWfUmqyJJFPyLOiYhnIuLbEfFiRPxOab8kIp6OiKmI+FJEvKu0n13Wp8r2dV33dW9pfyUiblixZyVJmlcve/pvA9dm5s8BlwObIuIa4FPAQ5n5M8Ax4M7S/07gWGl/qPQjIi4FbgN+FtgE/EFEnDHA5yJJWsKSRT87ZsrqWeWSwLXAY6V9F3BLWd5c1inbr4uIKO0Tmfl2Zr4GTAFXDeJJSJJ609OcfkScERHPA0eAPcB3gDcz80Tp8jqwuiyvBv4aoGw/Drynu32e20iSGhCZ2XvniFXAl4GPAY+UKRwiYi3wtcy8LCIOApsy8/Wy7TvA1cDHgacy84ul/eFym8fmPMZWYCvA6OjolRMTEwvmmZmZYWRkpOf8g3Bg+viSfUbPhcNvLd5nw+rzB5RoeYYxZr1qa7a25gKz9aOtuWBw2TZu3Lg/M8fm27asX87KzDcjYh/wC8CqiDiz7M2vAaZLt2lgLfB6RJwJnA98v6t9Vvdtuh9jB7ADYGxsLMfHxxfMMzk5yWLbV8IdPfxy1t0bTvDggcWH9tDt4wNKtDzDGLNetTVbW3OB2frR1lzQTLZejt55b9nDJyLOBX4JeBnYB9xaum0BHi/Lu8s6ZfuT2Xk7sRu4rRzdcwmwHnhmQM/jlLNu21feuUhSU3rZ078Y2FWOtPkp4NHMfCIiXgImIuJ+4Dng4dL/YeALETEFHKVzxA6Z+WJEPAq8BJwA7srMHw326UiSFrNk0c/MF4Cfn6f9u8xz9E1m/j3wawvc1wPAA8uPKUkaBL+RK0kVsehLUkUs+pJUEYu+JFXEoi9JFbHoS1JFLPqSVBGLviRVxKIvSRWx6EtSRZZ1ls1aeVI0SacL9/QlqSIWfUmqiEVfkipi0Zekilj0JakiFn1JqohFX5IqYtGXpIpY9CWpIn4jtwW6v/F7aPvNQ0wi6XTnnr4kVcSiL0kVsehLUkUs+pJUEYu+JFXEoi9JFbHoS1JFLPqSVBGLviRVxKIvSRWx6EtSRSz6klQRi74kVWTJoh8RayNiX0S8FBEvRsRHS/uFEbEnIl4t1xeU9oiIz0TEVES8EBFXdN3XltL/1YjYsnJPS5I0n1729E8Ad2fmpcA1wF0RcSmwDdibmeuBvWUd4EZgfblsBT4LnRcJ4D7gauAq4L7ZFwpJUjOWLPqZ+UZm/kVZ/jvgZWA1sBnYVbrtAm4py5uBz2fHU8CqiLgYuAHYk5lHM/MYsAfYNMgnI0laXGRm750j1gHfAC4D/ndmrirtARzLzFUR8QSwPTO/WbbtBe4BxoFzMvP+0v4x4K3M/PScx9hK5x0Co6OjV05MTCyYZ2ZmhpGRkZ7z9+vA9PFl9R89Fw6/1d9jbVh9fn837FFTY9aPtmZray4wWz/amgsGl23jxo37M3Nsvm09/3JWRIwAfwb8Rmb+bafOd2RmRkTvrx6LyMwdwA6AsbGxHB8fX7Dv5OQki20flDu6ftmqF3dvOMGDB/r7UbJDt4/3dbteNTVm/WhrtrbmArP1o625oJlsPR29ExFn0Sn4f5SZf16aD5dpG8r1kdI+Daztuvma0rZQuySpIb0cvRPAw8DLmfm7XZt2A7NH4GwBHu9q/1A5iuca4HhmvgF8Hbg+Ii4oH+BeX9okSQ3pZQ7i/cCvAwci4vnS9h+A7cCjEXEn8D3gg2XbV4GbgCngh8CHATLzaER8Eni29PtEZh4dxJOQJPVmyaJfPpCNBTZfN0//BO5a4L52AjuXE7A267o+Pzi0/eYhJpF0OvIbuZJUEYu+JFWkv+MKK7BumYdpStKpwD19SaqIRV+SKmLRl6SKWPQlqSIWfUmqiEVfkipi0Zekilj0JakiFn1JqohFX5IqYtGXpIpY9CWpIhZ9SaqIRV+SKmLRl6SKWPQlqSIWfUmqiEVfkirizyWeIrp/vvHQ9puHmETSqcw9fUmqiEVfkipi0Zekilj0JakifpDbpfvDUkk6HbmnL0kVcU//FOThm5L6ZdFvMaebJA2a0zuSVBGLviRVxKIvSRVZsuhHxM6IOBIRB7vaLoyIPRHxarm+oLRHRHwmIqYi4oWIuKLrNltK/1cjYsvKPB1J0mJ62dN/BNg0p20bsDcz1wN7yzrAjcD6ctkKfBY6LxLAfcDVwFXAfbMvFJKk5ixZ9DPzG8DROc2bgV1leRdwS1f757PjKWBVRFwM3ADsycyjmXkM2MNPvpBIklZYZObSnSLWAU9k5mVl/c3MXFWWAziWmasi4glge2Z+s2zbC9wDjAPnZOb9pf1jwFuZ+el5HmsrnXcJjI6OXjkxMbFgrpmZGUZGRnp+sks5MH18IPczei4cfmsgd7WkDavPX1b/QY/ZILU1W1tzgdn60dZcMLhsGzdu3J+ZY/NtO+nj9DMzI2LpV47e728HsANgbGwsx8fHF+w7OTnJYtuX644BHRd/94YTPHigoa9AHPjBO4u9fFFr0GM2SG3N1tZcYLZ+tDUXNJOt36N3DpdpG8r1kdI+Dazt6remtC3ULklqUL9FfzcwewTOFuDxrvYPlaN4rgGOZ+YbwNeB6yPigvIB7vWlTZLUoCXnICLiT+jMyV8UEa/TOQpnO/BoRNwJfA/4YOn+VeAmYAr4IfBhgMw8GhGfBJ4t/T6RmXM/HJYkrbAli35m/usFNl03T98E7lrgfnYCO5eVTpI0UNWfcM2TmkmqiadhkKSKWPQlqSIWfUmqSPVz+qcTf1FL0lLc05ekirinf5qae1SSe/6SwD19SaqKRV+SKlLl9I5fyJJUK/f0JakiVe7p12j23c3dG04wPtwokobIPX1JqohFX5IqUs30jh/eSlJFRV8/5ukapHo5vSNJFbHoS1JFLPqSVBGLviRVxA9yK+eHulJd3NOXpIqc1nv6Hpu/PO71S6c/9/QlqSIWfUmqiEVfkipyWs/pq3/O70unJ/f0Jaki7ulrSe71S6cPi76WxRcA6dTm9I4kVcQ9ffVtoS+/+Q5Aai/39CWpIu7pa+B8ByC1V+NFPyI2Ab8HnAF8LjO3N51Bw7HQh8Cz7XdvOMF406GkyjRa9CPiDOD3gV8CXgeejYjdmflSkzk0fAu9GziZk+T5TkJaWtN7+lcBU5n5XYCImAA2AxZ9nbSVOqvq3RtOcEe57/neofSqlxelNh4S28ZM6l9kZnMPFnErsCkz/01Z/3Xg6sz8SFefrcDWsvrPgVcWucuLgL9Zobgno625wGz9aGsuMFs/2poLBpftn2bme+fb0LoPcjNzB7Cjl74R8a3MHFvhSMvW1lxgtn60NReYrR9tzQXNZGv6kM1pYG3X+prSJklqQNNF/1lgfURcEhHvAm4DdjecQZKq1ej0TmaeiIiPAF+nc8jmzsx88STusqdpoCFoay4wWz/amgvM1o+25oIGsjX6Qa4kabg8DYMkVcSiL0kVOSWLfkRsiohXImIqIra1IM+hiDgQEc9HxLdK24URsSciXi3XFzSUZWdEHImIg11t82aJjs+UcXwhIq5oONfHI2K6jNvzEXFT17Z7S65XIuKGlcpVHmttROyLiJci4sWI+GhpH+q4LZJr6OMWEedExDMR8e2S7XdK+yUR8XTJ8KVywAYRcXZZnyrb1zWc65GIeK1rzC4v7Y39DXRlPCMinouIJ8p6s2OWmafUhc4HwN8B3ge8C/g2cOmQMx0CLprT9p+BbWV5G/CphrL8InAFcHCpLMBNwNeAAK4Bnm4418eB35qn76Xl3/Vs4JLy733GCma7GLiiLL8b+KuSYajjtkiuoY9bee4jZfks4OkyFo8Ct5X2PwT+bVn+d8AfluXbgC81nOsR4NZ5+jf2N9D1mL8J/DHwRFlvdMxOxT39d07lkJn/F5g9lUPbbAZ2leVdwC1NPGhmfgM42mOWzcDns+MpYFVEXNxgroVsBiYy8+3MfA2YovPvviIy843M/Iuy/HfAy8Bqhjxui+RaSGPjVp77TFk9q1wSuBZ4rLTPHbPZsXwMuC4iosFcC2nsbwAgItYANwOfK+tBw2N2Khb91cBfd62/zuJ/CE1I4H9GxP7onEYCYDQz3yjL/wcYHU60RbO0YSw/Ut5W7+yaAhtarvIW+ufp7CG2Ztzm5IIWjFuZpngeOALsofPO4s3MPDHP47+TrWw/DryniVyZOTtmD5Qxeygizp6ba57MK+G/AP8e+Iey/h4aHrNTsei30Qcy8wrgRuCuiPjF7o3ZeX/WimNj25QF+Czwz4DLgTeAB4cZJiJGgD8DfiMz/7Z72zDHbZ5crRi3zPxRZl5O55v1VwH/Yhg55pqbKyIuA+6lk+9fARcC9zSdKyJ+GTiSmfubfuxup2LRb92pHDJzulwfAb5M5w/g8OzbxHJ9ZHgJF8wy1LHMzMPlD/QfgP/Gj6ciGs8VEWfRKax/lJl/XpqHPm7z5WrTuJU8bwL7gF+gMz0y+6XP7sd/J1vZfj7w/YZybSpTZZmZbwP/neGM2fuBX4mIQ3Smpa+l89sijY7ZqVj0W3Uqh4g4LyLePbsMXA8cLJm2lG5bgMeHkxAWybIb+FA5guEa4HjXdMaKmzN3+qt0xm02123l6IVLgPXAMyuYI4CHgZcz83e7Ng113BbK1YZxi4j3RsSqsnwund/IeJlOkb21dJs7ZrNjeSvwZHn31ESuv+x68Q46c+bdY9bI30Bm3puZazJzHZ269WRm3k7TYzaIT4ObvtD5xP2v6Mwh/vaQs7yPzhET3wZenM1DZ+5tL/Aq8L+ACxvK8yd03vL/Pzrzg3culIXOEQu/X8bxADDWcK4vlMd9ofwHv7ir/2+XXK8AN67wmH2AztTNC8Dz5XLTsMdtkVxDHzfgXwLPlQwHgf/Y9ffwDJ0Pkf8UOLu0n1PWp8r29zWc68kyZgeBL/LjI3wa+xuYk3OcHx+90+iYeRoGSarIqTi9I0nqk0Vfkipi0Zekilj0JakiFn1JqohFX5IqYtGXpIr8f+lBuU9D5mHEAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_train.length.hist(bins=100)\n",
    "\n",
    "length_std = df_train.length.std()\n",
    "length_mean = df_train.length.mean()\n",
    "print(f'Confidence interval: {length_mean:.2f} ± 3 * {length_std:.2f}, so interval is [0; {round(length_mean + 3 * length_std)}]')\n",
    "\n",
    "df_train[['length']].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8accf595-2aaf-4cbf-b29d-b7911e4b2fc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Confidence interval: 3.84 ± 3 * 1.63, so interval is [1; 9]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>n_genres</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>50614.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>3.843759</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.625611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>3.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>4.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>5.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>10.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           n_genres\n",
       "count  50614.000000\n",
       "mean       3.843759\n",
       "std        1.625611\n",
       "min        1.000000\n",
       "25%        3.000000\n",
       "50%        4.000000\n",
       "75%        5.000000\n",
       "max       10.000000"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAD4CAYAAAAO9oqkAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAU4ElEQVR4nO3db4xd9X3n8fdncUmANjGE7ixrWzvWxqIi8WZDR0AXqRqFLhiIYh6kEREbTNZbPyhJaddSaroPkJKyItpSGrQNK29wMS2CsDQrrEJDLJKraKVAgCTlb1hmwQF7TUhjIJ1km3Sy330wP6/vDGNsz/03g98v6WrO+Z7fOed3f9zhM+fPPU5VIUk6vv2jUXdAkjR6hoEkyTCQJBkGkiQMA0kSsGLUHVis008/vcbHx0fdjZ78+Mc/5pRTThl1N5YEx2Iux2Mux+OQXsfiscce+9uq+uX59WUbBuPj4zz66KOj7kZPOp0Ok5OTo+7GkuBYzOV4zOV4HNLrWCT53kJ1TxNJkgwDSZJhIEnCMJAkYRhIkjAMJEkYBpIkjiIMkuxI8kqSJxdYtjVJJTm9zSfJzUmmkjye5OyutpuSPNdem7rqv5rkibbOzUnSrzcnSTo6R3NkcBuwYX4xyRrgQuDFrvLFwLr22gLc0tqeBlwHnAucA1yX5NS2zi3Ab3Wt94Z9SZIG64jfQK6qrycZX2DRTcCngHu7ahuB22v2X8x5KMnKJGcAk8DuqjoAkGQ3sCFJB3hHVT3U6rcDlwF/vdg3pIWNb7tvJPvdc8OlI9mvpGOzqMdRJNkI7Kuqv5l3VmcV8FLX/N5We7P63gXqh9vvFmaPOBgbG6PT6Sym+0vG9PT00N7D1vUzQ9nPfEf7/oY5FsuB4zGX43HIoMbimMMgycnAHzB7imioqmo7sB1gYmKilvuzSob5vJWrRnVkcMXkUbXz2TNzOR5zOR6HDGosFnM30T8H1gJ/k2QPsBr4VpJ/AuwD1nS1Xd1qb1ZfvUBdkjRExxwGVfVEVf3jqhqvqnFmT+2cXVUvA7uAK9tdRecBr1fVfuAB4MIkp7YLxxcCD7RlP0pyXruL6ErmXoOQJA3B0dxaeifwDeDMJHuTbH6T5vcDzwNTwH8FfhugXTj+DPBIe3364MXk1uYLbZ3/hRePJWnojuZuoo8eYfl413QBVx+m3Q5gxwL1R4H3HqkfkqTB8RvIkiTDQJJkGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJHEUYJNmR5JUkT3bV/lOS7yZ5PMl/T7Kya9m1SaaSPJvkoq76hlabSrKtq742ycOt/sUkJ/bx/UmSjsLRHBncBmyYV9sNvLeq/gXwP4FrAZKcBVwOvKet8/kkJyQ5AfhT4GLgLOCjrS3AZ4GbqurdwKvA5p7ekSTpmB0xDKrq68CBebWvVNVMm30IWN2mNwJ3VdVPq+oFYAo4p72mqur5qvoZcBewMUmADwD3tPV3Apf19pYkScdqRR+28W+BL7bpVcyGw0F7Ww3gpXn1c4F3Aa91BUt3+zdIsgXYAjA2Nkan0+m17yM1PT09tPewdf3MkRsNwNG+v2GOxXLgeMzleBwyqLHoKQyS/AdgBrijP915c1W1HdgOMDExUZOTk8PY7cB0Oh2G9R6u2nbfUPYz354rJo+q3TDHYjlwPOZyPA4Z1FgsOgySXAV8ELigqqqV9wFrupqtbjUOU/8hsDLJinZ00N1ekjQki7q1NMkG4FPAh6rqJ12LdgGXJ3lbkrXAOuCbwCPAunbn0InMXmTe1ULka8CH2/qbgHsX91YkSYt1NLeW3gl8Azgzyd4km4H/DPwSsDvJd5L8F4Cqegq4G3ga+DJwdVX9vP3V/wngAeAZ4O7WFuD3gX+fZIrZawi39vUdSpKO6IiniarqowuUD/s/7Kq6Hrh+gfr9wP0L1J9n9m4jSdKI9ONuIumwxo/ywvXW9TN9v8i954ZL+7o96a3Mx1FIkgwDSZJhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwmcTDdX85/QM4nk8krQYHhlIkgwDSZJhIEnCMJAkYRhIkjAMJEkYBpIkjiIMkuxI8kqSJ7tqpyXZneS59vPUVk+Sm5NMJXk8ydld62xq7Z9Lsqmr/qtJnmjr3Jwk/X6TkqQ3dzRHBrcBG+bVtgEPVtU64ME2D3AxsK69tgC3wGx4ANcB5wLnANcdDJDW5re61pu/L0nSgB0xDKrq68CBeeWNwM42vRO4rKt+e816CFiZ5AzgImB3VR2oqleB3cCGtuwdVfVQVRVwe9e2JElDsthrBmNVtb9NvwyMtelVwEtd7fa22pvV9y5QlyQNUc/PJqqqSlL96MyRJNnC7OknxsbG6HQ6w9ht32xdPzNnfuykN9aOV4MYi+X2+eg2PT29rPvfb47HIYMai8WGwfeTnFFV+9upnldafR+wpqvd6lbbB0zOq3daffUC7RdUVduB7QATExM1OTl5uKZL0vyH0m1dP8ONT/isQBjMWOy5YrKv2xumTqfDcvt8D5LjccigxmKxp4l2AQfvCNoE3NtVv7LdVXQe8Ho7nfQAcGGSU9uF4wuBB9qyHyU5r91FdGXXtiRJQ3LEP8WS3MnsX/WnJ9nL7F1BNwB3J9kMfA/4SGt+P3AJMAX8BPg4QFUdSPIZ4JHW7tNVdfCi9G8ze8fSScBft5ckaYiOGAZV9dHDLLpggbYFXH2Y7ewAdixQfxR475H6IUkaHL+BLEkyDCRJhoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgSaLHMEjye0meSvJkkjuTvD3J2iQPJ5lK8sUkJ7a2b2vzU235eNd2rm31Z5Nc1ON7kiQdo0WHQZJVwO8AE1X1XuAE4HLgs8BNVfVu4FVgc1tlM/Bqq9/U2pHkrLbee4ANwOeTnLDYfkmSjl2vp4lWACclWQGcDOwHPgDc05bvBC5r0xvbPG35BUnS6ndV1U+r6gVgCjinx35Jko7BisWuWFX7kvwR8CLwf4CvAI8Br1XVTGu2F1jVplcBL7V1Z5K8Dryr1R/q2nT3OnMk2QJsARgbG6PT6Sy2+yOxdf3MnPmxk95YO14NYiyW2+ej2/T09LLuf785HocMaiwWHQZJTmX2r/q1wGvAf2P2NM/AVNV2YDvAxMRETU5ODnJ3fXfVtvvmzG9dP8ONTyz6P8FbyiDGYs8Vk33d3jB1Oh2W2+d7kByPQwY1Fr2cJvoN4IWq+kFV/QPwJeB8YGU7bQSwGtjXpvcBawDa8ncCP+yuL7COJGkIegmDF4Hzkpzczv1fADwNfA34cGuzCbi3Te9q87TlX62qavXL291Ga4F1wDd76Jck6Rj1cs3g4ST3AN8CZoBvM3sK5z7griR/2Gq3tlVuBf48yRRwgNk7iKiqp5LczWyQzABXV9XPF9svSdKx6+kkbVVdB1w3r/w8C9wNVFV/D/zmYbZzPXB9L32RJC2e30CWJBkGkiTDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgSaLHB9VJS9n4vH9MaFj23HDpSPYr9cIjA0mSYSBJMgwkSRgGkiQMA0kShoEkCcNAkkSPYZBkZZJ7knw3yTNJfi3JaUl2J3mu/Ty1tU2Sm5NMJXk8ydld29nU2j+XZFOvb0qSdGx6PTL4HPDlqvoV4H3AM8A24MGqWgc82OYBLgbWtdcW4BaAJKcB1wHnAucA1x0MEEnScCw6DJK8E/h14FaAqvpZVb0GbAR2tmY7gcva9Ebg9pr1ELAyyRnARcDuqjpQVa8Cu4ENi+2XJOnY9fI4irXAD4A/S/I+4DHgGmCsqva3Ni8DY216FfBS1/p7W+1w9TdIsoXZowrGxsbodDo9dH/4tq6fmTM/dtIba8ert9JY9ONzOT09vew+34PkeBwyqLHoJQxWAGcDn6yqh5N8jkOnhACoqkpSvXRw3va2A9sBJiYmanJysl+bHoqr5j0rZ+v6GW58wsdDwVtrLPZcMdnzNjqdDsvt8z1IjschgxqLXq4Z7AX2VtXDbf4eZsPh++30D+3nK235PmBN1/qrW+1wdUnSkCw6DKrqZeClJGe20gXA08Au4OAdQZuAe9v0LuDKdlfRecDr7XTSA8CFSU5tF44vbDVJ0pD0elz+SeCOJCcCzwMfZzZg7k6yGfge8JHW9n7gEmAK+ElrS1UdSPIZ4JHW7tNVdaDHfkmSjkFPYVBV3wEmFlh0wQJtC7j6MNvZAezopS+SpMXzG8iSJMNAkmQYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIk+hAGSU5I8u0kf9Xm1yZ5OMlUki8mObHV39bmp9ry8a5tXNvqzya5qNc+SZKOTT+ODK4Bnuma/yxwU1W9G3gV2Nzqm4FXW/2m1o4kZwGXA+8BNgCfT3JCH/olSTpKPYVBktXApcAX2nyADwD3tCY7gcva9MY2T1t+QWu/Ebirqn5aVS8AU8A5vfRLknRsVvS4/p8AnwJ+qc2/C3itqmba/F5gVZteBbwEUFUzSV5v7VcBD3Vts3udOZJsAbYAjI2N0el0euz+cG1dPzNnfuykN9aOV2+lsejH53J6enrZfb4HyfE4ZFBjsegwSPJB4JWqeizJZN969CaqajuwHWBiYqImJ4ey2765att9c+a3rp/hxid6zeO3hrfSWOy5YrLnbXQ6HZbb53uQHI9DBjUWvfz2nQ98KMklwNuBdwCfA1YmWdGODlYD+1r7fcAaYG+SFcA7gR921Q/qXkeSNASLvmZQVddW1eqqGmf2AvBXq+oK4GvAh1uzTcC9bXpXm6ct/2pVVatf3u42WgusA7652H5Jko7dII7Lfx+4K8kfAt8Gbm31W4E/TzIFHGA2QKiqp5LcDTwNzABXV9XPB9AvSdJh9CUMqqoDdNr08yxwN1BV/T3wm4dZ/3rg+n70RRq18XnXhhZj6/qZN1xjOhp7bri0533r+OQ3kCVJhoEkyTCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkughDJKsSfK1JE8neSrJNa1+WpLdSZ5rP09t9SS5OclUkseTnN21rU2t/XNJNvX+tiRJx6KXI4MZYGtVnQWcB1yd5CxgG/BgVa0DHmzzABcD69prC3ALzIYHcB1wLnAOcN3BAJEkDceiw6Cq9lfVt9r03wHPAKuAjcDO1mwncFmb3gjcXrMeAlYmOQO4CNhdVQeq6lVgN7Bhsf2SJB27Ff3YSJJx4P3Aw8BYVe1vi14Gxtr0KuClrtX2ttrh6gvtZwuzRxWMjY3R6XT60f2h2bp+Zs782ElvrB2vHIu5Fjsey+134mhNT0+/Zd/bsRrUWPQcBkl+EfhL4Her6kdJ/v+yqqok1es+ura3HdgOMDExUZOTk/3a9FBcte2+OfNb189w4xN9yeNlz7GYa7HjseeKyf53ZgnodDost9/3QRnUWPT025fkF5gNgjuq6kut/P0kZ1TV/nYa6JVW3wes6Vp9davtAybn1Tu99Es6Xo3P+4NjWPbccOlI9qv+6eVuogC3As9U1R93LdoFHLwjaBNwb1f9ynZX0XnA6+100gPAhUlObReOL2w1SdKQ9HJkcD7wMeCJJN9ptT8AbgDuTrIZ+B7wkbbsfuASYAr4CfBxgKo6kOQzwCOt3aer6kAP/ZIkHaNFh0FV/Q8gh1l8wQLtC7j6MNvaAexYbF8kSb3xG8iSJMNAkmQYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCTRp3/PQNLxbdBPS926fuYNj4AHn5baTx4ZSJIMA0mSYSBJwjCQJGEYSJIwDCRJeGuppGVs0Le0vpm32m2tHhlIkjwykKTFGNVRyW0bThnIdj0ykCQtnTBIsiHJs0mmkmwbdX8k6XiyJMIgyQnAnwIXA2cBH01y1mh7JUnHj6VyzeAcYKqqngdIchewEXh6EDsb5R0IkrQUpapG3QeSfBjYUFX/rs1/DDi3qj4xr90WYEubPRN4dqgd7b/Tgb8ddSeWCMdiLsdjLsfjkF7H4p9V1S/PLy6VI4OjUlXbge2j7ke/JHm0qiZG3Y+lwLGYy/GYy/E4ZFBjsSSuGQD7gDVd86tbTZI0BEslDB4B1iVZm+RE4HJg14j7JEnHjSVxmqiqZpJ8AngAOAHYUVVPjbhbw/CWOeXVB47FXI7HXI7HIQMZiyVxAVmSNFpL5TSRJGmEDANJkmEwbEnWJPlakqeTPJXkmlH3aSlIckKSbyf5q1H3ZdSSrExyT5LvJnkmya+Nuk+jkuT32u/Jk0nuTPL2UfdpmJLsSPJKkie7aqcl2Z3kufbz1H7syzAYvhlga1WdBZwHXO2jNwC4Bnhm1J1YIj4HfLmqfgV4H8fpuCRZBfwOMFFV72X25pLLR9urobsN2DCvtg14sKrWAQ+2+Z4ZBkNWVfur6ltt+u+Y/UVfNdpejVaS1cClwBdG3ZdRS/JO4NeBWwGq6mdV9dpIOzVaK4CTkqwATgb+94j7M1RV9XXgwLzyRmBnm94JXNaPfRkGI5RkHHg/8PCIuzJqfwJ8Cvi/I+7HUrAW+AHwZ+202ReSDOYB9ktcVe0D/gh4EdgPvF5VXxltr5aEsara36ZfBsb6sVHDYESS/CLwl8DvVtWPRt2fUUnyQeCVqnps1H1ZIlYAZwO3VNX7gR/Tp9MAy007F76R2YD8p8ApSf7NaHu1tNTsdwP68v0Aw2AEkvwCs0FwR1V9adT9GbHzgQ8l2QPcBXwgyV+MtksjtRfYW1UHjxbvYTYcjke/AbxQVT+oqn8AvgT8qxH3aSn4fpIzANrPV/qxUcNgyJKE2fPBz1TVH4+6P6NWVddW1eqqGmf24uBXq+q4/euvql4GXkpyZitdwIAe5b4MvAicl+Tk9ntzAcfpxfR5dgGb2vQm4N5+bNQwGL7zgY8x+xfwd9rrklF3SkvKJ4E7kjwO/EvgP462O6PRjo7uAb4FPMHs/6+Oq8dSJLkT+AZwZpK9STYDNwD/OslzzB493dCXffk4CkmSRwaSJMNAkmQYSJIwDCRJGAaSJAwDSRKGgSQJ+H/TlYjYFRMpKwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_train['n_genres'] = df_train.tags.apply(lambda x: len(x.split(',')))\n",
    "df_train['n_genres'].hist(bins=10)\n",
    "\n",
    "length_std = df_train['n_genres'].std()\n",
    "length_mean = df_train['n_genres'].mean()\n",
    "print(f'Confidence interval: {length_mean:.2f} ± 3 * {length_std:.2f}, so interval is [1; {round(length_mean + 3 * length_std)}]')\n",
    "\n",
    "df_train[['n_genres']].describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd0534c8-d63e-421d-b4b0-439ab9c9c069",
   "metadata": {},
   "source": [
    "## Model Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd9f366b-f6d7-4371-89b1-2c05c9a83bd9",
   "metadata": {},
   "source": [
    "### Create Transformer Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c3ed15ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "# Custom Model for Genre Predictions\n",
    "class Trackformer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_classes=NUM_TAGS,\n",
    "        input_dim=768,\n",
    "        hidden_dim=128,\n",
    "        n_encoder_layers=1,\n",
    "        n_decoder_layers=6,\n",
    "        n_heads=4,\n",
    "        dropout=0.1\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.projector = nn.Linear(input_dim, hidden_dim)\n",
    "\n",
    "        self.transformer = nn.Transformer(\n",
    "            d_model=hidden_dim,\n",
    "            nhead=n_heads,\n",
    "            num_encoder_layers=n_encoder_layers,\n",
    "            num_decoder_layers=n_decoder_layers,\n",
    "            dropout=dropout,\n",
    "            batch_first=True\n",
    "        )\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.fc = nn.Linear(hidden_dim, num_classes)\n",
    "\n",
    "        self.n_params = sum(p.numel() for p in self.parameters())\n",
    "\n",
    "    def forward(self, embeds):\n",
    "        # Counting only the first 138 seconds of each track (by default)\n",
    "        x = [self.projector(x[:MAX_SECONDS]) for x in embeds]\n",
    "        x = pad_sequence(x, batch_first=True, padding_value=0) # add zero-padding\n",
    "        x = self.transformer(x, x[:, 0, :].unsqueeze(1)).squeeze(1)\n",
    "        out = self.fc(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "937aac94-ed6e-4f10-be00-26a8e506d598",
   "metadata": {},
   "source": [
    "### Add Train/Evaluation/Predict Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c70bf034-7966-4f44-9f2e-dcaf0f8a8184",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(model, train_loader, val_loader, criterion, optimizer):\n",
    "    model.train()\n",
    "\n",
    "    for iteration, data in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        embeds, target = data\n",
    "        embeds = [x.to(CUDA_DEV) for x in embeds]\n",
    "        target = target.to(CUDA_DEV).float()\n",
    "        pred_logits = model(embeds).sigmoid().float()\n",
    "        ce_loss = criterion(pred_logits, target)\n",
    "        ce_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if iteration % 100 == 0:\n",
    "            train_loss = ce_loss.item()\n",
    "            val_loss, val_acc = evaluate_model(model, val_loader)\n",
    "            if IS_WANDB:\n",
    "                wandb.log({\"batch\": iteration + 1, \"train_loss\": train_loss, \"val_loss\": val_loss, \"val_acc\": val_acc})\n",
    "            print(f'   {datetime.now()} batch {iteration + 1} train_loss {train_loss:.5f} val_loss {val_loss:.5f} val_acc {val_acc:.5f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3dd84262",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "from sklearn.metrics import average_precision_score\n",
    "\n",
    "def evaluate_model(model, val_loader):\n",
    "    warnings.filterwarnings(\"ignore\")\n",
    "    model.eval()\n",
    "\n",
    "    val_acc = 0.0\n",
    "    val_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        for iteration, data in enumerate(val_loader):\n",
    "            embeds, target = data\n",
    "            embeds = [x.to(CUDA_DEV) for x in embeds]\n",
    "            target = target.to(CUDA_DEV).float()\n",
    "            pred_logits = model(embeds).sigmoid().float()\n",
    "\n",
    "            y_pred, y_true = pred_logits.cpu().detach().numpy(), target.cpu().detach().numpy()\n",
    "            val_acc += average_precision_score(y_true, y_pred, average='macro', pos_label=1, sample_weight=None)\n",
    "\n",
    "            ce_loss = criterion(pred_logits, target)\n",
    "            val_loss += ce_loss.item()\n",
    "\n",
    "    warnings.filterwarnings(\"default\")\n",
    "    return val_loss / iteration, val_acc / iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2e624b0b-8daf-4702-a5de-c667fcd53121",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, loader):\n",
    "    model.eval()\n",
    "    track_idxs = []\n",
    "    predictions = []\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            track_idx, embeds = data\n",
    "            embeds = [x.to(CUDA_DEV) for x in embeds]\n",
    "            pred_probs = model(embeds).sigmoid()\n",
    "            predictions.append(pred_probs.cpu().numpy())\n",
    "            track_idxs.append(track_idx.numpy())\n",
    "    predictions = np.vstack(predictions)\n",
    "    track_idxs = np.vstack(track_idxs).ravel()\n",
    "    return track_idxs, predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92520883-098a-4d42-b538-478e199fd4f4",
   "metadata": {},
   "source": [
    "### Import W&B (if needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "70bbebc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if IS_WANDB:\n",
    "    import wandb\n",
    "    wandb.login()\n",
    "\n",
    "def wandb_init_if_needed():\n",
    "    if IS_WANDB:\n",
    "        wandb.init(\n",
    "            project=\"Trackformer\",\n",
    "            config={\n",
    "                \"architecture\": \"Transformer\",\n",
    "                \"learning_rate\": OPTIM_LR,\n",
    "                \"epochs\": N_EPOCH,\n",
    "            }\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "42b5be56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "433966f6-cf3f-401d-946e-1cbdea6c0ddb",
   "metadata": {},
   "source": [
    "## Train 1 Fold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0089e5e4-757f-41be-a06f-4fdfd06e0a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fold = 1 # Folds from 1 to 10 (by default)\n",
    "fold_size = df_train.shape[0] // N_FOLDS\n",
    "\n",
    "val_df = df_train[fold_size * (fold-1):fold_size * fold]\n",
    "train_df = df_train.loc[~df_train.index.isin(range(fold_size * (fold - 1), fold_size * fold))]\n",
    "\n",
    "test_dataset = TaggingDataset(df_test, True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn_test)\n",
    "\n",
    "train_dataset = TaggingDataset(train_df)\n",
    "val_dataset = TaggingDataset(val_df)\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9cca4b53-d7fc-478a-8235-551ea6480076",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total params: 4_680_960\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                            | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   2023-11-13 21:07:02.841561 batch 1 train_loss 0.72929 val_loss 0.61133 val_acc 0.04060\n",
      "   2023-11-13 21:07:14.194396 batch 101 train_loss 0.09590 val_loss 0.09712 val_acc 0.05238\n",
      "   2023-11-13 21:07:25.531234 batch 201 train_loss 0.07321 val_loss 0.07127 val_acc 0.05375\n",
      "   2023-11-13 21:07:36.841067 batch 301 train_loss 0.06190 val_loss 0.06550 val_acc 0.07175\n",
      "   2023-11-13 21:07:48.357905 batch 401 train_loss 0.05999 val_loss 0.06253 val_acc 0.08454\n",
      "   2023-11-13 21:07:59.936738 batch 501 train_loss 0.06027 val_loss 0.06063 val_acc 0.10087\n",
      "   2023-11-13 21:08:11.496572 batch 601 train_loss 0.05305 val_loss 0.05927 val_acc 0.11757\n",
      "   2023-11-13 21:08:22.811407 batch 701 train_loss 0.05560 val_loss 0.05766 val_acc 0.12854\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████| 1/1 [01:29<00:00, 89.11s/it]\n"
     ]
    }
   ],
   "source": [
    "model = Trackformer()\n",
    "criterion = nn.BCELoss()\n",
    "\n",
    "print(f'Total params: {model.n_params:_}')\n",
    "\n",
    "model = model.to(CUDA_DEV)\n",
    "criterion = criterion.to(CUDA_DEV)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=OPTIM_LR)\n",
    "\n",
    "wandb_init_if_needed()\n",
    "for epoch in tqdm(range(N_EPOCH)):\n",
    "    train_epoch(model, train_dataloader, val_dataloader, criterion, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9058b358-4ab4-4266-9c24-d6dbce01ad81",
   "metadata": {},
   "source": [
    "### Make Single Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f7222d0e-857f-426f-815b-308b67877a2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: total: 7.97 s\n",
      "Wall time: 7.95 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "track_idxs, predictions = predict(model, test_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2a792386-caa9-4bbc-8447-5448c5264b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_df = pd.DataFrame([\n",
    "    {'track': track, 'prediction': ','.join([str(p) for p in probs])}\n",
    "    for track, probs in zip(track_idxs, predictions)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "717c5830-848d-4f0e-8895-75d4601939a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_df.to_csv(f'pred_tformer_fold{fold}.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "20cf130f-4fd8-4fd6-bf65-26a89e283e70",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>track</th>\n",
       "      <th>prediction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>17730</td>\n",
       "      <td>0.37374988,0.20363799,0.2021929,0.3041287,0.03...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>32460</td>\n",
       "      <td>0.29124525,0.3581382,0.22837844,0.1856402,0.02...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>11288</td>\n",
       "      <td>0.14650339,0.22279377,0.18561621,0.23304635,0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>18523</td>\n",
       "      <td>0.54540426,0.13262658,0.28503522,0.11530596,0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71342</td>\n",
       "      <td>0.14968708,0.05640894,0.11685478,0.29343122,0....</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   track                                         prediction\n",
       "0  17730  0.37374988,0.20363799,0.2021929,0.3041287,0.03...\n",
       "1  32460  0.29124525,0.3581382,0.22837844,0.1856402,0.02...\n",
       "2  11288  0.14650339,0.22279377,0.18561621,0.23304635,0....\n",
       "3  18523  0.54540426,0.13262658,0.28503522,0.11530596,0....\n",
       "4  71342  0.14968708,0.05640894,0.11685478,0.29343122,0...."
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8717e0f9-931a-4ccf-8a17-66c013d97fae",
   "metadata": {},
   "source": [
    "## Train all 10 Folds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "09145e20",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for fold in range(1, N_FOLDS + 1):\n",
    "    print(f'FOLD-{fold}:')\n",
    "    ### Get Fold\n",
    "    fold_size = df_train.shape[0] // N_FOLDS\n",
    "\n",
    "    val_df = df_train[fold_size * (fold-1):fold_size * fold]\n",
    "    train_df = df_train.loc[~df_train.index.isin(range(fold_size * (fold - 1), fold_size * fold))]\n",
    "\n",
    "    test_dataset = TaggingDataset(df_test, True)\n",
    "    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn_test)\n",
    "\n",
    "    train_dataset = TaggingDataset(train_df)\n",
    "    val_dataset = TaggingDataset(val_df)\n",
    "\n",
    "    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)\n",
    "    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)\n",
    "    ### /Get Fold\n",
    "\n",
    "    ### Init new W&B session if needed:\n",
    "    wandb_init_if_needed()\n",
    "\n",
    "    ### Train Model\n",
    "    model = Trackformer()\n",
    "    criterion = nn.BCELoss()\n",
    "\n",
    "    print(f'Total params: {model.n_params:_}')\n",
    "\n",
    "    model = model.to(CUDA_DEV)\n",
    "    criterion = criterion.to(CUDA_DEV)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=OPTIM_LR)\n",
    "\n",
    "    for epoch in tqdm(range(N_EPOCH)):\n",
    "        train_epoch(model, train_dataloader, val_dataloader, criterion, optimizer)\n",
    "    ### /Train Model\n",
    "\n",
    "    ### Get Prediction\n",
    "    track_idxs, predictions = predict(model, test_dataloader)\n",
    "\n",
    "    predictions_df = pd.DataFrame([\n",
    "        {'track': track, 'prediction': ','.join([str(p) for p in probs])}\n",
    "        for track, probs in zip(track_idxs, predictions)\n",
    "    ])\n",
    "    ### /Get Prediction\n",
    "\n",
    "    ### Save predictions of current fold\n",
    "    predictions_df.to_csv(f'pred_tformer_fold{fold}.csv', index=False)\n",
    "    print(f'Fold-{fold} saved!\\n\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23b31e77-4682-4ff6-99fb-7866e7fb1485",
   "metadata": {},
   "source": [
    "### Bagging Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a6e0f9c1-a5d6-4bf1-87e3-3998d557fef8",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = []\n",
    "for i in tqdm(range(1, N_FOLDS + 1)):\n",
    "    pred = pd.read_csv(f'pred_tformer_fold{i}.csv')['prediction'].apply(lambda x: list(map(float, x.split(','))) ).values\n",
    "    pred = np.array([np.array(x) for x in pred])\n",
    "    preds.append(pred)\n",
    "preds = np.array(preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8d867ca5-fe67-4f52-8fe2-d463b46e9cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = preds.mean(axis=0)\n",
    "predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "bf45f69f-0177-477d-b099-7f74eb422875",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>track</th>\n",
       "      <th>prediction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>17730</td>\n",
       "      <td>0.025683811500000004,0.04249316646,0.121365329...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>32460</td>\n",
       "      <td>0.054375936999999985,0.4695480709999999,0.2681...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>11288</td>\n",
       "      <td>0.115668747,0.413396596,0.23798349599999996,0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>18523</td>\n",
       "      <td>0.658433986,0.02621826585,0.1576727775,0.38235...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71342</td>\n",
       "      <td>0.01329453447,0.00532173159,0.0136098291699999...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   track                                         prediction\n",
       "0  17730  0.025683811500000004,0.04249316646,0.121365329...\n",
       "1  32460  0.054375936999999985,0.4695480709999999,0.2681...\n",
       "2  11288  0.115668747,0.413396596,0.23798349599999996,0....\n",
       "3  18523  0.658433986,0.02621826585,0.1576727775,0.38235...\n",
       "4  71342  0.01329453447,0.00532173159,0.0136098291699999..."
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions_df = pd.DataFrame([\n",
    "    {'track': track, 'prediction': ','.join([str(p) for p in probs])}\n",
    "    for track, probs in zip(df_test.track.values, predictions)\n",
    "])\n",
    "\n",
    "predictions_df.to_csv(f'pred_tformer_fold_1_{N_FOLDS}.csv', index=False)\n",
    "predictions_df.head()"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}