HSUNEH_SAINT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m커널을 시작하지 못했습니다. \n",
      "\u001b[1;31mTraceback (most recent call last):\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/runpy.py\", line 196, in _run_module_as_main\n",
      "\u001b[1;31m    return _run_code(code, main_globals, None,\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/runpy.py\", line 86, in _run_code\n",
      "\u001b[1;31m    exec(code, run_globals)\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/ipykernel_launcher.py\", line 15, in <module>\n",
      "\u001b[1;31m    from ipykernel import kernelapp as app\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/ipykernel/kernelapp.py\", line 51, in <module>\n",
      "\u001b[1;31m    from .ipkernel import IPythonKernel\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 19, in <module>\n",
      "\u001b[1;31m    from .debugger import Debugger, _is_debugpy_available\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/ipykernel/debugger.py\", line 22, in <module>\n",
      "\u001b[1;31m    from debugpy.server import api  # noqa\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/debugpy/server/__init__.py\", line 9, in <module>\n",
      "\u001b[1;31m    import debugpy._vendored.force_pydevd  # noqa\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/debugpy/_vendored/force_pydevd.py\", line 37, in <module>\n",
      "\u001b[1;31m    pydevd_constants = import_module('_pydevd_bundle.pydevd_constants')\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/importlib/__init__.py\", line 126, in import_module\n",
      "\u001b[1;31m    return _bootstrap._gcd_import(name[level:], package, level)\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_constants.py\", line 365, in <module>\n",
      "\u001b[1;31m    from _pydev_imps._pydev_saved_modules import thread, threading\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/debugpy/_vendored/pydevd/_pydev_imps/_pydev_saved_modules.py\", line 93, in <module>\n",
      "\u001b[1;31m    import code as _code;    verify_shadowed.check(_code, ['compile_command', 'InteractiveInterpreter'])\n",
      "\u001b[1;31m  File \"/opt/conda/envs/LIGHTGCNenv/lib/python3.10/site-packages/debugpy/_vendored/pydevd/_pydev_imps/_pydev_saved_modules.py\", line 77, in check\n",
      "\u001b[1;31m    raise DebuggerInitializationError(msg)\n",
      "\u001b[1;31m_pydev_imps._pydev_saved_modules.DebuggerInitializationError: It was not possible to initialize the debugger due to a module name conflict.\n",
      "\u001b[1;31m\n",
      "\u001b[1;31mi.e.: the module \"code\" could not be imported because it is shadowed by:\n",
      "\u001b[1;31m/opt/ml/input/code/__init__.py\n",
      "\u001b[1;31mPlease rename this file/folder so that the original module from the standard library can be imported. \n",
      "\u001b[1;31m자세한 내용은 Jupyter <a href='command:jupyter.viewOutput'>로그</a>를 참조하세요."
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "from torch import nn\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from datetime import datetime\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'Self' from 'typing_extensions' (/opt/conda/lib/python3.8/site-packages/typing_extensions.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [150], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mpl\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/pytorch_lightning/__init__.py:35\u001b[0m\n\u001b[1;32m     32\u001b[0m     _logger\u001b[39m.\u001b[39mpropagate \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n\u001b[1;32m     34\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mlightning_lite\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutilities\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mseed\u001b[39;00m \u001b[39mimport\u001b[39;00m seed_everything  \u001b[39m# noqa: E402\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m \u001b[39mimport\u001b[39;00m Callback  \u001b[39m# noqa: E402\u001b[39;00m\n\u001b[1;32m     36\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcore\u001b[39;00m \u001b[39mimport\u001b[39;00m LightningDataModule, LightningModule  \u001b[39m# noqa: E402\u001b[39;00m\n\u001b[1;32m     37\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mtrainer\u001b[39;00m \u001b[39mimport\u001b[39;00m Trainer  \u001b[39m# noqa: E402\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/__init__.py:28\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mprediction_writer\u001b[39;00m \u001b[39mimport\u001b[39;00m BasePredictionWriter\n\u001b[1;32m     27\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mprogress\u001b[39;00m \u001b[39mimport\u001b[39;00m ProgressBarBase, RichProgressBar, TQDMProgressBar\n\u001b[0;32m---> 28\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mpruning\u001b[39;00m \u001b[39mimport\u001b[39;00m ModelPruning\n\u001b[1;32m     29\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mquantization\u001b[39;00m \u001b[39mimport\u001b[39;00m QuantizationAwareTraining\n\u001b[1;32m     30\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mrich_model_summary\u001b[39;00m \u001b[39mimport\u001b[39;00m RichModelSummary\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/pruning.py:31\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mpl\u001b[39;00m\n\u001b[1;32m     30\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallbacks\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcallback\u001b[39;00m \u001b[39mimport\u001b[39;00m Callback\n\u001b[0;32m---> 31\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcore\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mmodule\u001b[39;00m \u001b[39mimport\u001b[39;00m LightningModule\n\u001b[1;32m     32\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutilities\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexceptions\u001b[39;00m \u001b[39mimport\u001b[39;00m MisconfigurationException\n\u001b[1;32m     33\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutilities\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mrank_zero\u001b[39;00m \u001b[39mimport\u001b[39;00m rank_zero_debug, rank_zero_only\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/__init__.py:15\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[39m# Copyright The PyTorch Lightning team.\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[39m#\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[39m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[39m# See the License for the specific language governing permissions and\u001b[39;00m\n\u001b[1;32m     13\u001b[0m \u001b[39m# limitations under the License.\u001b[39;00m\n\u001b[0;32m---> 15\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcore\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mdatamodule\u001b[39;00m \u001b[39mimport\u001b[39;00m LightningDataModule\n\u001b[1;32m     16\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcore\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mmodule\u001b[39;00m \u001b[39mimport\u001b[39;00m LightningModule\n\u001b[1;32m     18\u001b[0m __all__ \u001b[39m=\u001b[39m [\u001b[39m\"\u001b[39m\u001b[39mLightningDataModule\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mLightningModule\u001b[39m\u001b[39m\"\u001b[39m]\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/datamodule.py:20\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtyping\u001b[39;00m \u001b[39mimport\u001b[39;00m Any, Dict, IO, List, Mapping, Optional, Sequence, Tuple, Union\n\u001b[1;32m     19\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorch\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutils\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mdata\u001b[39;00m \u001b[39mimport\u001b[39;00m DataLoader, Dataset, IterableDataset\n\u001b[0;32m---> 20\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtyping_extensions\u001b[39;00m \u001b[39mimport\u001b[39;00m Self\n\u001b[1;32m     22\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mpytorch_lightning\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mpl\u001b[39;00m\n\u001b[1;32m     23\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mlightning_lite\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutilities\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mtypes\u001b[39;00m \u001b[39mimport\u001b[39;00m _PATH\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'Self' from 'typing_extensions' (/opt/conda/lib/python3.8/site-packages/typing_extensions.py)"
     ]
    }
   ],
   "source": [
    "import pytorch_lightning as pl\n",
    "python -m ipykernel install --user --name recsys --display-name recsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config:\n",
    "    device = torch.device(\"cuda\")\n",
    "    MAX_SEQ = 100\n",
    "    EMBED_DIMS = 512\n",
    "    ENC_HEADS = DEC_HEADS = 8\n",
    "    NUM_ENCODER = NUM_DECODER = 4\n",
    "    BATCH_SIZE = 32\n",
    "    TRAIN_FILE = \"/opt/ml/input/data/train_data.csv\"\n",
    "    TEST_FILE = \"/opt/ml/input/data/test_data.csv\"\n",
    "    TOTAL_EXE = 13523\n",
    "    TOTAL_CAT = 10000\n",
    "    AUGMENTATION = 5\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1.데이터 불러오기"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "기존 SAINT : train의 Column 종류\n",
    "\n",
    "row_id, timestamp, user_id, content_id, content_type_id,\n",
    "\n",
    "task_container_id, user_answer, answered_correctly, prior_question_elapsed_time, prior_question_had_explanation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "우리 DKT\n",
    "\n",
    "userID,assessmentItemID,testId,answerCode,Timestamp,KnowledgeTag\n",
    "\n",
    "공통적으로 User id, timestamp 있고\n",
    "\n",
    "콘텐츠적으로 나누면 우리 dkt 는 3개정도 분류될거 같고,\n",
    "\n",
    "saint+이기 때문에 time 에 관해서 더 프로세싱하면 될거 같다.(elapsed, Lagged time 은 데이터상으로 구별하기 힘들다.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DKTDataset(Dataset):\n",
    "    def __init__(self, samples, max_seq):\n",
    "        super().__init__()\n",
    "        self.samples = samples\n",
    "        self.max_seq = max_seq\n",
    "        self.data = []\n",
    "        for id in self.samples.index:\n",
    "            exe_ids, answers, ela_time, categories = self.samples[id]\n",
    "            if len(exe_ids) > max_seq:\n",
    "                for l in range((len(exe_ids)+max_seq-1)//max_seq):\n",
    "                    self.data.append(\n",
    "                        (exe_ids[l:l+max_seq], answers[l:l+max_seq], ela_time[l:l+max_seq], categories[l:l+max_seq]))\n",
    "            elif len(exe_ids) < self.max_seq and len(exe_ids) > 50:\n",
    "                self.data.append((exe_ids, answers, ela_time, categories))\n",
    "            else:\n",
    "                continue\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        question_ids, answers, ela_time, exe_category = self.data[idx]\n",
    "        seq_len = len(question_ids)\n",
    "\n",
    "        exe_ids = np.zeros(self.max_seq, dtype=int)\n",
    "        ans = np.zeros(self.max_seq, dtype=int)\n",
    "        elapsed_time = np.zeros(self.max_seq, dtype=int)\n",
    "        exe_cat = np.zeros(self.max_seq, dtype=int)\n",
    "        if seq_len < self.max_seq:\n",
    "            exe_ids[-seq_len:] = question_ids\n",
    "            ans[-seq_len:] = answers\n",
    "            elapsed_time[-seq_len:] = ela_time\n",
    "            exe_cat[-seq_len:] = exe_category\n",
    "        else:\n",
    "            exe_ids[:] = question_ids[-self.max_seq:]\n",
    "            ans[:] = answers[-self.max_seq:]\n",
    "            elapsed_time[:] = ela_time[-self.max_seq:]\n",
    "            exe_cat[:] = exe_category[-self.max_seq:]\n",
    "\n",
    "        input_rtime = np.zeros(self.max_seq, dtype=int)\n",
    "        input_rtime = np.insert(elapsed_time, 0, 0)\n",
    "        input_rtime = np.delete(input_rtime, -1)\n",
    "\n",
    "        input = {\"input_ids\": exe_ids, \"input_rtime\": input_rtime.astype(\n",
    "            np.int), \"input_cat\": exe_cat}\n",
    "        return input, ans\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "전처리"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "'timestamp': 'int64',  -> Timestamp\n",
    "\n",
    "'user_id': 'int32',  -> userID\n",
    "\n",
    "'content_id': 'int16' -> assessmentItemID\n",
    "\n",
    "'answered_correctly': 'int8',   -> answerCode\n",
    "\n",
    "\n",
    "\n",
    "\"content_type_id\": \"int8\", -> 버림\n",
    "\n",
    "\n",
    "\n",
    "\"prior_question_elapsed_time\": \"float32\", -> Timestamp 변형\n",
    "\n",
    "\"task_container_id\": \"int16\" -> 누적 개수 (새로 생성 필요)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading csv.....\n",
      "shape of (train) dataframe : (2266586, 6)\n",
      "train elapsetime process  0/2266586 done \n",
      "train elapsetime process  100000/2266586 done \n",
      "train elapsetime process  200000/2266586 done \n",
      "train elapsetime process  300000/2266586 done \n",
      "train elapsetime process  400000/2266586 done \n",
      "train elapsetime process  500000/2266586 done \n",
      "train elapsetime process  600000/2266586 done \n",
      "train elapsetime process  700000/2266586 done \n",
      "train elapsetime process  800000/2266586 done \n",
      "train elapsetime process  900000/2266586 done \n",
      "train elapsetime process  1000000/2266586 done \n",
      "train elapsetime process  1100000/2266586 done \n",
      "train elapsetime process  1200000/2266586 done \n",
      "train elapsetime process  1300000/2266586 done \n",
      "train elapsetime process  1400000/2266586 done \n",
      "train elapsetime process  1500000/2266586 done \n",
      "train elapsetime process  1600000/2266586 done \n",
      "train elapsetime process  1700000/2266586 done \n",
      "train elapsetime process  1800000/2266586 done \n",
      "train elapsetime process  1900000/2266586 done \n",
      "train elapsetime process  2000000/2266586 done \n",
      "train elapsetime process  2100000/2266586 done \n",
      "train elapsetime process  2200000/2266586 done \n",
      "test elpasetime process  0/2266586 done \n",
      "test elpasetime process  100000/2266586 done \n",
      "test elpasetime process  200000/2266586 done \n",
      "train task process  100000/2266586 done \n",
      "train task process  200000/2266586 done \n",
      "train task process  300000/2266586 done \n",
      "train task process  400000/2266586 done \n",
      "train task process  500000/2266586 done \n",
      "train task process  600000/2266586 done \n",
      "train task process  700000/2266586 done \n",
      "train task process  800000/2266586 done \n",
      "train task process  900000/2266586 done \n",
      "train task process  1000000/2266586 done \n",
      "train task process  1100000/2266586 done \n",
      "train task process  1200000/2266586 done \n",
      "train task process  1300000/2266586 done \n",
      "train task process  1400000/2266586 done \n",
      "train task process  1500000/2266586 done \n",
      "train task process  1600000/2266586 done \n",
      "train task process  1700000/2266586 done \n",
      "train task process  1800000/2266586 done \n",
      "train task process  1900000/2266586 done \n",
      "train task process  2000000/2266586 done \n",
      "train task process  2100000/2266586 done \n",
      "train task process  2200000/2266586 done \n",
      "test task process  100000/2266586 done \n",
      "test task process  200000/2266586 done \n"
     ]
    }
   ],
   "source": [
    "print(\"loading csv.....\")\n",
    "train_df = pd.read_csv(Config.TRAIN_FILE)\n",
    "test_df = pd.read_csv(Config.TEST_FILE)\n",
    "print(\"shape of (train) dataframe :\", train_df.shape)\n",
    "\n",
    "# 기본적인 데이터 전처리 부분.\n",
    "# column : userID assessmentItemID  testId answerCode Timestamp KnowledgeTag grade\n",
    "\n",
    "train_df['grade'] = (train_df['assessmentItemID'].str[2]).astype('int')\n",
    "test_df['grade'] = (test_df['assessmentItemID'].str[2]).astype('int')\n",
    "train_df['testId'] = (train_df['assessmentItemID'].str[2] + train_df['assessmentItemID'].str[4:7])\n",
    "test_df['testId'] = test_df['assessmentItemID'].str[2] + test_df['assessmentItemID'].str[4:7]\n",
    "train_df[\"Timestamp\"] = pd.to_datetime(train_df[\"Timestamp\"])\n",
    "test_df[\"Timestamp\"] = pd.to_datetime(test_df[\"Timestamp\"])\n",
    "\n",
    "ItemID2idx = {v:k for k,v in enumerate(train_df['assessmentItemID'].unique())} # 9453개\n",
    "testId2idx = {v:k for k,v in enumerate(train_df['testId'].unique())} # 1536개\n",
    "Tag2idx = {v:k for k,v in enumerate(train_df['KnowledgeTag'].unique())} # 911개\n",
    "\n",
    "# args.n_item = train_df['assessmentItemID'].nunique()\n",
    "# args.n_test = train_df['testId'].nunique()\n",
    "# args.n_tag = train_df['KnowledgeTag'].nunique()\n",
    "# args.n_grade = train_df['grade'].nunique()\n",
    "\n",
    "train_df['assessmentItemID'] = train_df['assessmentItemID'].map(ItemID2idx)\n",
    "train_df['testId'] = train_df['testId'].map(testId2idx)\n",
    "train_df['KnowledgeTag'] = train_df['KnowledgeTag'].map(Tag2idx)\n",
    "\n",
    "test_df['assessmentItemID'] = test_df['assessmentItemID'].map(ItemID2idx)\n",
    "test_df['testId'] = test_df['testId'].map(testId2idx)\n",
    "test_df['KnowledgeTag'] = test_df['KnowledgeTag'].map(Tag2idx)\n",
    "\n",
    "\n",
    "\n",
    "#필요 Column 생성 - 'content_type_id'\n",
    "train_df['content_type_id'] = 0\n",
    "test_df['content_type_id'] = 0\n",
    "\n",
    "#필요 Column 생성 - 'prior_question_elapsed_time'\n",
    "train_df['prior_question_elapsed_time'] = 0\n",
    "test_df['prior_question_elapsed_time'] = 0\n",
    "\n",
    "for i in range(len(train_df)-1):\n",
    "    if train_df.at[i,'userID'] == train_df.at[(i+1),'userID']:\n",
    "        train_df.at[i,'prior_question_elapsed_time'] = (train_df.at[i+1,'Timestamp']-train_df.at[i,'Timestamp']).total_seconds()\n",
    "    if i %100000 == 0:\n",
    "        print(f'train elapsetime process  {i}/{len(train_df)} done ')\n",
    "\n",
    "for i in range(len(test_df)-1):\n",
    "    if test_df.at[i,'userID'] == test_df.at[(i+1),'userID']:\n",
    "        test_df.at[i,'prior_question_elapsed_time'] = (test_df.at[i+1,'Timestamp']-test_df.at[i,'Timestamp']).total_seconds()\n",
    "    if i %100000 == 0:\n",
    "        print(f'test elpasetime process  {i}/{len(train_df)} done ')\n",
    "\n",
    "# 필요 Column 생성 - 'task_container_id'\n",
    "train_df['task_container_id'] = 0\n",
    "test_df['task_container_id'] = 0\n",
    "\n",
    "for i in range(1,len(train_df)):\n",
    "    if train_df.at[i,'userID'] == train_df.at[i-1,'userID']:\n",
    "        train_df.at[i,'task_container_id'] = train_df.at[i-1,'task_container_id'] + 1\n",
    "    if i %100000 == 0:\n",
    "        print(f'train task process  {i}/{len(train_df)} done ')\n",
    "\n",
    "for i in range(1,len(test_df)):\n",
    "    if test_df.at[i,'userID'] == test_df.at[i-1,'userID']:\n",
    "        test_df.at[i,'task_container_id'] = test_df.at[i-1,'task_container_id'] + 1\n",
    "    if i %100000 == 0:\n",
    "        print(f'test task process  {i}/{len(train_df)} done ')\n",
    "\n",
    "# SAINT에 필요한 것만 남김\n",
    "train_df = train_df.drop(['testId','KnowledgeTag','grade'], axis=1)\n",
    "test_df = test_df.drop(['testId','KnowledgeTag','grade'], axis=1)\n",
    "\n",
    "# change to SAINT Input\n",
    "train_df  = train_df.rename(columns = {'userID': 'user_id', 'Timestamp':'timestamp', 'assessmentItemID' : 'content_id',\n",
    "                           'answerCode' : 'answered_correctly'  })\n",
    "test_df  = test_df.rename(columns = {'userID': 'user_id', 'Timestamp':'timestamp', 'assessmentItemID' : 'content_id',\n",
    "                           'answerCode' : 'answered_correctly'  })\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======origin length      : 2266586======\n",
      "START 1th AUGMENTATION\n",
      "END   1th AUGMENTATION\n",
      "START 2th AUGMENTATION\n",
      "END   2th AUGMENTATION\n",
      "START 3th AUGMENTATION\n",
      "END   3th AUGMENTATION\n",
      "START 4th AUGMENTATION\n",
      "END   4th AUGMENTATION\n",
      "START 5th AUGMENTATION\n",
      "END   5th AUGMENTATION\n",
      "======after augmentation : 13499046======\n"
     ]
    }
   ],
   "source": [
    "#data augmentation\n",
    "train_origin = train_df.copy()\n",
    "n= 1\n",
    "print(f'======origin length      : {len(train_df)}======')\n",
    "for i in range(Config.AUGMENTATION):\n",
    "    print(f'START {n}th AUGMENTATION')\n",
    "    tem = train_origin.drop_duplicates(subset = [\"user_id\"],keep = \"last\")\n",
    "    train_origin = train_origin.drop(index=tem.index)\n",
    "    train_origin['user_id'] += train_origin['user_id'].nunique()\n",
    "    train_df = pd.concat([train_df, train_origin], axis = 0)\n",
    "    print(f'END   {n}th AUGMENTATION')\n",
    "    n += 1\n",
    "print(f'======after augmentation : {len(train_df)}======')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no. of skills : 9454\n",
      "shape after exlusion: (13499046, 7)\n",
      "Grouping users...\n",
      "splitting\n",
      "train size:  (29754,) validation size:  (7439,)\n"
     ]
    }
   ],
   "source": [
    "train_df = train_df[train_df.content_type_id == 0]\n",
    "train_df.prior_question_elapsed_time.fillna(0, inplace=True)\n",
    "train_df.prior_question_elapsed_time /= 1000\n",
    "# train_df.prior_question_elapsed_time.clip(lower=0,upper=300,inplace=True)\n",
    "train_df.prior_question_elapsed_time = train_df.prior_question_elapsed_time.astype(\n",
    "    np.int)\n",
    "\n",
    "train_df = train_df.sort_values(\n",
    "    [\"timestamp\"], ascending=True).reset_index(drop=True)\n",
    "n_skills = train_df.content_id.nunique()\n",
    "print(\"no. of skills :\", n_skills)\n",
    "print(\"shape after exlusion:\", train_df.shape)\n",
    "\n",
    "# grouping based on user_id to get the data supplu\n",
    "print(\"Grouping users...\")\n",
    "group = train_df[[\"user_id\", \"content_id\", \"answered_correctly\", \"prior_question_elapsed_time\", \"task_container_id\"]]\\\n",
    "    .groupby(\"user_id\")\\\n",
    "    .apply(lambda r: (r.content_id.values, r.answered_correctly.values,\n",
    "                        r.prior_question_elapsed_time.values, r.task_container_id.values))\n",
    "# del train_df\n",
    "# gc.collect()\n",
    "print(\"splitting\")\n",
    "train, val = train_test_split(group, test_size=0.2)\n",
    "print(\"train size: \", train.shape, \"validation size: \", val.shape)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_dataset = DKTDataset(train, max_seq=Config.MAX_SEQ)\n",
    "val_dataset = DKTDataset(val, max_seq=Config.MAX_SEQ)\n",
    "\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                            batch_size=Config.BATCH_SIZE,\n",
    "                            num_workers=8,\n",
    "                            shuffle=True)\n",
    "val_loader = DataLoader(val_dataset,\n",
    "                        batch_size=Config.BATCH_SIZE,\n",
    "                        num_workers=8,\n",
    "                        shuffle=False)\n",
    "# del train_dataset, val_dataset\n",
    "# gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "train_loader, val_loader = get_dataloaders()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. 모델"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 멀티헤드어텐션"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StackedNMultiHeadAttention(nn.Module):\n",
    "    def __init__(self, n_stacks, n_dims, n_heads, seq_len, n_multihead=1, dropout=0.0):\n",
    "        super(StackedNMultiHeadAttention, self).__init__()\n",
    "        self.n_stacks = n_stacks\n",
    "        self.n_multihead = n_multihead\n",
    "        self.n_dims = n_dims\n",
    "        self.norm_layers = nn.LayerNorm(n_dims)\n",
    "        # n_stacks has n_multiheads each\n",
    "        self.multihead_layers = nn.ModuleList(n_stacks*[nn.ModuleList(n_multihead*[nn.MultiheadAttention(embed_dim=n_dims,\n",
    "                                                                                                         num_heads=n_heads,\n",
    "                                                                                                         dropout=dropout), ]), ])\n",
    "        self.ffn = nn.ModuleList(n_stacks*[FFN(n_dims)])\n",
    "        self.mask = torch.triu(torch.ones(seq_len, seq_len),\n",
    "                               diagonal=1).to(dtype=torch.bool)\n",
    "\n",
    "    def forward(self, input_q, input_k, input_v, encoder_output=None, break_layer=None):\n",
    "        for stack in range(self.n_stacks):\n",
    "            for multihead in range(self.n_multihead):\n",
    "                norm_q = self.norm_layers(input_q)\n",
    "                norm_k = self.norm_layers(input_k)\n",
    "                norm_v = self.norm_layers(input_v)\n",
    "                heads_output, _ = self.multihead_layers[stack][multihead](query=norm_q.permute(1, 0, 2),\n",
    "                                                                          key=norm_k.permute(\n",
    "                                                                              1, 0, 2),\n",
    "                                                                          value=norm_v.permute(\n",
    "                                                                              1, 0, 2),\n",
    "                                                                          attn_mask=self.mask.to(Config.device))\n",
    "                heads_output = heads_output.permute(1, 0, 2)\n",
    "                #assert encoder_output != None and break_layer is not None\n",
    "                if encoder_output != None and multihead == break_layer:\n",
    "                    assert break_layer <= multihead, \" break layer should be less than multihead layers and postive integer\"\n",
    "                    input_k = input_v = encoder_output\n",
    "                    input_q = input_q + heads_output\n",
    "                else:\n",
    "                    input_q = input_q + heads_output\n",
    "                    input_k = input_k + heads_output\n",
    "                    input_v = input_v + heads_output\n",
    "            last_norm = self.norm_layers(heads_output)\n",
    "            ffn_output = self.ffn[stack](last_norm)\n",
    "            ffn_output = ffn_output + heads_output\n",
    "        # after loops = input_q = input_k = input_v\n",
    "        return ffn_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 임베딩"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderEmbedding(nn.Module):\n",
    "    def __init__(self, n_exercises, n_categories, n_dims, seq_len):\n",
    "        super(EncoderEmbedding, self).__init__()\n",
    "        self.n_dims = n_dims\n",
    "        self.seq_len = seq_len\n",
    "        self.exercise_embed = nn.Embedding(n_exercises, n_dims)\n",
    "        self.category_embed = nn.Embedding(n_categories, n_dims)\n",
    "        self.position_embed = nn.Embedding(seq_len, n_dims)\n",
    "\n",
    "    def forward(self, exercises, categories):\n",
    "        e = self.exercise_embed(exercises)\n",
    "        c = self.category_embed(categories)\n",
    "        seq = torch.arange(self.seq_len, device=Config.device).unsqueeze(0)\n",
    "        p = self.position_embed(seq)\n",
    "        return p + c + e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderEmbedding(nn.Module):\n",
    "    def __init__(self, n_responses, n_dims, seq_len):\n",
    "        super(DecoderEmbedding, self).__init__()\n",
    "        self.n_dims = n_dims\n",
    "        self.seq_len = seq_len\n",
    "        self.response_embed = nn.Embedding(n_responses, n_dims)\n",
    "        self.time_embed = nn.Linear(1, n_dims, bias=False)\n",
    "        self.position_embed = nn.Embedding(seq_len, n_dims)\n",
    "\n",
    "    def forward(self, responses):\n",
    "        e = self.response_embed(responses)\n",
    "        seq = torch.arange(self.seq_len, device=Config.device).unsqueeze(0)\n",
    "        p = self.position_embed(seq)\n",
    "        return p + e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 메인"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlusSAINTModule(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        # n_encoder,n_detotal_responses,seq_len,max_time=300+1\n",
    "        super(PlusSAINTModule, self).__init__()\n",
    "        self.loss = nn.BCEWithLogitsLoss()\n",
    "        self.encoder_layer = StackedNMultiHeadAttention(n_stacks=Config.NUM_DECODER,\n",
    "                                                        n_dims=Config.EMBED_DIMS,\n",
    "                                                        n_heads=Config.DEC_HEADS,\n",
    "                                                        seq_len=Config.MAX_SEQ,\n",
    "                                                        n_multihead=1, dropout=0.0)\n",
    "        self.decoder_layer = StackedNMultiHeadAttention(n_stacks=Config.NUM_ENCODER,\n",
    "                                                        n_dims=Config.EMBED_DIMS,\n",
    "                                                        n_heads=Config.ENC_HEADS,\n",
    "                                                        seq_len=Config.MAX_SEQ,\n",
    "                                                        n_multihead=2, dropout=0.0)\n",
    "        self.encoder_embedding = EncoderEmbedding(n_exercises=Config.TOTAL_EXE,\n",
    "                                                  n_categories=Config.TOTAL_CAT,\n",
    "                                                  n_dims=Config.EMBED_DIMS, seq_len=Config.MAX_SEQ)\n",
    "        self.decoder_embedding = DecoderEmbedding(\n",
    "            n_responses=3, n_dims=Config.EMBED_DIMS, seq_len=Config.MAX_SEQ)\n",
    "        self.elapsed_time = nn.Linear(1, Config.EMBED_DIMS)\n",
    "        self.fc = nn.Linear(Config.EMBED_DIMS, 1)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        enc = self.encoder_embedding(\n",
    "            exercises=x[\"input_ids\"], categories=x['input_cat'])\n",
    "        dec = self.decoder_embedding(responses=y)\n",
    "        elapsed_time = x[\"input_rtime\"].unsqueeze(-1).float()\n",
    "        ela_time = self.elapsed_time(elapsed_time)\n",
    "        dec = dec + ela_time\n",
    "        # this encoder\n",
    "        encoder_output = self.encoder_layer(input_k=enc,\n",
    "                                            input_q=enc,\n",
    "                                            input_v=enc)\n",
    "        #this is decoder\n",
    "        decoder_output = self.decoder_layer(input_k=dec,\n",
    "                                            input_q=dec,\n",
    "                                            input_v=dec,\n",
    "                                            encoder_output=encoder_output,\n",
    "                                            break_layer=1)\n",
    "        # fully connected layer\n",
    "        out = self.fc(decoder_output)\n",
    "        return out.squeeze()\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters())\n",
    "\n",
    "    def training_step(self, batch, batch_ids):\n",
    "        input, labels = batch\n",
    "        target_mask = (input[\"input_ids\"] != 0)\n",
    "        out = self(input, labels)\n",
    "        loss = self.loss(out.float(), labels.float())\n",
    "        out = torch.masked_select(out, target_mask)\n",
    "        out = torch.sigmoid(out)\n",
    "        labels = torch.masked_select(labels, target_mask)\n",
    "        self.log(\"train_loss\", loss, on_step=True, prog_bar=True)\n",
    "        return {\"loss\": loss, \"outs\": out, \"labels\": labels}\n",
    "\n",
    "    def training_epoch_end(self, training_ouput):\n",
    "        out = np.concatenate([i[\"outs\"].cpu().detach().numpy()\n",
    "                              for i in training_ouput]).reshape(-1)\n",
    "        labels = np.concatenate([i[\"labels\"].cpu().detach().numpy()\n",
    "                                 for i in training_ouput]).reshape(-1)\n",
    "        auc = roc_auc_score(labels, out)\n",
    "        self.print(\"train auc\", auc)\n",
    "        self.log(\"train_auc\", auc)\n",
    "\n",
    "    def validation_step(self, batch, batch_ids):\n",
    "        input, labels = batch\n",
    "        target_mask = (input[\"input_ids\"] != 0)\n",
    "        out = self(input, labels)\n",
    "        loss = self.loss(out.float(), labels.float())\n",
    "        out = torch.masked_select(out, target_mask)\n",
    "        out = torch.sigmoid(out)\n",
    "        labels = torch.masked_select(labels, target_mask)\n",
    "        self.log(\"val_loss\", loss, on_step=True, prog_bar=True)\n",
    "        output = {\"outs\": out, \"labels\": labels}\n",
    "        return {\"val_loss\": loss, \"outs\": out, \"labels\": labels}\n",
    "\n",
    "    def validation_epoch_end(self, validation_ouput):\n",
    "        out = np.concatenate([i[\"outs\"].cpu().detach().numpy()\n",
    "                              for i in validation_ouput]).reshape(-1)\n",
    "        labels = np.concatenate([i[\"labels\"].cpu().detach().numpy()\n",
    "                                 for i in validation_ouput]).reshape(-1)\n",
    "        auc = roc_auc_score(labels, out)\n",
    "        self.print(\"val auc\", auc)\n",
    "        self.log(\"val_auc\", auc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saint_plus = PlusSAINTModule()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 학습"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(gpus=-1, max_epochs=5, progress_bar_refresh_rate=21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(model=saint_plus,\n",
    "            train_dataloader=train_loader,\n",
    "            val_dataloaders=[val_loader, ])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "saintenv",
   "language": "python",
   "name": "saintenv"
  },
  "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.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "96d24c1cd5bc0db60720ab4f882be5f45b591518c3be4d6e6a35957a64370f34"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}