KSI_CNN.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, CNN\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",
    "batch_size = 32\n",
    "n_epochs = 25\n",
    "save = True\n",
    "profile = False\n",
    "model_type = 'CNN'\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",
       "CNN                                      --                        --\n",
       "├─Embedding: 1-1                         [32, 2455, 100]           4,796,200\n",
       "├─Dropout: 1-2                           [32, 2455, 100]           --\n",
       "├─Conv1d: 1-3                            [32, 100, 2453]           30,100\n",
       "├─Conv1d: 1-4                            [32, 100, 2452]           40,100\n",
       "├─Conv1d: 1-5                            [32, 100, 2451]           50,100\n",
       "├─Linear: 1-6                            [32, 344]                 103,544\n",
       "==========================================================================================\n",
       "Total params: 5,020,044\n",
       "Trainable params: 5,020,044\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 9.60\n",
       "==========================================================================================\n",
       "Input size (MB): 1.87\n",
       "Forward/backward pass size (MB): 251.25\n",
       "Params size (MB): 20.08\n",
       "Estimated Total Size (MB): 273.20\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "base_model = CNN(n_words, n_wiki, n_embedding)\n",
    "base_model = base_model.to(DEVICE)\n",
    "base_summary = summary(base_model, [(batch_size, avg_note_size), (batch_size, n_vocab)], 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.6604, Val Recall@10: 0.6532, Train Micro F1: 0.5292, Val Micro F1: 0.5207, Train Macro F1: 0.0579, Val Macro F1: 0.0682, Train Micro AUC: 0.9578, Val Micro AUC: 0.9468, Train Macro AUC: 0.7287, Val Macro AUC: 0.7194\n",
      "Epoch: 002, Train Recall@10: 0.7283, Val Recall@10: 0.7156, Train Micro F1: 0.5933, Val Micro F1: 0.5803, Train Macro F1: 0.0950, Val Macro F1: 0.1080, Train Micro AUC: 0.9696, Val Micro AUC: 0.9595, Train Macro AUC: 0.8358, Val Macro AUC: 0.7841\n",
      "Epoch: 003, Train Recall@10: 0.7642, Val Recall@10: 0.7477, Train Micro F1: 0.6309, Val Micro F1: 0.6136, Train Macro F1: 0.1322, Val Macro F1: 0.1511, Train Micro AUC: 0.9749, Val Micro AUC: 0.9646, Train Macro AUC: 0.8968, Val Macro AUC: 0.8129\n",
      "Epoch: 004, Train Recall@10: 0.7843, Val Recall@10: 0.7647, Train Micro F1: 0.6521, Val Micro F1: 0.6293, Train Macro F1: 0.1557, Val Macro F1: 0.1709, Train Micro AUC: 0.9784, Val Micro AUC: 0.9672, Train Macro AUC: 0.9326, Val Macro AUC: 0.8302\n",
      "Epoch: 005, Train Recall@10: 0.7970, Val Recall@10: 0.7718, Train Micro F1: 0.6658, Val Micro F1: 0.6379, Train Macro F1: 0.1777, Val Macro F1: 0.1837, Train Micro AUC: 0.9809, Val Micro AUC: 0.9690, Train Macro AUC: 0.9505, Val Macro AUC: 0.8450\n",
      "Epoch: 006, Train Recall@10: 0.8069, Val Recall@10: 0.7757, Train Micro F1: 0.6755, Val Micro F1: 0.6420, Train Macro F1: 0.2078, Val Macro F1: 0.1958, Train Micro AUC: 0.9827, Val Micro AUC: 0.9700, Train Macro AUC: 0.9604, Val Macro AUC: 0.8515\n",
      "Epoch: 007, Train Recall@10: 0.8157, Val Recall@10: 0.7808, Train Micro F1: 0.6841, Val Micro F1: 0.6461, Train Macro F1: 0.2380, Val Macro F1: 0.2088, Train Micro AUC: 0.9843, Val Micro AUC: 0.9709, Train Macro AUC: 0.9663, Val Macro AUC: 0.8567\n",
      "Epoch: 008, Train Recall@10: 0.8238, Val Recall@10: 0.7839, Train Micro F1: 0.6930, Val Micro F1: 0.6494, Train Macro F1: 0.2663, Val Macro F1: 0.2167, Train Micro AUC: 0.9858, Val Micro AUC: 0.9717, Train Macro AUC: 0.9707, Val Macro AUC: 0.8595\n",
      "Epoch: 009, Train Recall@10: 0.8345, Val Recall@10: 0.7892, Train Micro F1: 0.7033, Val Micro F1: 0.6534, Train Macro F1: 0.3237, Val Macro F1: 0.2274, Train Micro AUC: 0.9872, Val Micro AUC: 0.9727, Train Macro AUC: 0.9740, Val Macro AUC: 0.8586\n",
      "Epoch: 010, Train Recall@10: 0.8431, Val Recall@10: 0.7891, Train Micro F1: 0.7110, Val Micro F1: 0.6548, Train Macro F1: 0.3869, Val Macro F1: 0.2338, Train Micro AUC: 0.9884, Val Micro AUC: 0.9731, Train Macro AUC: 0.9770, Val Macro AUC: 0.8591\n",
      "Epoch: 011, Train Recall@10: 0.8519, Val Recall@10: 0.7936, Train Micro F1: 0.7188, Val Micro F1: 0.6587, Train Macro F1: 0.4382, Val Macro F1: 0.2430, Train Micro AUC: 0.9895, Val Micro AUC: 0.9738, Train Macro AUC: 0.9796, Val Macro AUC: 0.8615\n",
      "Epoch: 012, Train Recall@10: 0.8586, Val Recall@10: 0.7928, Train Micro F1: 0.7258, Val Micro F1: 0.6558, Train Macro F1: 0.4881, Val Macro F1: 0.2463, Train Micro AUC: 0.9903, Val Micro AUC: 0.9737, Train Macro AUC: 0.9820, Val Macro AUC: 0.8598\n",
      "Epoch: 013, Train Recall@10: 0.8673, Val Recall@10: 0.7951, Train Micro F1: 0.7344, Val Micro F1: 0.6573, Train Macro F1: 0.5477, Val Macro F1: 0.2538, Train Micro AUC: 0.9912, Val Micro AUC: 0.9738, Train Macro AUC: 0.9841, Val Macro AUC: 0.8594\n",
      "Epoch: 014, Train Recall@10: 0.8748, Val Recall@10: 0.7946, Train Micro F1: 0.7415, Val Micro F1: 0.6535, Train Macro F1: 0.6010, Val Macro F1: 0.2570, Train Micro AUC: 0.9919, Val Micro AUC: 0.9738, Train Macro AUC: 0.9858, Val Macro AUC: 0.8584\n",
      "Epoch: 015, Train Recall@10: 0.8807, Val Recall@10: 0.7943, Train Micro F1: 0.7492, Val Micro F1: 0.6548, Train Macro F1: 0.6310, Val Macro F1: 0.2698, Train Micro AUC: 0.9926, Val Micro AUC: 0.9736, Train Macro AUC: 0.9873, Val Macro AUC: 0.8585\n",
      "Epoch: 016, Train Recall@10: 0.8878, Val Recall@10: 0.7958, Train Micro F1: 0.7580, Val Micro F1: 0.6549, Train Macro F1: 0.6628, Val Macro F1: 0.2727, Train Micro AUC: 0.9933, Val Micro AUC: 0.9736, Train Macro AUC: 0.9888, Val Macro AUC: 0.8579\n",
      "Epoch: 017, Train Recall@10: 0.8941, Val Recall@10: 0.7933, Train Micro F1: 0.7640, Val Micro F1: 0.6543, Train Macro F1: 0.6823, Val Macro F1: 0.2705, Train Micro AUC: 0.9938, Val Micro AUC: 0.9732, Train Macro AUC: 0.9899, Val Macro AUC: 0.8571\n",
      "Epoch: 018, Train Recall@10: 0.9001, Val Recall@10: 0.7922, Train Micro F1: 0.7705, Val Micro F1: 0.6515, Train Macro F1: 0.7010, Val Macro F1: 0.2681, Train Micro AUC: 0.9943, Val Micro AUC: 0.9733, Train Macro AUC: 0.9910, Val Macro AUC: 0.8547\n",
      "Epoch: 019, Train Recall@10: 0.9066, Val Recall@10: 0.7901, Train Micro F1: 0.7812, Val Micro F1: 0.6503, Train Macro F1: 0.7222, Val Macro F1: 0.2698, Train Micro AUC: 0.9948, Val Micro AUC: 0.9726, Train Macro AUC: 0.9919, Val Macro AUC: 0.8505\n",
      "Epoch: 020, Train Recall@10: 0.9120, Val Recall@10: 0.7900, Train Micro F1: 0.7874, Val Micro F1: 0.6514, Train Macro F1: 0.7362, Val Macro F1: 0.2750, Train Micro AUC: 0.9953, Val Micro AUC: 0.9727, Train Macro AUC: 0.9927, Val Macro AUC: 0.8528\n",
      "Epoch: 021, Train Recall@10: 0.9182, Val Recall@10: 0.7906, Train Micro F1: 0.7929, Val Micro F1: 0.6472, Train Macro F1: 0.7427, Val Macro F1: 0.2723, Train Micro AUC: 0.9957, Val Micro AUC: 0.9726, Train Macro AUC: 0.9934, Val Macro AUC: 0.8511\n",
      "Early stopping at epoch 20\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",
    "                        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.7959, Test Micro F1: 0.6550, Test Macro F1: 0.2527, Test Micro AUC: 0.9753, Test Macro AUC: 0.8500\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",
    "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",
       "CNN                                      --                        --\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                           [32, 2455, 100]           --\n",
       "├─Conv1d: 1-4                            [32, 100, 2453]           30,100\n",
       "├─Conv1d: 1-5                            [32, 100, 2452]           40,100\n",
       "├─Conv1d: 1-6                            [32, 100, 2451]           50,100\n",
       "├─Linear: 1-7                            [32, 344]                 103,544\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,247,645\n",
       "Trainable params: 6,247,645\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 9.63\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 268.95\n",
       "Params size (MB): 24.99\n",
       "Estimated Total Size (MB): 312.56\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ksi = KSI(n_embedding, n_vocab)\n",
    "ksi.to(DEVICE)\n",
    "model = CNN(n_words, n_wiki, n_embedding, 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.7342, Val Recall@10: 0.7240, Train Micro F1: 0.5626, Val Micro F1: 0.5508, Train Macro F1: 0.1101, Val Macro F1: 0.1261, Train Micro AUC: 0.9733, Val Micro AUC: 0.9648, Train Macro AUC: 0.8603, Val Macro AUC: 0.8419\n",
      "Epoch: 002, Train Recall@10: 0.7826, Val Recall@10: 0.7668, Train Micro F1: 0.6270, Val Micro F1: 0.6073, Train Macro F1: 0.1825, Val Macro F1: 0.1984, Train Micro AUC: 0.9801, Val Micro AUC: 0.9716, Train Macro AUC: 0.9123, Val Macro AUC: 0.8735\n",
      "Epoch: 003, Train Recall@10: 0.8036, Val Recall@10: 0.7793, Train Micro F1: 0.6557, Val Micro F1: 0.6294, Train Macro F1: 0.2318, Val Macro F1: 0.2234, Train Micro AUC: 0.9830, Val Micro AUC: 0.9738, Train Macro AUC: 0.9385, Val Macro AUC: 0.8843\n",
      "Epoch: 004, Train Recall@10: 0.8163, Val Recall@10: 0.7860, Train Micro F1: 0.6722, Val Micro F1: 0.6401, Train Macro F1: 0.2689, Val Macro F1: 0.2516, Train Micro AUC: 0.9850, Val Micro AUC: 0.9748, Train Macro AUC: 0.9547, Val Macro AUC: 0.8895\n",
      "Epoch: 005, Train Recall@10: 0.8282, Val Recall@10: 0.7933, Train Micro F1: 0.6853, Val Micro F1: 0.6447, Train Macro F1: 0.3141, Val Macro F1: 0.2577, Train Micro AUC: 0.9866, Val Micro AUC: 0.9755, Train Macro AUC: 0.9656, Val Macro AUC: 0.8931\n",
      "Epoch: 006, Train Recall@10: 0.8398, Val Recall@10: 0.7970, Train Micro F1: 0.6988, Val Micro F1: 0.6488, Train Macro F1: 0.3618, Val Macro F1: 0.2677, Train Micro AUC: 0.9880, Val Micro AUC: 0.9758, Train Macro AUC: 0.9723, Val Macro AUC: 0.8944\n",
      "Epoch: 007, Train Recall@10: 0.8490, Val Recall@10: 0.7969, Train Micro F1: 0.7098, Val Micro F1: 0.6507, Train Macro F1: 0.4096, Val Macro F1: 0.2817, Train Micro AUC: 0.9892, Val Micro AUC: 0.9755, Train Macro AUC: 0.9767, Val Macro AUC: 0.8929\n",
      "Epoch: 008, Train Recall@10: 0.8573, Val Recall@10: 0.7944, Train Micro F1: 0.7172, Val Micro F1: 0.6486, Train Macro F1: 0.4454, Val Macro F1: 0.2900, Train Micro AUC: 0.9902, Val Micro AUC: 0.9751, Train Macro AUC: 0.9800, Val Macro AUC: 0.8891\n",
      "Epoch: 009, Train Recall@10: 0.8678, Val Recall@10: 0.7917, Train Micro F1: 0.7279, Val Micro F1: 0.6478, Train Macro F1: 0.5044, Val Macro F1: 0.2884, Train Micro AUC: 0.9913, Val Micro AUC: 0.9744, Train Macro AUC: 0.9829, Val Macro AUC: 0.8852\n",
      "Epoch: 010, Train Recall@10: 0.8763, Val Recall@10: 0.7908, Train Micro F1: 0.7360, Val Micro F1: 0.6442, Train Macro F1: 0.5499, Val Macro F1: 0.2959, Train Micro AUC: 0.9922, Val Micro AUC: 0.9738, Train Macro AUC: 0.9853, Val Macro AUC: 0.8825\n",
      "Epoch: 011, Train Recall@10: 0.8844, Val Recall@10: 0.7873, Train Micro F1: 0.7453, Val Micro F1: 0.6414, Train Macro F1: 0.6018, Val Macro F1: 0.2895, Train Micro AUC: 0.9930, Val Micro AUC: 0.9729, Train Macro AUC: 0.9874, Val Macro AUC: 0.8769\n",
      "Early stopping at epoch 10\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",
    "                       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.7946, Test Micro F1: 0.6482, Test Macro F1: 0.2571, Test Micro AUC: 0.9771, Test Macro AUC: 0.8923\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",
    "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": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "CNN                                      --                        --\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                           [32, 2455, 100]           --\n",
       "├─Conv1d: 1-4                            [32, 100, 2453]           30,100\n",
       "├─Conv1d: 1-5                            [32, 100, 2452]           40,100\n",
       "├─Conv1d: 1-6                            [32, 100, 2451]           50,100\n",
       "├─Linear: 1-7                            [32, 344]                 103,544\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,247,647\n",
       "Trainable params: 6,247,647\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 9.63\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 1340.95\n",
       "Params size (MB): 24.99\n",
       "Estimated Total Size (MB): 1384.57\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_ksi = ModifiedKSI(n_embedding, n_vocab)\n",
    "mod_ksi.to(DEVICE)\n",
    "mod_model = CNN(n_words, n_wiki, n_embedding, 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": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.7488, Val Recall@10: 0.7409, Train Micro F1: 0.5869, Val Micro F1: 0.5776, Train Macro F1: 0.1438, Val Macro F1: 0.1753, Train Micro AUC: 0.9756, Val Micro AUC: 0.9690, Train Macro AUC: 0.8802, Val Macro AUC: 0.8805\n",
      "Epoch: 002, Train Recall@10: 0.7833, Val Recall@10: 0.7714, Train Micro F1: 0.6306, Val Micro F1: 0.6188, Train Macro F1: 0.1849, Val Macro F1: 0.2148, Train Micro AUC: 0.9802, Val Micro AUC: 0.9736, Train Macro AUC: 0.9105, Val Macro AUC: 0.8982\n",
      "Epoch: 003, Train Recall@10: 0.7993, Val Recall@10: 0.7873, Train Micro F1: 0.6495, Val Micro F1: 0.6348, Train Macro F1: 0.2218, Val Macro F1: 0.2409, Train Micro AUC: 0.9824, Val Micro AUC: 0.9755, Train Macro AUC: 0.9309, Val Macro AUC: 0.9037\n",
      "Epoch: 004, Train Recall@10: 0.8095, Val Recall@10: 0.7938, Train Micro F1: 0.6634, Val Micro F1: 0.6442, Train Macro F1: 0.2493, Val Macro F1: 0.2592, Train Micro AUC: 0.9841, Val Micro AUC: 0.9766, Train Macro AUC: 0.9482, Val Macro AUC: 0.9061\n",
      "Epoch: 005, Train Recall@10: 0.8176, Val Recall@10: 0.7988, Train Micro F1: 0.6735, Val Micro F1: 0.6490, Train Macro F1: 0.2766, Val Macro F1: 0.2682, Train Micro AUC: 0.9853, Val Micro AUC: 0.9772, Train Macro AUC: 0.9586, Val Macro AUC: 0.9074\n",
      "Epoch: 006, Train Recall@10: 0.8253, Val Recall@10: 0.8016, Train Micro F1: 0.6824, Val Micro F1: 0.6543, Train Macro F1: 0.3078, Val Macro F1: 0.2844, Train Micro AUC: 0.9864, Val Micro AUC: 0.9777, Train Macro AUC: 0.9659, Val Macro AUC: 0.9103\n",
      "Epoch: 007, Train Recall@10: 0.8325, Val Recall@10: 0.8057, Train Micro F1: 0.6908, Val Micro F1: 0.6563, Train Macro F1: 0.3486, Val Macro F1: 0.2892, Train Micro AUC: 0.9874, Val Micro AUC: 0.9780, Train Macro AUC: 0.9711, Val Macro AUC: 0.9093\n",
      "Epoch: 008, Train Recall@10: 0.8400, Val Recall@10: 0.8056, Train Micro F1: 0.6979, Val Micro F1: 0.6565, Train Macro F1: 0.3848, Val Macro F1: 0.2972, Train Micro AUC: 0.9883, Val Micro AUC: 0.9782, Train Macro AUC: 0.9749, Val Macro AUC: 0.9084\n",
      "Epoch: 009, Train Recall@10: 0.8453, Val Recall@10: 0.8065, Train Micro F1: 0.7051, Val Micro F1: 0.6574, Train Macro F1: 0.4342, Val Macro F1: 0.3016, Train Micro AUC: 0.9891, Val Micro AUC: 0.9783, Train Macro AUC: 0.9778, Val Macro AUC: 0.9095\n",
      "Epoch: 010, Train Recall@10: 0.8527, Val Recall@10: 0.8067, Train Micro F1: 0.7142, Val Micro F1: 0.6591, Train Macro F1: 0.4756, Val Macro F1: 0.3106, Train Micro AUC: 0.9900, Val Micro AUC: 0.9784, Train Macro AUC: 0.9805, Val Macro AUC: 0.9072\n",
      "Epoch: 011, Train Recall@10: 0.8603, Val Recall@10: 0.8091, Train Micro F1: 0.7220, Val Micro F1: 0.6595, Train Macro F1: 0.5135, Val Macro F1: 0.3135, Train Micro AUC: 0.9908, Val Micro AUC: 0.9785, Train Macro AUC: 0.9828, Val Macro AUC: 0.9080\n",
      "Epoch: 012, Train Recall@10: 0.8674, Val Recall@10: 0.8087, Train Micro F1: 0.7311, Val Micro F1: 0.6580, Train Macro F1: 0.5639, Val Macro F1: 0.3206, Train Micro AUC: 0.9915, Val Micro AUC: 0.9783, Train Macro AUC: 0.9848, Val Macro AUC: 0.9054\n",
      "Epoch: 013, Train Recall@10: 0.8743, Val Recall@10: 0.8078, Train Micro F1: 0.7400, Val Micro F1: 0.6587, Train Macro F1: 0.6024, Val Macro F1: 0.3230, Train Micro AUC: 0.9922, Val Micro AUC: 0.9782, Train Macro AUC: 0.9864, Val Macro AUC: 0.9030\n",
      "Epoch: 014, Train Recall@10: 0.8805, Val Recall@10: 0.8063, Train Micro F1: 0.7485, Val Micro F1: 0.6559, Train Macro F1: 0.6337, Val Macro F1: 0.3269, Train Micro AUC: 0.9928, Val Micro AUC: 0.9779, Train Macro AUC: 0.9879, Val Macro AUC: 0.8966\n",
      "Epoch: 015, Train Recall@10: 0.8888, Val Recall@10: 0.8051, Train Micro F1: 0.7585, Val Micro F1: 0.6566, Train Macro F1: 0.6601, Val Macro F1: 0.3250, Train Micro AUC: 0.9935, Val Micro AUC: 0.9777, Train Macro AUC: 0.9893, Val Macro AUC: 0.8962\n",
      "Epoch: 016, Train Recall@10: 0.8946, Val Recall@10: 0.8028, Train Micro F1: 0.7659, Val Micro F1: 0.6549, Train Macro F1: 0.6930, Val Macro F1: 0.3171, Train Micro AUC: 0.9940, Val Micro AUC: 0.9771, Train Macro AUC: 0.9904, Val Macro AUC: 0.8922\n",
      "Early stopping at epoch 15\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",
    "                           early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.8066, Test Micro F1: 0.6571, Test Macro F1: 0.3023, Test Micro AUC: 0.9799, Test Macro AUC: 0.9061\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",
    "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": 22,
   "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": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "CNN                                      --                        --\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                           [32, 2455, 100]           --\n",
       "├─Conv1d: 1-4                            [32, 100, 2453]           30,100\n",
       "├─Conv1d: 1-5                            [32, 100, 2452]           40,100\n",
       "├─Conv1d: 1-6                            [32, 100, 2451]           50,100\n",
       "├─Linear: 1-7                            [32, 344]                 103,544\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,247,647\n",
       "Trainable params: 6,247,647\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 9.63\n",
       "==========================================================================================\n",
       "Input size (MB): 18.62\n",
       "Forward/backward pass size (MB): 1340.95\n",
       "Params size (MB): 24.99\n",
       "Estimated Total Size (MB): 1384.57\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mod_ksi2 = ModifiedKSI(n_embedding, n_vocab)\n",
    "mod_ksi2.to(DEVICE)\n",
    "tfidf_model = CNN(n_words, n_wiki, n_embedding, 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": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Recall@10: 0.7539, Val Recall@10: 0.7474, Train Micro F1: 0.5916, Val Micro F1: 0.5822, Train Macro F1: 0.1546, Val Macro F1: 0.1834, Train Micro AUC: 0.9764, Val Micro AUC: 0.9700, Train Macro AUC: 0.8784, Val Macro AUC: 0.8774\n",
      "Epoch: 002, Train Recall@10: 0.7839, Val Recall@10: 0.7715, Train Micro F1: 0.6310, Val Micro F1: 0.6197, Train Macro F1: 0.1956, Val Macro F1: 0.2235, Train Micro AUC: 0.9807, Val Micro AUC: 0.9741, Train Macro AUC: 0.9125, Val Macro AUC: 0.8949\n",
      "Epoch: 003, Train Recall@10: 0.8003, Val Recall@10: 0.7870, Train Micro F1: 0.6491, Val Micro F1: 0.6331, Train Macro F1: 0.2255, Val Macro F1: 0.2452, Train Micro AUC: 0.9829, Val Micro AUC: 0.9759, Train Macro AUC: 0.9331, Val Macro AUC: 0.9023\n",
      "Epoch: 004, Train Recall@10: 0.8121, Val Recall@10: 0.7958, Train Micro F1: 0.6633, Val Micro F1: 0.6442, Train Macro F1: 0.2570, Val Macro F1: 0.2615, Train Micro AUC: 0.9846, Val Micro AUC: 0.9772, Train Macro AUC: 0.9478, Val Macro AUC: 0.9044\n",
      "Epoch: 005, Train Recall@10: 0.8215, Val Recall@10: 0.8010, Train Micro F1: 0.6742, Val Micro F1: 0.6498, Train Macro F1: 0.2861, Val Macro F1: 0.2692, Train Micro AUC: 0.9859, Val Micro AUC: 0.9779, Train Macro AUC: 0.9593, Val Macro AUC: 0.9060\n",
      "Epoch: 006, Train Recall@10: 0.8283, Val Recall@10: 0.8019, Train Micro F1: 0.6826, Val Micro F1: 0.6531, Train Macro F1: 0.3208, Val Macro F1: 0.2807, Train Micro AUC: 0.9870, Val Micro AUC: 0.9783, Train Macro AUC: 0.9666, Val Macro AUC: 0.9080\n",
      "Epoch: 007, Train Recall@10: 0.8361, Val Recall@10: 0.8041, Train Micro F1: 0.6905, Val Micro F1: 0.6541, Train Macro F1: 0.3517, Val Macro F1: 0.2853, Train Micro AUC: 0.9879, Val Micro AUC: 0.9786, Train Macro AUC: 0.9718, Val Macro AUC: 0.9075\n",
      "Epoch: 008, Train Recall@10: 0.8432, Val Recall@10: 0.8051, Train Micro F1: 0.6989, Val Micro F1: 0.6548, Train Macro F1: 0.3968, Val Macro F1: 0.2954, Train Micro AUC: 0.9888, Val Micro AUC: 0.9787, Train Macro AUC: 0.9757, Val Macro AUC: 0.9051\n",
      "Epoch: 009, Train Recall@10: 0.8494, Val Recall@10: 0.8065, Train Micro F1: 0.7065, Val Micro F1: 0.6565, Train Macro F1: 0.4371, Val Macro F1: 0.2977, Train Micro AUC: 0.9897, Val Micro AUC: 0.9787, Train Macro AUC: 0.9789, Val Macro AUC: 0.9033\n",
      "Epoch: 010, Train Recall@10: 0.8566, Val Recall@10: 0.8067, Train Micro F1: 0.7160, Val Micro F1: 0.6578, Train Macro F1: 0.4923, Val Macro F1: 0.3053, Train Micro AUC: 0.9905, Val Micro AUC: 0.9787, Train Macro AUC: 0.9815, Val Macro AUC: 0.9028\n",
      "Epoch: 011, Train Recall@10: 0.8644, Val Recall@10: 0.8058, Train Micro F1: 0.7242, Val Micro F1: 0.6548, Train Macro F1: 0.5307, Val Macro F1: 0.3125, Train Micro AUC: 0.9912, Val Micro AUC: 0.9786, Train Macro AUC: 0.9835, Val Macro AUC: 0.9006\n",
      "Epoch: 012, Train Recall@10: 0.8721, Val Recall@10: 0.8075, Train Micro F1: 0.7342, Val Micro F1: 0.6570, Train Macro F1: 0.5764, Val Macro F1: 0.3182, Train Micro AUC: 0.9920, Val Micro AUC: 0.9785, Train Macro AUC: 0.9856, Val Macro AUC: 0.8985\n",
      "Epoch: 013, Train Recall@10: 0.8783, Val Recall@10: 0.8042, Train Micro F1: 0.7426, Val Micro F1: 0.6546, Train Macro F1: 0.6208, Val Macro F1: 0.3185, Train Micro AUC: 0.9927, Val Micro AUC: 0.9782, Train Macro AUC: 0.9873, Val Macro AUC: 0.8988\n",
      "Epoch: 014, Train Recall@10: 0.8862, Val Recall@10: 0.8026, Train Micro F1: 0.7532, Val Micro F1: 0.6545, Train Macro F1: 0.6586, Val Macro F1: 0.3195, Train Micro AUC: 0.9933, Val Micro AUC: 0.9779, Train Macro AUC: 0.9888, Val Macro AUC: 0.8969\n",
      "Epoch: 015, Train Recall@10: 0.8935, Val Recall@10: 0.8020, Train Micro F1: 0.7625, Val Micro F1: 0.6498, Train Macro F1: 0.6904, Val Macro F1: 0.3182, Train Micro AUC: 0.9939, Val Micro AUC: 0.9777, Train Macro AUC: 0.9901, Val Macro AUC: 0.8952\n",
      "Epoch: 016, Train Recall@10: 0.9002, Val Recall@10: 0.7976, Train Micro F1: 0.7714, Val Micro F1: 0.6503, Train Macro F1: 0.7193, Val Macro F1: 0.3242, Train Micro AUC: 0.9944, Val Micro AUC: 0.9771, Train Macro AUC: 0.9912, Val Macro AUC: 0.8915\n",
      "Epoch: 017, Train Recall@10: 0.9066, Val Recall@10: 0.7966, Train Micro F1: 0.7805, Val Micro F1: 0.6485, Train Macro F1: 0.7394, Val Macro F1: 0.3193, Train Micro AUC: 0.9949, Val Micro AUC: 0.9768, Train Macro AUC: 0.9922, Val Macro AUC: 0.8909\n",
      "Early stopping at epoch 16\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",
    "                             early_stopping=early_stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Recall@10: 0.8060, Test Micro F1: 0.6548, Test Macro F1: 0.3156, Test Micro AUC: 0.9800, Test Macro AUC: 0.9003\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",
    "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
}