{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import gc\nimport math\nimport random\nimport joblib\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import train_test_split\n\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_skill = 13523\nmax_seq = 130\n\nbatch_size = 256\nembed_dim = 256\nnum_head = 8\nnum_layer = 2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load data"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ndtype = {'timestamp': 'int64', \n         'user_id': 'int32' ,\n         'content_id': 'int16',\n         'content_type_id': 'int8',\n         'answered_correctly':'int8'}\n\ntrain_df = pd.read_csv(\"/kaggle/input/riiid-test-answer-prediction/train.csv\", usecols=[1,2,3,4,7], dtype=dtype)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"question_df = pd.read_csv(\"../input/riiid-test-answer-prediction/questions.csv\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Preprocess"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = train_df[train_df.content_type_id == False]\ntrain_df = pd.merge(train_df, question_df[[\"question_id\", \"part\"]], left_on=\"content_id\", right_on=\"question_id\", how=\"left\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"group = train_df[['user_id', 'content_id', 'answered_correctly', 'part']].groupby('user_id').apply(lambda r: (\n        r['content_id'].values,\n        r['answered_correctly'].values,\n        r['part'].values))\n\ndel train_df\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RIIIDDataset(Dataset):\n    def __init__(self, group, n_skill, subset=\"train\", max_seq=100):\n        super(RIIIDDataset, self).__init__()\n        self.max_seq = max_seq\n        self.n_skill = n_skill # 13523\n        self.group = group\n        self.subset = subset\n        \n        # self.user_ids = [x for x in group.index]\n        self.user_ids = []\n        for user_id in group.index:\n            '''\n            q: question_id\n            qa: question answer correct or not\n            '''\n            content_id_, correct_, part_ = group[user_id]\n            if len(content_id_) < 2: # 2 interactions minimum\n                continue\n            self.user_ids.append(user_id) # user_ids indexes\n\n    def __len__(self):\n        return len(self.user_ids)\n\n    def __getitem__(self, index):\n        user_id = self.user_ids[index] # Pick a user\n        content_id_, correct_, part_ = self.group[user_id]\n        seq_len = len(content_id_)\n\n        content_id = np.zeros(self.max_seq, dtype=int)\n        part = np.zeros(self.max_seq, dtype=int)\n        correct = np.zeros(self.max_seq, dtype=int)\n\n        if seq_len >= self.max_seq:\n            if self.subset == \"train\":\n                if random.random() > 0.1:\n                    random_start_index = random.randint(0, seq_len - self.max_seq)\n                    '''\n                    Pick 100 questions, answers, prior question time, \n                    priori question explain from a random index\n                    '''\n                    end_index = random_start_index + self.max_seq\n                    content_id[:] = content_id_[random_start_index:end_index] \n                    part[:] = part_[random_start_index:end_index] \n                    correct[:] = correct_[random_start_index:end_index] \n                else:\n                    content_id[:] = content_id_[-self.max_seq:]\n                    part[:] = part_[-self.max_seq:]\n                    correct[:] = correct_[-self.max_seq:]\n            else:\n                content_id[:] = content_id_[-self.max_seq:]\n                part[:] = part_[-self.max_seq:]\n                correct[:] = correct_[-self.max_seq:]\n        else:\n            if random.random()>0.1:\n                seq_len = random.randint(2,seq_len)\n                content_id[-seq_len:] = content_id_[:seq_len]\n                part[-seq_len:] = part_[:seq_len]\n                correct[-seq_len:] = correct_[:seq_len]\n            else:\n                content_id[-seq_len:] = content_id_\n                part[-seq_len:] = part_\n                correct[-seq_len:] = correct_\n                \n        response = correct[:-1]\n\n        content_id = content_id[1:]\n        part = part[1:]\n        correct = correct[1:]\n\n        return response, content_id, part, correct","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train, val = train_test_split(group, test_size=0.1)\n\ntrain_dataset = RIIIDDataset(train, n_skill, max_seq=max_seq, subset=\"train\")\ntrain_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=8)\ndel train\n\nval_dataset = RIIIDDataset(val, n_skill, max_seq=max_seq,subset=\"val\")\nval_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=8)\ndel val\n\nprint(\"train dataset \", len(train_dataset), \" validation dataset \", len(val_dataset))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model and training"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTModel(nn.Module):\n    def __init__(self, n_skill, max_seq=100,\n                 embed_dim=512, num_head=8, num_layer=4):\n        super(SAINTModel, self).__init__()\n        self.n_skill = n_skill\n        \n        self.res_embedding = nn.Embedding(2, embed_dim)\n        self.res_pos = nn.Embedding(max_seq-1, embed_dim)\n\n        self.ex_embedding = nn.Embedding(n_skill+1, embed_dim)\n        self.ex_pos = nn.Embedding(max_seq-1, embed_dim)\n        \n        self.part_embedding = nn.Embedding(8, embed_dim)\n\n        self.transformer = nn.Transformer(d_model=embed_dim, nhead=num_head,\n                                        num_decoder_layers=num_layer, num_encoder_layers=num_layer)\n        \n        self.pred = nn.Linear(embed_dim, 1)\n\n        self._reset_parameters()\n    \n    def _reset_parameters(self):\n        r\"\"\"Initiate parameters in the model.\"\"\"\n        for p in self.parameters():\n            if p.dim() > 1:\n                nn.init.xavier_uniform_(p)\n\n    def forward(self, r, e, part):\n        device = r.device\n        src = self.ex_embedding(e)\n        src_mask = self.transformer.generate_square_subsequent_mask(e.size(1)).to(device)\n\n        ex_pos_id = torch.arange(e.size(1)).unsqueeze(0).to(device)\n        ex_pos = self.ex_pos(ex_pos_id)\n        src += ex_pos\n        \n        part = self.part_embedding(part)\n        src += part\n\n        tgt = self.res_embedding(r)\n        tgt_mask = self.transformer.generate_square_subsequent_mask(r.size(1)).to(device)\n        \n        res_pos_id = torch.arange(r.size(1)).unsqueeze(0).to(device)\n        res_pos = self.res_pos(res_pos_id)\n        tgt += res_pos\n\n        src = src.permute(1, 0, 2) # x: [bs, s_len, embed] => [s_len, bs, embed]\n        tgt = tgt.permute(1, 0, 2)\n        x = self.transformer(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)\n        x = x.permute(1, 0, 2)\n\n        output = self.pred(x)\n\n        return output.squeeze(-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = SAINTModel(n_skill, max_seq=max_seq, embed_dim=embed_dim, num_head=num_head, num_layer=num_layer)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_cosine_schedule_with_warmup(\n    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1\n):\n    \"\"\"\n    Create a schedule with a learning rate that decreases following the values of the cosine function between the\n    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the\n    initial lr set in the optimizer.\n\n    Args:\n        optimizer (:class:`~torch.optim.Optimizer`):\n            The optimizer for which to schedule the learning rate.\n        num_warmup_steps (:obj:`int`):\n            The number of steps for the warmup phase.\n        num_training_steps (:obj:`int`):\n            The total number of training steps.\n        num_cycles (:obj:`float`, `optional`, defaults to 0.5):\n            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0\n            following a half-cosine).\n        last_epoch (:obj:`int`, `optional`, defaults to -1):\n            The index of the last epoch when resuming training.\n\n    Return:\n        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n    \"\"\"\n\n    def lr_lambda(current_step):\n        if current_step < num_warmup_steps:\n            return float(current_step) / float(max(1, num_warmup_steps))\n        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))\n\n    return LambdaLR(optimizer, lr_lambda, last_epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)\n\nnum_warmup_steps = (len(train_dataset)/batch_size) * 4\nnum_training_steps = (len(train_dataset)/batch_size) * 50\nscheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)\ncriterion = nn.BCEWithLogitsLoss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel.to(device)\ncriterion.to(device)\nprint('model to gpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(model, train_iterator, optim, criterion, scheduler, device=\"cpu\"):\n    model.train()\n\n    train_loss = []\n    num_corrects = 0\n    num_total = 0\n    labels = []\n    outs = []\n\n    tbar = tqdm(train_iterator)\n    for item in tbar:  \n        r = item[0].to(device).long()\n        e = item[1].to(device).long()\n        p = item[2].to(device).long()\n        label = item[3].to(device).float()        \n\n        tgt_mask = (p != 0)\n\n        optim.zero_grad()\n        output = model(r, e, p)\n        \n        output_ = torch.masked_select(output, tgt_mask)\n        label_ = torch.masked_select(label, tgt_mask)\n        \n        output = output[:, -1]\n        label = label[:, -1]\n\n        loss = criterion(output_, label_)\n        loss.backward()\n        optim.step()\n        scheduler.step()\n        train_loss.append(loss.item())\n\n        pred = (torch.sigmoid(output) >= 0.5).long()\n        \n        num_corrects += (pred == label).sum().item()\n        num_total += len(label)\n\n        labels.extend(label.view(-1).data.cpu().numpy())\n        outs.extend(output.view(-1).data.cpu().numpy())\n\n        tbar.set_description('loss - {:.4f}'.format(loss))\n\n    acc = num_corrects / num_total\n    auc = roc_auc_score(labels, outs)\n    loss = np.average(train_loss)\n\n    return loss, acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def validation(model, val_iterator, criterion, device=\"cpu\"):\n    model.eval()\n\n    val_loss = []\n    num_corrects = 0\n    num_total = 0\n    labels = []\n    outs = []\n    \n    tbar = tqdm(val_iterator)\n    for item in tbar:\n        r = item[0].to(device).long()\n        e = item[1].to(device).long()\n        p = item[2].to(device).long()\n        label = item[3].to(device).float()\n        tgt_mask = (p != 0)\n\n        with torch.no_grad():\n            output = model(r, e, p)\n            \n        output = torch.masked_select(output, tgt_mask)\n        label = torch.masked_select(label, tgt_mask)\n\n        loss = criterion(output, label)\n        val_loss.append(loss.item())\n\n        pred = (torch.sigmoid(output) >= 0.5).long()\n\n        num_corrects += (pred == label).sum().item()\n        num_total += len(label)\n\n        labels.extend(label.squeeze(-1).data.cpu().numpy())\n        outs.extend(output.squeeze(-1).data.cpu().numpy())\n\n        tbar.set_description('loss - {:.4f}'.format(loss))\n\n    acc = num_corrects / num_total\n    auc = roc_auc_score(labels, outs)\n    loss = np.average(val_loss)\n\n    return loss, acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 50\n\nlast_auc = 0\n\nfor epoch in range(epochs):\n    loss, train_acc, train_auc = train_epoch(model, train_dataloader, optimizer, criterion, scheduler, device)\n    print(\"\\nepoch - {} train_loss - {:.3f} acc - {:.3f} auc - {:.4f}\".format(epoch, loss, train_acc, train_auc))\n\n    val_loss, val_acc, val_auc = validation(model, val_dataloader, criterion, device)\n    print(\"\\nepoch - {} vall_loss - {:.3f} acc - {:.3f} auc - {:.4f}\".format(epoch, val_loss, val_acc, val_auc))\n\n    # test_loss, test_acc, test_auc = validation(model, test_dataloader, criterion, device)\n    # print(\"\\nepoch - {} test_loss - {:.3f} acc - {:.3f} auc - {:.4f}\".format(epoch, test_loss, test_acc, test_auc))\n    if last_auc > val_auc:\n        print(\"early stop epoch \", epoch)\n        break\n    else:\n        last_auc = val_auc","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Testing"},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, group, test_df, n_skill=13523, max_seq=130):\n        super(TestDataset, self).__init__()\n        self.group = group\n        self.user_ids = [x for x in test_df[\"user_id\"].unique()]\n        self.test_df = test_df\n        self.n_skill = n_skill\n        self.max_seq = max_seq\n\n    def __len__(self):\n        return self.test_df.shape[0]\n\n    def __getitem__(self, index):\n        test_info = self.test_df.iloc[index]\n        user_id = test_info[\"user_id\"]\n        target_id = test_info[\"content_id\"]\n\n        content_id = np.zeros(self.max_seq, dtype=int)\n        task_id = np.zeros(self.max_seq, dtype=int)\n        part = np.zeros(self.max_seq, dtype=int)\n        et = np.zeros(self.max_seq, dtype=int)\n        lt = np.zeros(self.max_seq, dtype=int)\n        correct = np.zeros(self.max_seq, dtype=int)\n        \n        if user_id in self.group.index:\n            content_id_, correct_, part_ = self.group[user_id]\n            seq_len = len(content_id_)\n        \n            if seq_len >= self.max_seq:\n                content_id = content_id_[-self.max_seq:]\n                part = part_[-self.max_seq:]\n                correct = correct_[-self.max_seq:]\n\n            else:\n                content_id[-seq_len:] = content_id_\n                part[-seq_len:] = part_ \n                correct[-seq_len:] = correct_\n        \n        response = correct[1:]\n        \n        question = np.append(content_id[2:], [target_id])\n        part = np.append(part[2:], [test_info[\"part\"]])\n\n        return response, question, part","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import riiideducation\n\nenv = riiideducation.make_env()\niter_test = env.iter_test()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import psutil\n\nprev_test_df = None\n\nmodel.eval()\nfor (test_df, sample_prediction_df) in tqdm(iter_test):\n    if (prev_test_df is not None) & (psutil.virtual_memory().percent<90):\n        print(psutil.virtual_memory().percent)\n        prev_test_df['answered_correctly'] = eval(test_df['prior_group_answers_correct'].iloc[0])\n        \n        prev_test_df = prev_test_df[prev_test_df.content_type_id == False]\n        prev_test_df = pd.merge(prev_test_df, question_df[[\"question_id\", \"part\"]], \n                                left_on=\"content_id\", right_on=\"question_id\", how=\"left\")\n                \n        \n        #content_id, task_id, correct, lag_time, elapsed_time, had_e, part\n        prev_group = prev_test_df[['user_id', 'content_id', 'answered_correctly', 'part']].groupby('user_id').apply(lambda r: (\n            r['content_id'].values,\n            r['answered_correctly'].values,\n            r['part'].values))\n        \n        for prev_user_id in prev_group.index:\n            prev_group_content = prev_group[prev_user_id][0]\n            prev_group_ac = prev_group[prev_user_id][1]\n            prev_group_part = prev_group[prev_user_id][2]\n            if prev_user_id in group.index:\n                group[prev_user_id] = (np.append(group[prev_user_id][0], prev_group_content), \n                                       np.append(group[prev_user_id][1], prev_group_ac),\n                                       np.append(group[prev_user_id][2], prev_group_part))\n \n            else:\n                group[prev_user_id] = (prev_group_content, \n                                       prev_group_ac, \n                                       prev_group_part)\n            \n            if len(group[prev_user_id][0]) > max_seq:\n                new_group_content = group[prev_user_id][0][-max_seq:]\n                new_group_ac = group[prev_user_id][1][-max_seq:]\n                new_group_part = group[prev_user_id][2][-max_seq:]\n\n                group[prev_user_id] = (new_group_content, \n                                       new_group_ac, \n                                       new_group_part)\n                \n    prev_test_df = test_df.copy()\n    \n    test_df = test_df[test_df.content_type_id == False]\n    \n    test_df = pd.merge(test_df, question_df[[\"question_id\", \"part\"]], \n                       left_on=\"content_id\", right_on=\"question_id\", how=\"left\")\n    \n    test_dataset = TestDataset(group, test_df, max_seq=max_seq)\n    test_dataloader = DataLoader(test_dataset, batch_size=512, shuffle=False)\n    \n    outs = []\n\n    for item in tqdm(test_dataloader):\n        r = item[0].to(device).long()\n        c = item[1].to(device).long()\n        p = item[2].to(device).long()\n\n        with torch.no_grad():\n            output = model(r, c, p)\n        \n        output = torch.sigmoid(output)\n        output = output[:, -1]\n\n        outs.extend(output.view(-1).data.cpu().numpy())\n        \n    test_df['answered_correctly'] =  outs\n    \n    env.predict(test_df.loc[test_df['content_type_id'] == 0, ['row_id', 'answered_correctly']])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nplt.hist(outs)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}