demo/supervised-sequence-to-target-transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "03a90ff7",
   "metadata": {},
   "source": [
    "# Supervised task with transformer sequence encoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "862a9986",
   "metadata": {},
   "source": [
    "## Prepare your data\n",
    "\n",
    "- Use `Pyspark` in local or cluster mode for big dataset and `Pandas` for small.\n",
    "- Split data into required parts (train, valid, test, ...).\n",
    "- Use `ptls.preprocessing` for simple data preparation. \n",
    "- Transform features to compatible format using `Pyspark` or `Pandas` functions. \n",
    "You can also use `ptls.data_load.preprocessing` for common data transformation patterns.\n",
    "- Split sequences to `ptls-data` format with `ptls.data_load.split_tools`. Save prepared data into `Parquet` format or \n",
    "keep it in memory (`Pickle` also works).\n",
    "- Use one of the available `ptls.data_load.datasets` to define input for the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8bb3ccb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from ptls.preprocessing import PandasDataPreprocessor\n",
    "from ptls.data_load.datasets import MemoryMapDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81dcabf4",
   "metadata": {},
   "source": [
    "Read target data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2a0303dc",
   "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>client_id</th>\n",
       "      <th>bins</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>24662</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1046</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>34089</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>34848</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>47076</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   client_id  bins\n",
       "0      24662     2\n",
       "1       1046     0\n",
       "2      34089     2\n",
       "3      34848     1\n",
       "4      47076     3"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_target = pd.read_csv('https://huggingface.co/datasets/dllllb/age-group-prediction/resolve/main/train_target.csv?download=true')\n",
    "df_target.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d66e0d07",
   "metadata": {},
   "source": [
    "Split target data into train, test and validatation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2cc01e53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split 30000 records to train: 20000, valid: 3000, test: 7000\n"
     ]
    }
   ],
   "source": [
    "df_target_train, df_target_test = train_test_split(\n",
    "    df_target, test_size=7000, stratify=df_target['bins'], random_state=142)\n",
    "df_target_train, df_target_valid = train_test_split(\n",
    "    df_target_train, test_size=3000, stratify=df_target_train['bins'], random_state=142)\n",
    "print('Split {} records to train: {}, valid: {}, test: {}'.format(\n",
    "    *[len(df) for df in [df_target, df_target_train, df_target_valid, df_target_test]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25e0039",
   "metadata": {},
   "source": [
    "Load data with transactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2146a1ac",
   "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>client_id</th>\n",
       "      <th>trans_date</th>\n",
       "      <th>small_group</th>\n",
       "      <th>amount_rur</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>33172</td>\n",
       "      <td>6</td>\n",
       "      <td>4</td>\n",
       "      <td>71.463</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>33172</td>\n",
       "      <td>6</td>\n",
       "      <td>35</td>\n",
       "      <td>45.017</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>33172</td>\n",
       "      <td>8</td>\n",
       "      <td>11</td>\n",
       "      <td>13.887</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>33172</td>\n",
       "      <td>9</td>\n",
       "      <td>11</td>\n",
       "      <td>15.983</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>33172</td>\n",
       "      <td>10</td>\n",
       "      <td>11</td>\n",
       "      <td>21.341</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   client_id  trans_date  small_group  amount_rur\n",
       "0      33172           6            4      71.463\n",
       "1      33172           6           35      45.017\n",
       "2      33172           8           11      13.887\n",
       "3      33172           9           11      15.983\n",
       "4      33172          10           11      21.341"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_trx = pd.read_csv('https://huggingface.co/datasets/dllllb/age-group-prediction/resolve/main/transactions_train.csv.gz?download=true', compression='gzip')\n",
    "df_trx.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "01f61caa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split 26450577 transactions to train: 17622321, valid: 2634248, test: 6194008\n"
     ]
    }
   ],
   "source": [
    "df_trx_train = pd.merge(df_trx, df_target_train['client_id'], on='client_id', how='inner')\n",
    "df_trx_valid = pd.merge(df_trx, df_target_valid['client_id'], on='client_id', how='inner')\n",
    "df_trx_test = pd.merge(df_trx, df_target_test['client_id'], on='client_id', how='inner')\n",
    "print('Split {} transactions to train: {}, valid: {}, test: {}'.format(\n",
    "    *[len(df) for df in [df_trx, df_trx_train, df_trx_valid, df_trx_test]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "37bfd82d",
   "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>client_id</th>\n",
       "      <th>trans_date</th>\n",
       "      <th>small_group</th>\n",
       "      <th>amount_rur</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4159434</th>\n",
       "      <td>42191</td>\n",
       "      <td>308</td>\n",
       "      <td>2</td>\n",
       "      <td>76.361</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10477921</th>\n",
       "      <td>30412</td>\n",
       "      <td>628</td>\n",
       "      <td>11</td>\n",
       "      <td>20.237</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2682069</th>\n",
       "      <td>1095</td>\n",
       "      <td>353</td>\n",
       "      <td>1</td>\n",
       "      <td>64.957</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10469276</th>\n",
       "      <td>20216</td>\n",
       "      <td>160</td>\n",
       "      <td>91</td>\n",
       "      <td>14.234</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4149024</th>\n",
       "      <td>48112</td>\n",
       "      <td>580</td>\n",
       "      <td>3</td>\n",
       "      <td>4.579</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          client_id  trans_date  small_group  amount_rur\n",
       "4159434       42191         308            2      76.361\n",
       "10477921      30412         628           11      20.237\n",
       "2682069        1095         353            1      64.957\n",
       "10469276      20216         160           91      14.234\n",
       "4149024       48112         580            3       4.579"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_trx_train.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "11aa0d26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform flat table to dictionaries with client features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "129f938e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ptls.preprocessing import PandasDataPreprocessor\n",
    "\n",
    "preprocessor = PandasDataPreprocessor(\n",
    "    col_id='client_id',\n",
    "    col_event_time='trans_date',\n",
    "    event_time_transformation='none',\n",
    "    cols_category=['small_group'],\n",
    "    cols_numerical=['amount_rur'],\n",
    "    return_records=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3b4f6f65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30.2 s, sys: 6.14 s, total: 36.3 s\n",
      "Wall time: 36.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df_data_train = preprocessor.fit_transform(df_trx_train)\n",
    "df_data_valid = preprocessor.transform(df_trx_valid)\n",
    "df_data_test = preprocessor.transform(df_trx_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "bf17bf2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Record in dataset, train 20000, valid 3000, test 7000\n",
      "Each record is a client with list of transactions\n"
     ]
    }
   ],
   "source": [
    "print('Record in dataset, train {}, valid {}, test {}\\nEach record is a client with list of transactions'.format(\n",
    "    *[len(df) for df in [df_data_train, df_data_valid, df_data_test]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b48cca65",
   "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>client_id</th>\n",
       "      <th>trans_date</th>\n",
       "      <th>event_time</th>\n",
       "      <th>small_group</th>\n",
       "      <th>amount_rur</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>[tensor(0), tensor(5), tensor(10), tensor(11),...</td>\n",
       "      <td>[tensor(0), tensor(5), tensor(10), tensor(11),...</td>\n",
       "      <td>[tensor(4), tensor(3), tensor(1), tensor(3), t...</td>\n",
       "      <td>[tensor(4.0540, dtype=torch.float64), tensor(1...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7</td>\n",
       "      <td>[tensor(1), tensor(2), tensor(12), tensor(13),...</td>\n",
       "      <td>[tensor(1), tensor(2), tensor(12), tensor(13),...</td>\n",
       "      <td>[tensor(3), tensor(53), tensor(1), tensor(5), ...</td>\n",
       "      <td>[tensor(18.3190, dtype=torch.float64), tensor(...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>12</td>\n",
       "      <td>[tensor(3), tensor(6), tensor(6), tensor(6), t...</td>\n",
       "      <td>[tensor(3), tensor(6), tensor(6), tensor(6), t...</td>\n",
       "      <td>[tensor(1), tensor(19), tensor(13), tensor(6),...</td>\n",
       "      <td>[tensor(3.0220, dtype=torch.float64), tensor(2...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   client_id                                         trans_date  \\\n",
       "0          6  [tensor(0), tensor(5), tensor(10), tensor(11),...   \n",
       "1          7  [tensor(1), tensor(2), tensor(12), tensor(13),...   \n",
       "2         12  [tensor(3), tensor(6), tensor(6), tensor(6), t...   \n",
       "\n",
       "                                          event_time  \\\n",
       "0  [tensor(0), tensor(5), tensor(10), tensor(11),...   \n",
       "1  [tensor(1), tensor(2), tensor(12), tensor(13),...   \n",
       "2  [tensor(3), tensor(6), tensor(6), tensor(6), t...   \n",
       "\n",
       "                                         small_group  \\\n",
       "0  [tensor(4), tensor(3), tensor(1), tensor(3), t...   \n",
       "1  [tensor(3), tensor(53), tensor(1), tensor(5), ...   \n",
       "2  [tensor(1), tensor(19), tensor(13), tensor(6),...   \n",
       "\n",
       "                                          amount_rur  \n",
       "0  [tensor(4.0540, dtype=torch.float64), tensor(1...  \n",
       "1  [tensor(18.3190, dtype=torch.float64), tensor(...  \n",
       "2  [tensor(3.0220, dtype=torch.float64), tensor(2...  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_data_train.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "514a6f27",
   "metadata": {},
   "source": [
    "To learn our model, we need to add prefix target to target column due to feature naming rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f10d9339",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_target = df_target.rename(columns={'bins': 'target_bin'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9e9d418e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_data_train = pd.merge(df_data_train, df_target, on='client_id')\n",
    "df_data_valid = pd.merge(df_data_valid, df_target, on='client_id')\n",
    "df_data_test = pd.merge(df_data_test, df_target, on='client_id')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "baf6c255",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_data_train = df_data_train.to_dict(orient='records')\n",
    "df_data_valid = df_data_valid.to_dict(orient='records')\n",
    "df_data_test = df_data_test.to_dict(orient='records')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f5c2ed62",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'client_id': 6,\n",
       " 'trans_date': tensor([ 0,  5, 10, 11, 15, 15, 16, 16, 17, 18]),\n",
       " 'small_group': tensor([ 4,  3,  1,  3,  4,  1,  4,  3, 18,  2]),\n",
       " 'amount_rur': tensor([ 4.0540, 13.7380, 20.7010, 21.5640, 13.4990, 23.7220,  4.3040,  8.6250,\n",
       "         12.9380, 28.1620], dtype=torch.float64),\n",
       " 'event_time': tensor([ 0,  5, 10, 11, 15, 15, 16, 16, 17, 18]),\n",
       " 'target_bin': 1}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show first 10 transactions from one record\n",
    "rec = df_data_train[0]\n",
    "{k: v[:10] if type(v) is torch.Tensor else v for k, v in rec.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ba4f831",
   "metadata": {},
   "source": [
    "Memory map dataset is a torch dataset, but also use filters to preprocess our data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "bb94744d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = MemoryMapDataset(df_data_train)\n",
    "dataset_valid = MemoryMapDataset(df_data_valid)\n",
    "dataset_test = MemoryMapDataset(df_data_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01d06ea",
   "metadata": {},
   "source": [
    "## Build encoder\n",
    "\n",
    "- All parts are available in `ptls.nn`.\n",
    "- You can also use pretrained layers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c59e65",
   "metadata": {},
   "source": [
    "In this task we will use TransformerSeqEncoder based on transformer architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "fd9840d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchmetrics\n",
    "from ptls.nn import TrxEncoder, TransformerSeqEncoder, Head"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f800708",
   "metadata": {},
   "source": [
    "Define TrxEncoder to learn embedding for single transaction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "d228aea0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trx_encoder=TrxEncoder(\n",
    "        embeddings={\n",
    "            'small_group': {'in': 150, 'out': 31},\n",
    "        },\n",
    "        numeric_values={\n",
    "            'amount_rur': 'log',\n",
    "        },\n",
    "        embeddings_noise=0.001\n",
    ")\n",
    "\n",
    "trx_encoder.output_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be25a39",
   "metadata": {},
   "source": [
    "We can choose parameters for our transformer encoder, for example the number of heads in the multiheadattention, the number of sub-encoder-layers in the encoder and dimension of linear layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9a594a6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_params = {\n",
    "    \"n_heads\": 1,\n",
    "    \"dim_hidden\": 128,\n",
    "    \"n_layers\": 4,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d27b62a",
   "metadata": {},
   "source": [
    "Define sequence encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "7d51e5b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_encoder = TransformerSeqEncoder(\n",
    "    trx_encoder=trx_encoder,\n",
    "    **transformer_params\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0cc5f19",
   "metadata": {},
   "source": [
    "## Choose framework for encoder train\n",
    "\n",
    "- There are both supervised of unsupervised frameworks in `ptls.frames`.\n",
    "- Keep in mind that each framework requires his own batch format.\n",
    "Tools for batch collate can be found in the selected framework package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "34e999a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/conda/lib/python3.9/site-packages/pydantic/_internal/_fields.py:149: UserWarning: Field \"model_server_url\" has conflict with protected namespace \"model_\".\n",
      "\n",
      "You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ()`.\n",
      "  warnings.warn(\n",
      "/home/user/conda/lib/python3.9/site-packages/pydantic/_internal/_config.py:318: UserWarning: Valid config keys have changed in V2:\n",
      "* 'schema_extra' has been renamed to 'json_schema_extra'\n",
      "  warnings.warn(message, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "from functools import partial\n",
    "from ptls.frames.supervised import SeqToTargetDataset, SequenceToTarget\n",
    "from ptls.frames import PtlsDataModule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd1ce7bc",
   "metadata": {},
   "source": [
    "To define a model for supervised learning, we need sequence encoder to get embedding for one user, head to transform embeddings to solve our task. In this task we will use linear layer and softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "3c661043",
   "metadata": {},
   "outputs": [],
   "source": [
    "sup_module = SequenceToTarget(\n",
    "    seq_encoder=seq_encoder,\n",
    "    head=Head(input_size=seq_encoder.embedding_size, objective='classification', num_classes=4),\n",
    "    loss=torch.nn.NLLLoss(),\n",
    "    metric_list=torchmetrics.Accuracy(),\n",
    "    optimizer_partial=partial(torch.optim.Adam),\n",
    "    lr_scheduler_partial=partial(torch.optim.lr_scheduler.StepLR, step_size=4, gamma=0.5),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9323ba6",
   "metadata": {},
   "source": [
    "Data module to define data for training, validating and testing, batch sizes and number of workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "953afcf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sup_data = PtlsDataModule(\n",
    "    train_data=SeqToTargetDataset(dataset_train, target_col_name='target_bin', target_dtype=torch.long),\n",
    "    valid_data=SeqToTargetDataset(dataset_valid, target_col_name='target_bin', target_dtype=torch.long),\n",
    "    test_data=SeqToTargetDataset(dataset_test, target_col_name='target_bin', target_dtype=torch.long),\n",
    "    train_batch_size=128,\n",
    "    valid_batch_size=1024,\n",
    "    train_num_workers=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9a780c4",
   "metadata": {},
   "source": [
    "## Train your encoder with selected framework and `pytorch_lightning`\n",
    "\n",
    "- Provide data with one of the DataLoaders that is compatible with selected framework. \n",
    "- Monitor the progress on tensorboard.\n",
    "- Optionally tune hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "353ac8b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "2dca56e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the TensorBoard notebook extension\n",
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "b2a9920d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-ef90855983924e36\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-ef90855983924e36\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir lightning_logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "a768ecfd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    max_epochs=2,\n",
    "    gpus=1 if torch.cuda.is_available() else 0,\n",
    "    enable_progress_bar=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "6289402a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name          | Type                  | Params\n",
      "--------------------------------------------------------\n",
      "0 | seq_encoder   | TransformerSeqEncoder | 55.6 K\n",
      "1 | head          | Head                  | 132   \n",
      "2 | loss          | NLLLoss               | 0     \n",
      "3 | train_metrics | ModuleDict            | 0     \n",
      "4 | valid_metrics | ModuleDict            | 0     \n",
      "5 | test_metrics  | ModuleDict            | 0     \n",
      "--------------------------------------------------------\n",
      "55.7 K    Trainable params\n",
      "0         Non-trainable params\n",
      "55.7 K    Total params\n",
      "0.223     Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logger.version = 1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6bfa897d94b44bc4a99c92567b36fffe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6h 38min 45s, sys: 2h 38min 26s, total: 9h 17min 11s\n",
      "Wall time: 1h 15min 14s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(f'logger.version = {trainer.logger.version}')\n",
    "trainer.fit(sup_module, sup_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "285e7114",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': tensor(1.0247), 'seq_len': tensor(851.9688), 'y': tensor(1.6250), 'val_loss': tensor(1.0196), 'valid/Accuracy': tensor(0.5547), 'train/Accuracy': tensor(0.5416)}\n"
     ]
    }
   ],
   "source": [
    "# train and validation metrics\n",
    "print(trainer.logged_metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "55b93bc9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Restoring states from the checkpoint path at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=1-step=314.ckpt\n",
      "Loaded model weights from checkpoint at /home/jovyan/lightning_logs/version_1/checkpoints/epoch=1-step=314.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05c28bcf9b7d4eac9aa44690f955fb36",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">       test/Accuracy       </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.5495714545249939     </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m      test/Accuracy      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.5495714545249939    \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'test/Accuracy': 0.5495714545249939}]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test metrics\n",
    "trainer.test(ckpt_path='best', dataloaders=sup_data.test_dataloader())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538a5c57",
   "metadata": {},
   "source": [
    "# Make predict\n",
    "\n",
    "Let's make predict to check metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "46b36509",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ptls.data_load.utils import collate_feature_dict\n",
    "from ptls.frames.inference_module import InferenceModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "fd672b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_dl = torch.utils.data.DataLoader(\n",
    "    dataset=dataset_test,\n",
    "    collate_fn=collate_feature_dict,\n",
    "    shuffle=False,\n",
    "    batch_size=1000,\n",
    "    num_workers=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "69fadd37",
   "metadata": {},
   "outputs": [],
   "source": [
    "inf_module = InferenceModule(\n",
    "    torch.nn.Sequential(\n",
    "        sup_module,\n",
    "        torch.nn.Softmax(dim=1),\n",
    "    ),\n",
    "    model_out_name='prob',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "9abe09cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd3a4b57c3e74c95ba0fb0e2b5a5ba1b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 157it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_predict = trainer.predict(inf_module, inference_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "e959950c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_predict = pd.concat(df_predict, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "c19c6b4f",
   "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>client_id</th>\n",
       "      <th>target_bin</th>\n",
       "      <th>prob_0000</th>\n",
       "      <th>prob_0001</th>\n",
       "      <th>prob_0002</th>\n",
       "      <th>prob_0003</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>530</th>\n",
       "      <td>3769</td>\n",
       "      <td>2</td>\n",
       "      <td>0.065010</td>\n",
       "      <td>0.004084</td>\n",
       "      <td>0.922117</td>\n",
       "      <td>0.008790</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>958</th>\n",
       "      <td>13871</td>\n",
       "      <td>2</td>\n",
       "      <td>0.262393</td>\n",
       "      <td>0.031538</td>\n",
       "      <td>0.622771</td>\n",
       "      <td>0.083299</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>671</th>\n",
       "      <td>11754</td>\n",
       "      <td>3</td>\n",
       "      <td>0.068451</td>\n",
       "      <td>0.667012</td>\n",
       "      <td>0.005306</td>\n",
       "      <td>0.259231</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>708</th>\n",
       "      <td>19111</td>\n",
       "      <td>3</td>\n",
       "      <td>0.099277</td>\n",
       "      <td>0.610973</td>\n",
       "      <td>0.006524</td>\n",
       "      <td>0.283227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>368</th>\n",
       "      <td>31272</td>\n",
       "      <td>1</td>\n",
       "      <td>0.213297</td>\n",
       "      <td>0.450815</td>\n",
       "      <td>0.043134</td>\n",
       "      <td>0.292753</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>520</th>\n",
       "      <td>39318</td>\n",
       "      <td>2</td>\n",
       "      <td>0.045078</td>\n",
       "      <td>0.004748</td>\n",
       "      <td>0.942754</td>\n",
       "      <td>0.007420</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>272</th>\n",
       "      <td>44799</td>\n",
       "      <td>0</td>\n",
       "      <td>0.511363</td>\n",
       "      <td>0.033735</td>\n",
       "      <td>0.281010</td>\n",
       "      <td>0.173892</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>615</th>\n",
       "      <td>25731</td>\n",
       "      <td>0</td>\n",
       "      <td>0.501576</td>\n",
       "      <td>0.047266</td>\n",
       "      <td>0.249107</td>\n",
       "      <td>0.202051</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>811</th>\n",
       "      <td>12658</td>\n",
       "      <td>0</td>\n",
       "      <td>0.229762</td>\n",
       "      <td>0.408421</td>\n",
       "      <td>0.107694</td>\n",
       "      <td>0.254123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>72</th>\n",
       "      <td>21866</td>\n",
       "      <td>2</td>\n",
       "      <td>0.371182</td>\n",
       "      <td>0.013171</td>\n",
       "      <td>0.534654</td>\n",
       "      <td>0.080994</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     client_id  target_bin  prob_0000  prob_0001  prob_0002  prob_0003\n",
       "530       3769           2   0.065010   0.004084   0.922117   0.008790\n",
       "958      13871           2   0.262393   0.031538   0.622771   0.083299\n",
       "671      11754           3   0.068451   0.667012   0.005306   0.259231\n",
       "708      19111           3   0.099277   0.610973   0.006524   0.283227\n",
       "368      31272           1   0.213297   0.450815   0.043134   0.292753\n",
       "520      39318           2   0.045078   0.004748   0.942754   0.007420\n",
       "272      44799           0   0.511363   0.033735   0.281010   0.173892\n",
       "615      25731           0   0.501576   0.047266   0.249107   0.202051\n",
       "811      12658           0   0.229762   0.408421   0.107694   0.254123\n",
       "72       21866           2   0.371182   0.013171   0.534654   0.080994"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_predict.sample(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "db5e2b4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 0, 1, ..., 1, 2, 2])"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_pred = df_predict[[f'prob_{i:04d}' for i in range(4)]].values.argmax(axis=1)\n",
    "y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "bddd6d04",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3, 3, 1, ..., 1, 2, 2])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_true = df_predict['target_bin'].values\n",
    "y_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "ad4d074a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix, accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "cb10cd53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5495714285714286"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy_score(y_true, y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "3b4df3c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 667,  155,  698,  215],\n",
       "       [ 221, 1148,  101,  279],\n",
       "       [ 155,   35, 1557,   17],\n",
       "       [ 557,  501,  219,  475]])"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "confusion_matrix(y_true, y_pred)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root] *",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}