KSI_LSTMattn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torchinfo import summary\n",
    "\n",
    "from KSI_models import KSI, ModifiedKSI, LSTMattn\n",
    "from KSI_utils import load_KSI_data, train_model, test_model\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_embedding = 100\n",
    "n_hidden = 100 # 200 in paper, but too intensive for my machine\n",
    "batch_size = 32\n",
    "n_epochs = 25\n",
    "save = True\n",
    "profile = False\n",
    "model_type = 'LSTMattn'\n",
    "early_stopping = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir = 'data/original/'\n",
    "loaders, wikivec, word_to_ix = load_KSI_data(dir=dir, \n",
    "                                             batch_size=batch_size, \n",
    "                                             train=True, \n",
    "                                             val=True, \n",
    "                                             test=True, \n",
    "                                             device=DEVICE)\n",
    "train_dataloader = loaders['train']\n",
    "val_dataloader = loaders['val']\n",
    "test_dataloader = loaders['test']\n",
    "\n",
    "n_wiki, n_vocab = wikivec.shape\n",
    "n_words = len(word_to_ix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note_lengths = []\n",
    "# for data in train_dataloader:\n",
    "#     n, _, _ = data\n",
    "#     note_lengths.append(n.shape[1])\n",
    "# avg_note_size = np.round(np.array(note_lengths).mean()).astype(int)\n",
    "\n",
    "avg_note_size = 2455"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "LSTMattn                                 --                        --\n",
       "├─Embedding: 1-1                         [32, 2455, 100]           4,796,200\n",
       "├─Dropout: 1-2                           [2455, 32, 100]           --\n",
       "├─LSTM: 1-3                              [2455, 32, 100]           80,800\n",
       "├─Linear: 1-4                            --                        34,400\n",
       "├─Linear: 1-5                            --                        34,744\n",
       "==========================================================================================\n",
       "Total params: 4,946,144\n",
       "Trainable params: 4,946,144\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 6.50\n",
       "==========================================================================================\n",
       "Input size (MB): 1.87\n",
       "Forward/backward pass size (MB): 125.70\n",
       "Params size (MB): 19.78\n",
       "Estimated Total Size (MB): 147.35\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "base_model = LSTMattn(n_words, n_wiki, n_embedding, n_hidden, batch_size)\n",
    "base_model = base_model.to(DEVICE)\n",
    "base_summary = summary(base_model, [(batch_size, avg_note_size), \n",
    "                                    (batch_size, n_vocab)], \n",
    "                       dtypes=[torch.int, torch.float])\n",
    "\n",
    "base_summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.4405, Val Recall@10: 0.4460, Train Micro F1: 0.0000, Val Micro F1: 0.0000, Train Macro F1: 0.0000, Val Macro F1: 0.0000, Train Micro AUC: 0.9269, Val Micro AUC: 0.9123, Train Macro AUC: 0.5821, Val Macro AUC: 0.5883\n",
      "Epoch: 002, Train Recall@10: 0.5541, Val Recall@10: 0.5566, Train Micro F1: 0.4543, Val Micro F1: 0.4571, Train Macro F1: 0.0207, Val Macro F1: 0.0252, Train Micro AUC: 0.9415, Val Micro AUC: 0.9295, Train Macro AUC: 0.6246, Val Macro AUC: 0.6363\n",
      "Epoch: 003, Train Recall@10: 0.6379, Val Recall@10: 0.6377, Train Micro F1: 0.5660, Val Micro F1: 0.5633, Train Macro F1: 0.0509, Val Macro F1: 0.0610, Train Micro AUC: 0.9536, Val Micro AUC: 0.9435, Train Macro AUC: 0.6631, Val Macro AUC: 0.6796\n",
      "Epoch: 004, Train Recall@10: 0.7006, Val Recall@10: 0.6949, Train Micro F1: 0.6051, Val Micro F1: 0.5996, Train Macro F1: 0.0719, Val Macro F1: 0.0858, Train Micro AUC: 0.9614, Val Micro AUC: 0.9526, Train Macro AUC: 0.6839, Val Macro AUC: 0.7050\n",
      "Epoch: 005, Train Recall@10: 0.7329, Val Recall@10: 0.7262, Train Micro F1: 0.6298, Val Micro F1: 0.6228, Train Macro F1: 0.0864, Val Macro F1: 0.1046, Train Micro AUC: 0.9667, Val Micro AUC: 0.9589, Train Macro AUC: 0.7133, Val Macro AUC: 0.7366\n",
      "Epoch: 006, Train Recall@10: 0.7645, Val Recall@10: 0.7564, Train Micro F1: 0.6505, Val Micro F1: 0.6408, Train Macro F1: 0.1102, Val Macro F1: 0.1325, Train Micro AUC: 0.9720, Val Micro AUC: 0.9648, Train Macro AUC: 0.7402, Val Macro AUC: 0.7615\n",
      "Epoch: 007, Train Recall@10: 0.7813, Val Recall@10: 0.7684, Train Micro F1: 0.6652, Val Micro F1: 0.6524, Train Macro F1: 0.1280, Val Macro F1: 0.1515, Train Micro AUC: 0.9751, Val Micro AUC: 0.9681, Train Macro AUC: 0.7698, Val Macro AUC: 0.7786\n",
      "Epoch: 008, Train Recall@10: 0.7919, Val Recall@10: 0.7786, Train Micro F1: 0.6734, Val Micro F1: 0.6586, Train Macro F1: 0.1366, Val Macro F1: 0.1590, Train Micro AUC: 0.9771, Val Micro AUC: 0.9702, Train Macro AUC: 0.7942, Val Macro AUC: 0.7892\n",
      "Epoch: 009, Train Recall@10: 0.8072, Val Recall@10: 0.7938, Train Micro F1: 0.6847, Val Micro F1: 0.6653, Train Macro F1: 0.1534, Val Macro F1: 0.1779, Train Micro AUC: 0.9796, Val Micro AUC: 0.9728, Train Macro AUC: 0.8090, Val Macro AUC: 0.8051\n",
      "Epoch: 010, Train Recall@10: 0.8186, Val Recall@10: 0.8011, Train Micro F1: 0.6913, Val Micro F1: 0.6689, Train Macro F1: 0.1677, Val Macro F1: 0.1946, Train Micro AUC: 0.9813, Val Micro AUC: 0.9744, Train Macro AUC: 0.8227, Val Macro AUC: 0.8143\n",
      "Epoch: 011, Train Recall@10: 0.8268, Val Recall@10: 0.8069, Train Micro F1: 0.6995, Val Micro F1: 0.6735, Train Macro F1: 0.1878, Val Macro F1: 0.2109, Train Micro AUC: 0.9826, Val Micro AUC: 0.9754, Train Macro AUC: 0.8367, Val Macro AUC: 0.8261\n",
      "Epoch: 012, Train Recall@10: 0.8337, Val Recall@10: 0.8111, Train Micro F1: 0.7068, Val Micro F1: 0.6758, Train Macro F1: 0.2031, Val Macro F1: 0.2251, Train Micro AUC: 0.9838, Val Micro AUC: 0.9763, Train Macro AUC: 0.8478, Val Macro AUC: 0.8351\n",
      "Epoch: 013, Train Recall@10: 0.8410, Val Recall@10: 0.8153, Train Micro F1: 0.7129, Val Micro F1: 0.6766, Train Macro F1: 0.2138, Val Macro F1: 0.2345, Train Micro AUC: 0.9848, Val Micro AUC: 0.9772, Train Macro AUC: 0.8615, Val Macro AUC: 0.8472\n",
      "Epoch: 014, Train Recall@10: 0.8453, Val Recall@10: 0.8160, Train Micro F1: 0.7175, Val Micro F1: 0.6791, Train Macro F1: 0.2206, Val Macro F1: 0.2427, Train Micro AUC: 0.9856, Val Micro AUC: 0.9777, Train Macro AUC: 0.8755, Val Macro AUC: 0.8497\n",
      "Epoch: 015, Train Recall@10: 0.8511, Val Recall@10: 0.8188, Train Micro F1: 0.7221, Val Micro F1: 0.6815, Train Macro F1: 0.2276, Val Macro F1: 0.2471, Train Micro AUC: 0.9866, Val Micro AUC: 0.9783, Train Macro AUC: 0.8856, Val Macro AUC: 0.8574\n",
      "Epoch: 016, Train Recall@10: 0.8570, Val Recall@10: 0.8232, Train Micro F1: 0.7275, Val Micro F1: 0.6818, Train Macro F1: 0.2389, Val Macro F1: 0.2588, Train Micro AUC: 0.9874, Val Micro AUC: 0.9791, Train Macro AUC: 0.8944, Val Macro AUC: 0.8611\n",
      "Epoch: 017, Train Recall@10: 0.8612, Val Recall@10: 0.8230, Train Micro F1: 0.7322, Val Micro F1: 0.6834, Train Macro F1: 0.2492, Val Macro F1: 0.2631, Train Micro AUC: 0.9880, Val Micro AUC: 0.9793, Train Macro AUC: 0.9006, Val Macro AUC: 0.8603\n",
      "Epoch: 018, Train Recall@10: 0.8657, Val Recall@10: 0.8235, Train Micro F1: 0.7384, Val Micro F1: 0.6848, Train Macro F1: 0.2589, Val Macro F1: 0.2691, Train Micro AUC: 0.9887, Val Micro AUC: 0.9794, Train Macro AUC: 0.9062, Val Macro AUC: 0.8632\n",
      "Epoch: 019, Train Recall@10: 0.8696, Val Recall@10: 0.8261, Train Micro F1: 0.7435, Val Micro F1: 0.6868, Train Macro F1: 0.2689, Val Macro F1: 0.2773, Train Micro AUC: 0.9892, Val Micro AUC: 0.9796, Train Macro AUC: 0.9117, Val Macro AUC: 0.8684\n",
      "Epoch: 020, Train Recall@10: 0.8742, Val Recall@10: 0.8248, Train Micro F1: 0.7473, Val Micro F1: 0.6846, Train Macro F1: 0.2734, Val Macro F1: 0.2772, Train Micro AUC: 0.9898, Val Micro AUC: 0.9796, Train Macro AUC: 0.9182, Val Macro AUC: 0.8701\n",
      "Epoch: 021, Train Recall@10: 0.8779, Val Recall@10: 0.8254, Train Micro F1: 0.7515, Val Micro F1: 0.6842, Train Macro F1: 0.2838, Val Macro F1: 0.2807, Train Micro AUC: 0.9903, Val Micro AUC: 0.9797, Train Macro AUC: 0.9234, Val Macro AUC: 0.8716\n",
      "Epoch: 022, Train Recall@10: 0.8810, Val Recall@10: 0.8256, Train Micro F1: 0.7556, Val Micro F1: 0.6842, Train Macro F1: 0.2926, Val Macro F1: 0.2855, Train Micro AUC: 0.9907, Val Micro AUC: 0.9796, Train Macro AUC: 0.9272, Val Macro AUC: 0.8742\n",
      "Epoch: 023, Train Recall@10: 0.8839, Val Recall@10: 0.8246, Train Micro F1: 0.7578, Val Micro F1: 0.6827, Train Macro F1: 0.2977, Val Macro F1: 0.2891, Train Micro AUC: 0.9911, Val Micro AUC: 0.9796, Train Macro AUC: 0.9314, Val Macro AUC: 0.8746\n",
      "Epoch: 024, Train Recall@10: 0.8877, Val Recall@10: 0.8251, Train Micro F1: 0.7645, Val Micro F1: 0.6815, Train Macro F1: 0.3131, Val Macro F1: 0.2931, Train Micro AUC: 0.9916, Val Micro AUC: 0.9795, Train Macro AUC: 0.9361, Val Macro AUC: 0.8764\n",
      "Early stopping at epoch 24\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(base_model.parameters())\n",
    "prof_base = train_model(base_model, \n",
    "                        train_dataloader=train_dataloader,\n",
    "                        val_dataloader=val_dataloader,\n",
    "                        optimizer=optimizer,\n",
    "                        n_epochs=n_epochs,\n",
    "                        profile=profile, \n",
    "                        log_path=f'./log/{model_type}',\n",
    "                        device=DEVICE,\n",
    "                        init_hidden=True,\n",
    "                        early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save:\n",
    "    torch.save(base_model, f'{dir}{model_type}_model.pt')\n",
    "if profile:\n",
    "    print(prof_base.key_averages(group_by_stack_n=5).table(sort_by='self_cuda_time_total'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.8242, Test Micro F1: 0.6848, Test Macro F1: 0.2594, Test Micro AUC: 0.9805, Test Macro AUC: 0.8545\n"
     ]
    }
   ],
   "source": [
    "tt_recall_at_k, tt_micro_f1, tt_macro_f1, tt_micro_auc, tt_macro_auc, label_aucs_base = test_model(base_model, \n",
    "                                                                                                   test_dataloader, \n",
    "                                                                                                   wikivec,\n",
    "                                                                                                   device=DEVICE,\n",
    "                                                                                                   init_hidden=True)\n",
    "print(f'Test Recall@10: {tt_recall_at_k:.4f}, Test Micro F1: {tt_micro_f1:.4f}, Test Macro F1: {tt_macro_f1:.4f}' +\n",
    "      f', Test Micro AUC: {tt_micro_auc:.4f}, Test Macro AUC: {tt_macro_auc:.4f}')\n",
    "del base_model\n",
    "gc.collect()\n",
    "if DEVICE == 'cuda':\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "LSTMattn                                 --                        --\n",
       "├─KSI: 1-1                               --                        --\n",
       "│    └─Linear: 2-1                       --                        (recursive)\n",
       "│    └─Linear: 2-2                       --                        (recursive)\n",
       "│    └─Linear: 2-3                       --                        (recursive)\n",
       "├─Embedding: 1-2                         [32, 2455, 100]           4,796,200\n",
       "├─Dropout: 1-3                           [2455, 32, 100]           --\n",
       "├─LSTM: 1-4                              [2455, 32, 100]           80,800\n",
       "├─KSI: 1-1                               --                        --\n",
       "│    └─Linear: 2-4                       [32, 344, 100]            1,217,400\n",
       "│    └─Linear: 2-5                       [32, 344, 100]            10,100\n",
       "│    └─Linear: 2-6                       [32, 344, 1]              101\n",
       "==========================================================================================\n",
       "Total params: 6,104,601\n",
       "Trainable params: 6,104,601\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 6.54\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 143.40\n",
       "Params size (MB): 24.42\n",
       "Estimated Total Size (MB): 186.44\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ksi = KSI(n_embedding, n_vocab)\n",
    "ksi.to(DEVICE)\n",
    "model = LSTMattn(n_words, n_wiki, n_embedding, n_hidden, ksi=ksi)\n",
    "model = model.to(DEVICE)\n",
    "ksi_summary = summary(model, [(batch_size, avg_note_size), \n",
    "                              (batch_size, n_vocab),\n",
    "                              (n_wiki, n_vocab)], \n",
    "                      dtypes=[torch.int, torch.float, torch.float])\n",
    "\n",
    "ksi_summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.6576, Val Recall@10: 0.6528, Train Micro F1: 0.4033, Val Micro F1: 0.3975, Train Macro F1: 0.0596, Val Macro F1: 0.0709, Train Micro AUC: 0.9669, Val Micro AUC: 0.9588, Train Macro AUC: 0.8290, Val Macro AUC: 0.8168\n",
      "Epoch: 002, Train Recall@10: 0.7019, Val Recall@10: 0.6958, Train Micro F1: 0.4814, Val Micro F1: 0.4733, Train Macro F1: 0.0872, Val Macro F1: 0.0991, Train Micro AUC: 0.9720, Val Micro AUC: 0.9636, Train Macro AUC: 0.8622, Val Macro AUC: 0.8375\n",
      "Epoch: 003, Train Recall@10: 0.7398, Val Recall@10: 0.7277, Train Micro F1: 0.5297, Val Micro F1: 0.5138, Train Macro F1: 0.1134, Val Macro F1: 0.1225, Train Micro AUC: 0.9764, Val Micro AUC: 0.9680, Train Macro AUC: 0.8833, Val Macro AUC: 0.8478\n",
      "Epoch: 004, Train Recall@10: 0.7664, Val Recall@10: 0.7507, Train Micro F1: 0.5782, Val Micro F1: 0.5582, Train Macro F1: 0.1451, Val Macro F1: 0.1492, Train Micro AUC: 0.9793, Val Micro AUC: 0.9704, Train Macro AUC: 0.8978, Val Macro AUC: 0.8553\n",
      "Epoch: 005, Train Recall@10: 0.7887, Val Recall@10: 0.7657, Train Micro F1: 0.6126, Val Micro F1: 0.5861, Train Macro F1: 0.1778, Val Macro F1: 0.1713, Train Micro AUC: 0.9817, Val Micro AUC: 0.9722, Train Macro AUC: 0.9093, Val Macro AUC: 0.8601\n",
      "Epoch: 006, Train Recall@10: 0.8014, Val Recall@10: 0.7720, Train Micro F1: 0.6297, Val Micro F1: 0.5949, Train Macro F1: 0.2042, Val Macro F1: 0.1853, Train Micro AUC: 0.9834, Val Micro AUC: 0.9730, Train Macro AUC: 0.9199, Val Macro AUC: 0.8640\n",
      "Epoch: 007, Train Recall@10: 0.8149, Val Recall@10: 0.7755, Train Micro F1: 0.6470, Val Micro F1: 0.6065, Train Macro F1: 0.2320, Val Macro F1: 0.2003, Train Micro AUC: 0.9851, Val Micro AUC: 0.9736, Train Macro AUC: 0.9292, Val Macro AUC: 0.8682\n",
      "Epoch: 008, Train Recall@10: 0.8267, Val Recall@10: 0.7802, Train Micro F1: 0.6637, Val Micro F1: 0.6121, Train Macro F1: 0.2595, Val Macro F1: 0.2070, Train Micro AUC: 0.9865, Val Micro AUC: 0.9739, Train Macro AUC: 0.9371, Val Macro AUC: 0.8704\n",
      "Epoch: 009, Train Recall@10: 0.8358, Val Recall@10: 0.7809, Train Micro F1: 0.6759, Val Micro F1: 0.6155, Train Macro F1: 0.2852, Val Macro F1: 0.2177, Train Micro AUC: 0.9876, Val Micro AUC: 0.9736, Train Macro AUC: 0.9441, Val Macro AUC: 0.8711\n",
      "Epoch: 010, Train Recall@10: 0.8444, Val Recall@10: 0.7809, Train Micro F1: 0.6885, Val Micro F1: 0.6162, Train Macro F1: 0.3245, Val Macro F1: 0.2222, Train Micro AUC: 0.9886, Val Micro AUC: 0.9731, Train Macro AUC: 0.9502, Val Macro AUC: 0.8704\n",
      "Epoch: 011, Train Recall@10: 0.8511, Val Recall@10: 0.7783, Train Micro F1: 0.6969, Val Micro F1: 0.6140, Train Macro F1: 0.3468, Val Macro F1: 0.2237, Train Micro AUC: 0.9894, Val Micro AUC: 0.9718, Train Macro AUC: 0.9555, Val Macro AUC: 0.8680\n",
      "Epoch: 012, Train Recall@10: 0.8579, Val Recall@10: 0.7716, Train Micro F1: 0.7031, Val Micro F1: 0.6098, Train Macro F1: 0.3800, Val Macro F1: 0.2266, Train Micro AUC: 0.9900, Val Micro AUC: 0.9701, Train Macro AUC: 0.9597, Val Macro AUC: 0.8650\n",
      "Epoch: 013, Train Recall@10: 0.8632, Val Recall@10: 0.7680, Train Micro F1: 0.7083, Val Micro F1: 0.6052, Train Macro F1: 0.3988, Val Macro F1: 0.2299, Train Micro AUC: 0.9905, Val Micro AUC: 0.9679, Train Macro AUC: 0.9631, Val Macro AUC: 0.8603\n",
      "Epoch: 014, Train Recall@10: 0.8673, Val Recall@10: 0.7645, Train Micro F1: 0.7126, Val Micro F1: 0.6016, Train Macro F1: 0.4180, Val Macro F1: 0.2274, Train Micro AUC: 0.9907, Val Micro AUC: 0.9655, Train Macro AUC: 0.9660, Val Macro AUC: 0.8576\n",
      "Epoch: 015, Train Recall@10: 0.8732, Val Recall@10: 0.7601, Train Micro F1: 0.7238, Val Micro F1: 0.6038, Train Macro F1: 0.4440, Val Macro F1: 0.2362, Train Micro AUC: 0.9911, Val Micro AUC: 0.9640, Train Macro AUC: 0.9687, Val Macro AUC: 0.8540\n",
      "Early stopping at epoch 15\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters())\n",
    "prof_ksi = train_model(model, \n",
    "                       train_dataloader=train_dataloader,\n",
    "                       val_dataloader=val_dataloader,\n",
    "                       wikivec=wikivec,\n",
    "                       optimizer=optimizer,\n",
    "                       n_epochs=n_epochs, \n",
    "                       profile=profile, \n",
    "                       log_path=f'./log/{model_type}_KSI',\n",
    "                       device=DEVICE,\n",
    "                       init_hidden=True,\n",
    "                       early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save:\n",
    "    torch.save(model, f'{dir}{model_type}_KSI_model.pt')\n",
    "if profile:\n",
    "    print(prof_ksi.key_averages(group_by_stack_n=5).table(sort_by='self_cuda_time_total'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.7755, Test Micro F1: 0.6124, Test Macro F1: 0.2103, Test Micro AUC: 0.9745, Test Macro AUC: 0.8801\n"
     ]
    }
   ],
   "source": [
    "tt_recall_at_k, tt_micro_f1, tt_macro_f1, tt_micro_auc, tt_macro_auc, label_aucs_ksi = test_model(model, \n",
    "                                                                                                  test_dataloader, \n",
    "                                                                                                  wikivec,\n",
    "                                                                                                  device=DEVICE,\n",
    "                                                                                                  init_hidden=True)\n",
    "print(f'Test Recall@10: {tt_recall_at_k:.4f}, Test Micro F1: {tt_micro_f1:.4f}, Test Macro F1: {tt_macro_f1:.4f}' +\n",
    "      f', Test Micro AUC: {tt_micro_auc:.4f}, Test Macro AUC: {tt_macro_auc:.4f}')\n",
    "del model\n",
    "gc.collect()\n",
    "if DEVICE == 'cuda':\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modified KSI using frequency vectors rather than binary vectors\n",
    "dir = 'data/original_freqs/'\n",
    "loaders, wikivec, word_to_ix = load_KSI_data(dir=dir, \n",
    "                                             batch_size=batch_size, \n",
    "                                             train=True, \n",
    "                                             val=True, \n",
    "                                             test=True, \n",
    "                                             device=DEVICE)\n",
    "train_dataloader = loaders['train']\n",
    "val_dataloader = loaders['val']\n",
    "test_dataloader = loaders['test']\n",
    "\n",
    "n_wiki, n_vocab = wikivec.shape\n",
    "n_words = len(word_to_ix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "LSTMattn                                 --                        --\n",
       "├─ModifiedKSI: 1-1                       --                        --\n",
       "│    └─Linear: 2-1                       --                        (recursive)\n",
       "│    └─Linear: 2-2                       --                        (recursive)\n",
       "│    └─Linear: 2-3                       --                        (recursive)\n",
       "│    └─Linear: 2-4                       --                        (recursive)\n",
       "├─Embedding: 1-2                         [32, 2455, 100]           4,796,200\n",
       "├─Dropout: 1-3                           [2455, 32, 100]           --\n",
       "├─LSTM: 1-4                              [2455, 32, 100]           80,800\n",
       "├─ModifiedKSI: 1-1                       --                        --\n",
       "│    └─Linear: 2-5                       [32, 344, 12173, 1]       2\n",
       "│    └─Linear: 2-6                       [32, 344, 100]            1,217,400\n",
       "│    └─Linear: 2-7                       [32, 344, 100]            10,100\n",
       "│    └─Linear: 2-8                       [32, 344, 1]              101\n",
       "==========================================================================================\n",
       "Total params: 6,104,603\n",
       "Trainable params: 6,104,603\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 6.54\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 1215.40\n",
       "Params size (MB): 24.42\n",
       "Estimated Total Size (MB): 1258.44\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_ksi = ModifiedKSI(n_embedding, n_vocab)\n",
    "mod_ksi.to(DEVICE)\n",
    "mod_model = LSTMattn(n_words, n_wiki, n_embedding, n_hidden, ksi=mod_ksi)\n",
    "mod_model = mod_model.to(DEVICE)\n",
    "mod_summary = summary(mod_model, [(batch_size, avg_note_size), \n",
    "                                  (batch_size, n_vocab),\n",
    "                                  (n_wiki, n_vocab)], \n",
    "                      dtypes=[torch.int, torch.float, torch.float])\n",
    "\n",
    "mod_summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.6862, Val Recall@10: 0.6851, Train Micro F1: 0.4554, Val Micro F1: 0.4497, Train Macro F1: 0.0518, Val Macro F1: 0.0617, Train Micro AUC: 0.9693, Val Micro AUC: 0.9628, Train Macro AUC: 0.8360, Val Macro AUC: 0.8402\n",
      "Epoch: 002, Train Recall@10: 0.7156, Val Recall@10: 0.7135, Train Micro F1: 0.4894, Val Micro F1: 0.4830, Train Macro F1: 0.0976, Val Macro F1: 0.1166, Train Micro AUC: 0.9739, Val Micro AUC: 0.9679, Train Macro AUC: 0.8599, Val Macro AUC: 0.8598\n",
      "Epoch: 003, Train Recall@10: 0.7322, Val Recall@10: 0.7269, Train Micro F1: 0.5043, Val Micro F1: 0.4980, Train Macro F1: 0.1174, Val Macro F1: 0.1403, Train Micro AUC: 0.9759, Val Micro AUC: 0.9701, Train Macro AUC: 0.8692, Val Macro AUC: 0.8652\n",
      "Epoch: 004, Train Recall@10: 0.7597, Val Recall@10: 0.7556, Train Micro F1: 0.5603, Val Micro F1: 0.5520, Train Macro F1: 0.1394, Val Macro F1: 0.1483, Train Micro AUC: 0.9783, Val Micro AUC: 0.9727, Train Macro AUC: 0.8816, Val Macro AUC: 0.8738\n",
      "Epoch: 005, Train Recall@10: 0.7749, Val Recall@10: 0.7689, Train Micro F1: 0.5819, Val Micro F1: 0.5732, Train Macro F1: 0.1534, Val Macro F1: 0.1660, Train Micro AUC: 0.9798, Val Micro AUC: 0.9743, Train Macro AUC: 0.8904, Val Macro AUC: 0.8792\n",
      "Epoch: 006, Train Recall@10: 0.7854, Val Recall@10: 0.7780, Train Micro F1: 0.5957, Val Micro F1: 0.5845, Train Macro F1: 0.1668, Val Macro F1: 0.1766, Train Micro AUC: 0.9813, Val Micro AUC: 0.9757, Train Macro AUC: 0.8978, Val Macro AUC: 0.8856\n",
      "Epoch: 007, Train Recall@10: 0.7942, Val Recall@10: 0.7844, Train Micro F1: 0.6068, Val Micro F1: 0.5919, Train Macro F1: 0.1791, Val Macro F1: 0.1859, Train Micro AUC: 0.9822, Val Micro AUC: 0.9765, Train Macro AUC: 0.9038, Val Macro AUC: 0.8895\n",
      "Epoch: 008, Train Recall@10: 0.8050, Val Recall@10: 0.7938, Train Micro F1: 0.6241, Val Micro F1: 0.6065, Train Macro F1: 0.2028, Val Macro F1: 0.1982, Train Micro AUC: 0.9834, Val Micro AUC: 0.9774, Train Macro AUC: 0.9098, Val Macro AUC: 0.8924\n",
      "Epoch: 009, Train Recall@10: 0.8116, Val Recall@10: 0.7984, Train Micro F1: 0.6366, Val Micro F1: 0.6160, Train Macro F1: 0.2092, Val Macro F1: 0.2043, Train Micro AUC: 0.9842, Val Micro AUC: 0.9781, Train Macro AUC: 0.9156, Val Macro AUC: 0.8963\n",
      "Epoch: 010, Train Recall@10: 0.8181, Val Recall@10: 0.8009, Train Micro F1: 0.6457, Val Micro F1: 0.6214, Train Macro F1: 0.2188, Val Macro F1: 0.2088, Train Micro AUC: 0.9849, Val Micro AUC: 0.9786, Train Macro AUC: 0.9213, Val Macro AUC: 0.8989\n",
      "Epoch: 011, Train Recall@10: 0.8247, Val Recall@10: 0.8052, Train Micro F1: 0.6587, Val Micro F1: 0.6297, Train Macro F1: 0.2292, Val Macro F1: 0.2134, Train Micro AUC: 0.9857, Val Micro AUC: 0.9791, Train Macro AUC: 0.9267, Val Macro AUC: 0.9007\n",
      "Epoch: 012, Train Recall@10: 0.8300, Val Recall@10: 0.8085, Train Micro F1: 0.6665, Val Micro F1: 0.6336, Train Macro F1: 0.2456, Val Macro F1: 0.2203, Train Micro AUC: 0.9863, Val Micro AUC: 0.9794, Train Macro AUC: 0.9318, Val Macro AUC: 0.9021\n",
      "Epoch: 013, Train Recall@10: 0.8365, Val Recall@10: 0.8120, Train Micro F1: 0.6769, Val Micro F1: 0.6395, Train Macro F1: 0.2673, Val Macro F1: 0.2296, Train Micro AUC: 0.9871, Val Micro AUC: 0.9796, Train Macro AUC: 0.9365, Val Macro AUC: 0.9034\n",
      "Epoch: 014, Train Recall@10: 0.8419, Val Recall@10: 0.8127, Train Micro F1: 0.6882, Val Micro F1: 0.6458, Train Macro F1: 0.2864, Val Macro F1: 0.2382, Train Micro AUC: 0.9877, Val Micro AUC: 0.9799, Train Macro AUC: 0.9407, Val Macro AUC: 0.9044\n",
      "Epoch: 015, Train Recall@10: 0.8471, Val Recall@10: 0.8156, Train Micro F1: 0.6936, Val Micro F1: 0.6461, Train Macro F1: 0.3015, Val Macro F1: 0.2421, Train Micro AUC: 0.9882, Val Micro AUC: 0.9800, Train Macro AUC: 0.9445, Val Macro AUC: 0.9047\n",
      "Epoch: 016, Train Recall@10: 0.8525, Val Recall@10: 0.8177, Train Micro F1: 0.7040, Val Micro F1: 0.6496, Train Macro F1: 0.3190, Val Macro F1: 0.2524, Train Micro AUC: 0.9887, Val Micro AUC: 0.9800, Train Macro AUC: 0.9485, Val Macro AUC: 0.9055\n",
      "Epoch: 017, Train Recall@10: 0.8573, Val Recall@10: 0.8154, Train Micro F1: 0.7100, Val Micro F1: 0.6510, Train Macro F1: 0.3302, Val Macro F1: 0.2549, Train Micro AUC: 0.9892, Val Micro AUC: 0.9800, Train Macro AUC: 0.9511, Val Macro AUC: 0.9055\n",
      "Epoch: 018, Train Recall@10: 0.8617, Val Recall@10: 0.8154, Train Micro F1: 0.7180, Val Micro F1: 0.6536, Train Macro F1: 0.3429, Val Macro F1: 0.2638, Train Micro AUC: 0.9897, Val Micro AUC: 0.9800, Train Macro AUC: 0.9540, Val Macro AUC: 0.9076\n",
      "Epoch: 019, Train Recall@10: 0.8670, Val Recall@10: 0.8159, Train Micro F1: 0.7252, Val Micro F1: 0.6538, Train Macro F1: 0.3626, Val Macro F1: 0.2676, Train Micro AUC: 0.9902, Val Micro AUC: 0.9799, Train Macro AUC: 0.9568, Val Macro AUC: 0.9073\n",
      "Epoch: 020, Train Recall@10: 0.8712, Val Recall@10: 0.8143, Train Micro F1: 0.7315, Val Micro F1: 0.6558, Train Macro F1: 0.3800, Val Macro F1: 0.2766, Train Micro AUC: 0.9906, Val Micro AUC: 0.9798, Train Macro AUC: 0.9591, Val Macro AUC: 0.9062\n",
      "Epoch: 021, Train Recall@10: 0.8752, Val Recall@10: 0.8136, Train Micro F1: 0.7371, Val Micro F1: 0.6555, Train Macro F1: 0.4025, Val Macro F1: 0.2749, Train Micro AUC: 0.9910, Val Micro AUC: 0.9795, Train Macro AUC: 0.9613, Val Macro AUC: 0.9061\n",
      "Early stopping at epoch 21\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(mod_model.parameters())\n",
    "prof_mod_ksi = train_model(mod_model, \n",
    "                           train_dataloader=train_dataloader,\n",
    "                           val_dataloader=val_dataloader,\n",
    "                           wikivec=wikivec,\n",
    "                           optimizer=optimizer,\n",
    "                           n_epochs=n_epochs, \n",
    "                           profile=profile, \n",
    "                           log_path=f'./log/{model_type}_ModifiedKSI',\n",
    "                           device=DEVICE,\n",
    "                           init_hidden=True,\n",
    "                           early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save:\n",
    "    torch.save(mod_model, f'{dir}{model_type}_ModifiedKSI_model.pt')\n",
    "if profile:\n",
    "    print(prof_mod_ksi.key_averages(group_by_stack_n=5).table(sort_by='self_cuda_time_total'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.8119, Test Micro F1: 0.6483, Test Macro F1: 0.2476, Test Micro AUC: 0.9812, Test Macro AUC: 0.9063\n"
     ]
    }
   ],
   "source": [
    "tt_recall_at_k, tt_micro_f1, tt_macro_f1, tt_micro_auc, tt_macro_auc, label_aucs_mod = test_model(mod_model, \n",
    "                                                                                                  test_dataloader, \n",
    "                                                                                                  wikivec,\n",
    "                                                                                                  device=DEVICE,\n",
    "                                                                                                  init_hidden=True)\n",
    "print(f'Test Recall@10: {tt_recall_at_k:.4f}, Test Micro F1: {tt_micro_f1:.4f}, Test Macro F1: {tt_macro_f1:.4f}' +\n",
    "      f', Test Micro AUC: {tt_micro_auc:.4f}, Test Macro AUC: {tt_macro_auc:.4f}')\n",
    "del mod_model\n",
    "gc.collect()\n",
    "if DEVICE == 'cuda':\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modified KSI using tfidf vectors rather than binary vectors\n",
    "dir = 'data/original_tfidf/'\n",
    "loaders, wikivec, word_to_ix = load_KSI_data(dir=dir, \n",
    "                                             batch_size=batch_size, \n",
    "                                             train=True, \n",
    "                                             val=True, \n",
    "                                             test=True, \n",
    "                                             device=DEVICE)\n",
    "train_dataloader = loaders['train']\n",
    "val_dataloader = loaders['val']\n",
    "test_dataloader = loaders['test']\n",
    "\n",
    "n_wiki, n_vocab = wikivec.shape\n",
    "n_words = len(word_to_ix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "LSTMattn                                 --                        --\n",
       "├─ModifiedKSI: 1-1                       --                        --\n",
       "│    └─Linear: 2-1                       --                        (recursive)\n",
       "│    └─Linear: 2-2                       --                        (recursive)\n",
       "│    └─Linear: 2-3                       --                        (recursive)\n",
       "│    └─Linear: 2-4                       --                        (recursive)\n",
       "├─Embedding: 1-2                         [32, 2455, 100]           4,796,200\n",
       "├─Dropout: 1-3                           [2455, 32, 100]           --\n",
       "├─LSTM: 1-4                              [2455, 32, 100]           80,800\n",
       "├─ModifiedKSI: 1-1                       --                        --\n",
       "│    └─Linear: 2-5                       [32, 344, 12173, 1]       2\n",
       "│    └─Linear: 2-6                       [32, 344, 100]            1,217,400\n",
       "│    └─Linear: 2-7                       [32, 344, 100]            10,100\n",
       "│    └─Linear: 2-8                       [32, 344, 1]              101\n",
       "==========================================================================================\n",
       "Total params: 6,104,603\n",
       "Trainable params: 6,104,603\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 6.54\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 1215.40\n",
       "Params size (MB): 24.42\n",
       "Estimated Total Size (MB): 1258.44\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_ksi2 = ModifiedKSI(n_embedding, n_vocab)\n",
    "mod_ksi2.to(DEVICE)\n",
    "tfidf_model = LSTMattn(n_words, n_wiki, n_embedding, n_hidden, ksi=mod_ksi2)\n",
    "tfidf_model = tfidf_model.to(DEVICE)\n",
    "tfidf_summary = summary(tfidf_model, [(batch_size, avg_note_size), \n",
    "                                      (batch_size, n_vocab),\n",
    "                                      (n_wiki, n_vocab)], \n",
    "                        dtypes=[torch.int, torch.float, torch.float])\n",
    "\n",
    "tfidf_summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.6786, Val Recall@10: 0.6752, Train Micro F1: 0.4677, Val Micro F1: 0.4649, Train Macro F1: 0.0703, Val Macro F1: 0.0857, Train Micro AUC: 0.9691, Val Micro AUC: 0.9628, Train Macro AUC: 0.8340, Val Macro AUC: 0.8403\n",
      "Epoch: 002, Train Recall@10: 0.7019, Val Recall@10: 0.6985, Train Micro F1: 0.4768, Val Micro F1: 0.4731, Train Macro F1: 0.0923, Val Macro F1: 0.1065, Train Micro AUC: 0.9723, Val Micro AUC: 0.9662, Train Macro AUC: 0.8561, Val Macro AUC: 0.8549\n",
      "Epoch: 003, Train Recall@10: 0.7151, Val Recall@10: 0.7138, Train Micro F1: 0.4888, Val Micro F1: 0.4825, Train Macro F1: 0.1151, Val Macro F1: 0.1228, Train Micro AUC: 0.9742, Val Micro AUC: 0.9681, Train Macro AUC: 0.8675, Val Macro AUC: 0.8641\n",
      "Epoch: 004, Train Recall@10: 0.7256, Val Recall@10: 0.7228, Train Micro F1: 0.5020, Val Micro F1: 0.4953, Train Macro F1: 0.1327, Val Macro F1: 0.1420, Train Micro AUC: 0.9759, Val Micro AUC: 0.9697, Train Macro AUC: 0.8771, Val Macro AUC: 0.8693\n",
      "Epoch: 005, Train Recall@10: 0.7469, Val Recall@10: 0.7420, Train Micro F1: 0.5368, Val Micro F1: 0.5283, Train Macro F1: 0.1505, Val Macro F1: 0.1563, Train Micro AUC: 0.9778, Val Micro AUC: 0.9718, Train Macro AUC: 0.8861, Val Macro AUC: 0.8764\n",
      "Epoch: 006, Train Recall@10: 0.7590, Val Recall@10: 0.7527, Train Micro F1: 0.5504, Val Micro F1: 0.5408, Train Macro F1: 0.1683, Val Macro F1: 0.1665, Train Micro AUC: 0.9791, Val Micro AUC: 0.9731, Train Macro AUC: 0.8936, Val Macro AUC: 0.8822\n",
      "Epoch: 007, Train Recall@10: 0.7709, Val Recall@10: 0.7631, Train Micro F1: 0.5663, Val Micro F1: 0.5546, Train Macro F1: 0.1903, Val Macro F1: 0.1722, Train Micro AUC: 0.9804, Val Micro AUC: 0.9743, Train Macro AUC: 0.9012, Val Macro AUC: 0.8878\n",
      "Epoch: 008, Train Recall@10: 0.7833, Val Recall@10: 0.7741, Train Micro F1: 0.5848, Val Micro F1: 0.5704, Train Macro F1: 0.2023, Val Macro F1: 0.1761, Train Micro AUC: 0.9817, Val Micro AUC: 0.9756, Train Macro AUC: 0.9085, Val Macro AUC: 0.8933\n",
      "Epoch: 009, Train Recall@10: 0.7938, Val Recall@10: 0.7822, Train Micro F1: 0.5992, Val Micro F1: 0.5815, Train Macro F1: 0.2188, Val Macro F1: 0.1899, Train Micro AUC: 0.9830, Val Micro AUC: 0.9769, Train Macro AUC: 0.9162, Val Macro AUC: 0.8965\n",
      "Epoch: 010, Train Recall@10: 0.8014, Val Recall@10: 0.7859, Train Micro F1: 0.6094, Val Micro F1: 0.5878, Train Macro F1: 0.2354, Val Macro F1: 0.1979, Train Micro AUC: 0.9839, Val Micro AUC: 0.9775, Train Macro AUC: 0.9223, Val Macro AUC: 0.8991\n",
      "Epoch: 011, Train Recall@10: 0.8087, Val Recall@10: 0.7916, Train Micro F1: 0.6231, Val Micro F1: 0.5962, Train Macro F1: 0.2500, Val Macro F1: 0.2040, Train Micro AUC: 0.9847, Val Micro AUC: 0.9779, Train Macro AUC: 0.9275, Val Macro AUC: 0.8994\n",
      "Epoch: 012, Train Recall@10: 0.8144, Val Recall@10: 0.7939, Train Micro F1: 0.6315, Val Micro F1: 0.6020, Train Macro F1: 0.2653, Val Macro F1: 0.2116, Train Micro AUC: 0.9853, Val Micro AUC: 0.9781, Train Macro AUC: 0.9321, Val Macro AUC: 0.8996\n",
      "Epoch: 013, Train Recall@10: 0.8207, Val Recall@10: 0.7963, Train Micro F1: 0.6416, Val Micro F1: 0.6063, Train Macro F1: 0.2848, Val Macro F1: 0.2301, Train Micro AUC: 0.9861, Val Micro AUC: 0.9783, Train Macro AUC: 0.9373, Val Macro AUC: 0.8998\n",
      "Epoch: 014, Train Recall@10: 0.8284, Val Recall@10: 0.8003, Train Micro F1: 0.6517, Val Micro F1: 0.6135, Train Macro F1: 0.3036, Val Macro F1: 0.2432, Train Micro AUC: 0.9868, Val Micro AUC: 0.9786, Train Macro AUC: 0.9426, Val Macro AUC: 0.9015\n",
      "Epoch: 015, Train Recall@10: 0.8341, Val Recall@10: 0.8007, Train Micro F1: 0.6600, Val Micro F1: 0.6169, Train Macro F1: 0.3192, Val Macro F1: 0.2477, Train Micro AUC: 0.9874, Val Micro AUC: 0.9786, Train Macro AUC: 0.9473, Val Macro AUC: 0.9013\n",
      "Epoch: 016, Train Recall@10: 0.8399, Val Recall@10: 0.8016, Train Micro F1: 0.6682, Val Micro F1: 0.6193, Train Macro F1: 0.3396, Val Macro F1: 0.2525, Train Micro AUC: 0.9880, Val Micro AUC: 0.9786, Train Macro AUC: 0.9514, Val Macro AUC: 0.9016\n",
      "Epoch: 017, Train Recall@10: 0.8458, Val Recall@10: 0.8038, Train Micro F1: 0.6815, Val Micro F1: 0.6290, Train Macro F1: 0.3611, Val Macro F1: 0.2590, Train Micro AUC: 0.9886, Val Micro AUC: 0.9786, Train Macro AUC: 0.9546, Val Macro AUC: 0.9016\n",
      "Epoch: 018, Train Recall@10: 0.8518, Val Recall@10: 0.8024, Train Micro F1: 0.6909, Val Micro F1: 0.6319, Train Macro F1: 0.3741, Val Macro F1: 0.2636, Train Micro AUC: 0.9892, Val Micro AUC: 0.9784, Train Macro AUC: 0.9571, Val Macro AUC: 0.9008\n",
      "Epoch: 019, Train Recall@10: 0.8568, Val Recall@10: 0.8020, Train Micro F1: 0.6994, Val Micro F1: 0.6338, Train Macro F1: 0.3941, Val Macro F1: 0.2651, Train Micro AUC: 0.9897, Val Micro AUC: 0.9782, Train Macro AUC: 0.9595, Val Macro AUC: 0.8995\n",
      "Epoch: 020, Train Recall@10: 0.8620, Val Recall@10: 0.8003, Train Micro F1: 0.7073, Val Micro F1: 0.6337, Train Macro F1: 0.4164, Val Macro F1: 0.2676, Train Micro AUC: 0.9902, Val Micro AUC: 0.9779, Train Macro AUC: 0.9618, Val Macro AUC: 0.8987\n",
      "Epoch: 021, Train Recall@10: 0.8676, Val Recall@10: 0.8009, Train Micro F1: 0.7145, Val Micro F1: 0.6348, Train Macro F1: 0.4249, Val Macro F1: 0.2767, Train Micro AUC: 0.9907, Val Micro AUC: 0.9776, Train Macro AUC: 0.9639, Val Macro AUC: 0.8985\n",
      "Epoch: 022, Train Recall@10: 0.8724, Val Recall@10: 0.7994, Train Micro F1: 0.7249, Val Micro F1: 0.6379, Train Macro F1: 0.4553, Val Macro F1: 0.2775, Train Micro AUC: 0.9912, Val Micro AUC: 0.9772, Train Macro AUC: 0.9659, Val Macro AUC: 0.8978\n",
      "Early stopping at epoch 22\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(tfidf_model.parameters())\n",
    "prof_tfidf_ksi = train_model(tfidf_model, \n",
    "                             train_dataloader=train_dataloader,\n",
    "                             val_dataloader=val_dataloader,\n",
    "                             wikivec=wikivec,\n",
    "                             optimizer=optimizer,\n",
    "                             n_epochs=n_epochs, \n",
    "                             profile=profile, \n",
    "                             log_path=f'./log/{model_type}_ModifiedKSI_tfidf',\n",
    "                             device=DEVICE,\n",
    "                             init_hidden=True,\n",
    "                             early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "if save:\n",
    "    torch.save(tfidf_model, f'{dir}{model_type}_ModifiedKSI_tfidf_model.pt')\n",
    "if profile:\n",
    "    print(prof_tfidf_ksi.key_averages(group_by_stack_n=5).table(sort_by='self_cuda_time_total'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.7973, Test Micro F1: 0.6257, Test Macro F1: 0.2481, Test Micro AUC: 0.9795, Test Macro AUC: 0.8985\n"
     ]
    }
   ],
   "source": [
    "tt_recall_at_k, tt_micro_f1, tt_macro_f1, tt_micro_auc, tt_macro_auc, label_aucs_mod = test_model(tfidf_model, \n",
    "                                                                                                  test_dataloader, \n",
    "                                                                                                  wikivec,\n",
    "                                                                                                  device=DEVICE,\n",
    "                                                                                                  init_hidden=True)\n",
    "print(f'Test Recall@10: {tt_recall_at_k:.4f}, Test Micro F1: {tt_micro_f1:.4f}, Test Macro F1: {tt_macro_f1:.4f}' +\n",
    "      f', Test Micro AUC: {tt_micro_auc:.4f}, Test Macro AUC: {tt_macro_auc:.4f}')\n",
    "del tfidf_model\n",
    "gc.collect()\n",
    "if DEVICE == 'cuda':\n",
    "    torch.cuda.empty_cache()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "57e07979f6a7af2a0b0e861d549d9c40e5b4b1911b131063753718048dd868ae"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('deepl')",
   "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.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}