{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Main contribution\n\nThis is based on many great kernels of SAKT model (see Refs).\n\nThe main change is inspired by the current 2nd place @mamasinkgs 's comment: https://www.kaggle.com/c/riiid-test-answer-prediction/discussion/204801#1128342\n\n> I can't share detailed information, but I think no difficult technique is required to get to 0.810. What are truly necessary are stable validation, careful implementation for dataloader/model/training loop, and a stable system that avoids bug in inference.\n\nSo I dissect into the inference loop using @its7171 's `iter_env`'s simulation to check what happens. Since the test is given in a chronological order, there may only be a few interactions of a single user in one `test_df`, as such, we have to code the train loader to capture this. The illustration of this is in the inference part.\n\nSo the main change is the train loader:\n- The train loader now uses the `timestamp` information to cut the sequences. There is a timestamp thresh to determine whether in the current seqence, the last and the second last entries in train could be in the same batch if they are from test.\n- Additionally shifts the sequence as well so there are much much more sequences in 1 epoch. Using this train loader the model can reach > 0.78 CV in a few epochs.\n- Unfortunately this training strategy does not fit Kaggle's memory, so I have to train the model locally.\n\n\nReference:\n- https://www.kaggle.com/gilfernandes/riiid-self-attention-transformer\n- https://www.kaggle.com/manikanthr5/riiid-sakt-model-training-public\n- https://www.kaggle.com/leadbest/sakt-with-randomization-state-updates\n- https://www.kaggle.com/wangsg/a-self-attentive-model-for-knowledge-tracing"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import gc\nimport random\nfrom tqdm.notebook import tqdm\nimport numpy as np \nimport pandas as pd \nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import train_test_split\nimport psutil\nimport seaborn as sns\nsns.set()\nimport matplotlib.pyplot as plt\n\nimport pickle\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader\nfrom pathlib import Path\nimport datatable as dt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN = False\nn_skill = 13523\n\nNUM_SKILLS = 13523 # number of problems\nMAX_SEQ = 180\nACCEPTED_USER_CONTENT_SIZE = 5\nNUM_EMBED = 128\nRECENT_SIZE = 10 # recent data in the training loop\nTIMESTAMP_GAP = 20_000\nNUM_HEADS = 8\nBATCH_SIZE = 64\nVAL_BATCH_SIZE = 512\nTEST_SIZE = 0.05\nDROPOUT = 0.1\nSEED = 1127","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ntrain_dtypes = {\n    'content_type_id': 'bool',\n    'timestamp': 'int64',\n    'user_id': 'int32', \n    'content_id': 'int16', \n    'answered_correctly': 'int8', \n    'prior_question_elapsed_time': 'float32', \n    'prior_question_had_explanation': 'bool'\n}\ntarget = 'answered_correctly'\ntrain_df = dt.fread('../input/riiid-test-answer-prediction/train.csv', \n                   columns=set(train_dtypes.keys())).to_pandas().astype(train_dtypes)\n\n# train_df = pd.read_parquet('../input/cv-strategy-in-the-kaggle-environment/cv3_train.parquet')\n# train_df = train_df[train_dtypes.keys()]\n# train_df = train_df.astype(train_dtypes)\n\n\ntrain_df = train_df[train_df.content_type_id == False]\n#arrange by timestamp\ntrain_df = train_df.sort_values(['timestamp'], ascending=True).reset_index(drop = True)\n\n# del train_df['timestamp']\ndel train_df['content_type_id']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Pre-process\nAdding the `timestamp`."},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ngroup = train_df[['user_id', 'content_id', 'answered_correctly','timestamp']]\\\n            .groupby('user_id')\\\n            .apply(lambda r: (r['content_id'].values, \n                              r['answered_correctly'].values,\n                              r['timestamp'].values))\n\ndel train_df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# %%time\n# with open('../input/riiid-sakt-baseline-user-group/group.pickle', 'rb') as f:\n#     group = pickle.load(f)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"group, val_group = train_test_split(group, test_size = TEST_SIZE, random_state=SEED)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Main change: train loaders"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAKTDataset(Dataset):\n    def __init__(self, group, n_skill, \n                        max_seq=MAX_SEQ, \n                        min_seq=ACCEPTED_USER_CONTENT_SIZE,\n                        recent_seq=RECENT_SIZE,\n                        gap=TIMESTAMP_GAP):\n        super(SAKTDataset, self).__init__()\n        self.samples = {}\n        self.n_skill = n_skill\n        self.max_seq = max_seq\n        self.min_seq = min_seq\n        self.recent_seq = recent_seq\n        self.gap = gap\n        \n        self.user_ids = []\n        for i, user_id in enumerate(group.index):\n            if i%10000==0: print(f\"{i} users processed\")\n            \n            content_id, answered_correctly, timestamp = group[user_id]\n\n            if len(content_id) >= self.min_seq:\n\n                if len(content_id) > self.max_seq:\n                    '''\n                    Case 1: longer than max_seq, break into sub-seqs.\n                    '''\n                    total_questions = len(content_id)\n                    num_seq_user = total_questions // self.max_seq\n                    for seq in range(num_seq_user):\n                        index = f\"{user_id}_{seq}\"\n                        self.user_ids.append(index)\n                        start = seq * self.max_seq\n                        end = (seq + 1) * self.max_seq\n                        '''\n                        New contribution #2:\n                        timestamp[end-1] the original last entry's timestamp\n                        the difference with previous entry should be bigger than a threshold\n                        otherwise they may be in the same test_df batch\n                        '''\n                        idx_same_bundle = []\n                        while timestamp[end-1] - timestamp[end-2] < self.gap and end-start >= self.min_seq:\n                            if timestamp[end-1] == timestamp[end-2]:\n                                idx_same_bundle.append(end-1)\n                            end -= 1     \n                        self.samples[index] = (content_id[start:end], \n                                               answered_correctly[start:end],\n                                               timestamp[start:end]\n                                               )\n                        \n                        if idx_same_bundle and end-1-start >= self.min_seq: \n                            # seeing multiple questions at a time, this list is not empty\n                            for j, idx in enumerate(idx_same_bundle):\n                                index = f\"{user_id}_{seq}_{j}\"\n                                self.user_ids.append(index)\n                                self.samples[index] = (np.r_[content_id[start:end-1], content_id[idx]], \n                                                       np.r_[answered_correctly[start:end-1], answered_correctly[idx]],\n                                                       np.r_[timestamp[start:end-1], timestamp[idx]]\n                                                       )\n                    '''\n                    left-over sequence\n                    '''\n                    content_id_last = content_id[end:]\n                    answered_correctly_last = answered_correctly[end:]\n                    timestamp_last = timestamp[end:]\n                    end = len(content_id_last)\n\n                    if end >= self.min_seq and end <= self.max_seq:\n                        idx_same_bundle = []\n                        while timestamp_last[end-1] - timestamp_last[end-2] < self.gap and end >= self.min_seq:\n                            if timestamp_last[end-1] == timestamp_last[end-2]:\n                                idx_same_bundle.append(end-1)\n                            end -= 1\n\n                        index = f\"{user_id}_{num_seq_user + 1}\"\n                        self.user_ids.append(index)\n                        self.samples[index] = (content_id_last[:end], \n                                               answered_correctly_last[:end],\n                                               timestamp_last[:end]\n                                               )\n\n                        if idx_same_bundle and end >= self.min_seq: \n                            # seeing multiple questions at a time, this list is not empty\n                            for j, idx in enumerate(idx_same_bundle):\n                                index = f\"{user_id}_{num_seq_user + 1}_{j}\"\n                                self.user_ids.append(index)\n                                self.samples[index] = (np.r_[content_id_last[:end-1], content_id_last[idx]], \n                                                       np.r_[answered_correctly_last[:end-1], answered_correctly_last[idx]],\n                                                       np.r_[timestamp_last[:end-1], timestamp_last[idx]]\n                                                       )\n                else: # len(content_id) <= self.max_seq\n                    '''\n                    Case 2: shorter than max_seq, keep them all\n                    '''\n                    index = f'{user_id}'\n                    end = len(timestamp)\n                    idx_same_bundle = []\n                    # last time stamp diff should be bigger than a threshold\n                    while timestamp[end-1] - timestamp[end-2] < self.gap and end >= self.min_seq:\n                        if timestamp[end-1] == timestamp[end-2]:\n                                idx_same_bundle.append(end-1)\n                        end -= 1\n                    self.user_ids.append(index)\n                    self.samples[index] = (content_id[:end], \n                                           answered_correctly[:end],\n                                           timestamp[:end],\n                                           )\n\n                    if idx_same_bundle and end >= self.min_seq: \n                        # seeing multiple questions at a time, this list is not empty\n                        for j, idx in enumerate(idx_same_bundle):\n                            index = f\"{user_id}_{j}\"\n                            self.user_ids.append(index)\n                            self.samples[index] = (np.r_[content_id[:end-1], content_id[idx]], \n                                                   np.r_[answered_correctly[:end-1], answered_correctly[idx]],\n                                                   np.r_[timestamp[:end-1], timestamp[idx]]\n                                                   )\n            '''\n            New contribution #1\n            Adding a shifted sequence, now train_loader has much more sequences per epoch\n            '''\n            if self.recent_seq is None: self.recent_seq = self.max_seq + 1\n                \n            if len(content_id) >= 2*self.recent_seq: #\n                for i in range(1, self.recent_seq): # adding a shifted sequence\n                    '''\n                    Shifting cases:\n                    generating much much more sequences by shifting the last few entries\n                    '''\n                    content_id_shift = content_id[:-i]\n                    answered_correctly_shift = answered_correctly[:-i]\n                    timestamp_shift = timestamp[:-i]\n                    if len(content_id_shift) >= self.min_seq:\n                        if len(content_id_shift) > self.max_seq:\n                            '''\n                            Case S 1: shifted seq greater than max_seq, break into pieces\n                            '''\n                            total_questions_2 = len(content_id_shift)\n                            num_seq_user = total_questions_2 // self.max_seq\n\n                            for seq in range(num_seq_user):\n\n                                index = f\"{user_id}_{seq}_{i}_s\"\n                                self.user_ids.append(index)\n                                start = seq * self.max_seq\n                                end = (seq + 1) * self.max_seq\n\n                                idx_same_bundle = []\n                                while timestamp_shift[end-1] - timestamp_shift[end-2] < self.gap and end-start >= self.min_seq:\n                                    if timestamp_shift[end-1] == timestamp_shift[end-2]:\n                                        idx_same_bundle.append(end-1)\n                                    end -= 1\n\n                                self.samples[index] = (content_id_shift[start:end], \n                                                       answered_correctly_shift[start:end],\n                                                       timestamp_shift[start:end]\n                                                       )\n\n                                if idx_same_bundle and end-1-start >= self.min_seq: \n                                    # seeing multiple questions at a time, this list is not empty\n                                    for j, idx in enumerate(idx_same_bundle):\n                                        index = f\"{user_id}_{seq}_{i}_s_{j}\"\n                                        self.user_ids.append(index)\n                                        self.samples[index] = (np.r_[content_id_shift[start:end-1], content_id_shift[idx]], \n                                        np.r_[answered_correctly_shift[start:end-1], answered_correctly_shift[idx]],\n                                        np.r_[timestamp_shift[start:end-1], timestamp_shift[idx]]\n                                        )\n                            '''\n                            left-over sequence\n                            '''\n                            content_id_last = content_id_shift[end:]\n                            answered_correctly_last = answered_correctly_shift[end:]\n                            timestamp_last = timestamp_shift[end:]\n                            end = len(content_id_last)              \n                            if end >= self.min_seq and end <= self.max_seq:\n                                idx_same_bundle = []\n                                while timestamp_last[end-1] - timestamp_last[end-2] < self.gap and end >= self.min_seq:\n                                    if timestamp_last[end-1] == timestamp_last[end-2]:\n                                        idx_same_bundle.append(end-1)\n                                    end -= 1\n\n                                index = f\"{user_id}_{num_seq_user + 1}_{i}_s\"\n                                self.user_ids.append(index)\n                                self.samples[index] = (content_id_last[:end], \n                                                       answered_correctly_last[:end],\n                                                       timestamp_last[:end]\n                                                       )\n\n                                if idx_same_bundle and end >= self.min_seq: \n                                    # seeing multiple questions at a time, this list is not empty\n                                    for j, idx in enumerate(idx_same_bundle):\n                                        index = f\"{user_id}_{num_seq_user + 1}_{i}_s_{j}\"\n                                        self.user_ids.append(index)\n                                        self.samples[index] = (np.r_[content_id_last[:end-1], content_id_last[idx]], \n                                        np.r_[answered_correctly_last[:end-1], answered_correctly_last[idx]],\n                                        np.r_[timestamp_last[:end-1], timestamp_last[idx]]\n                                        )\n                        else: #len(content_id_shift) <= self.max_seq\n                            '''\n                            Case S 2: shifted seq less than or equal to max_seq\n                            '''\n                            index = f'{user_id}_{i}_s'\n                            end = len(timestamp_shift)\n                            idx_same_bundle = []\n                            # last time stamp diff should be bigger than a threshold\n                            while timestamp_shift[end-1] - timestamp_shift[end-2] < self.gap and end >= self.min_seq:\n                                if timestamp_shift[end-1] == timestamp_shift[end-2]:\n                                        idx_same_bundle.append(end-1)\n                                end -= 1\n                            self.user_ids.append(index)\n                            self.samples[index] = (content_id_shift[:end], \n                                                   answered_correctly_shift[:end],\n                                                   timestamp_shift[:end]\n                                                   )\n\n                            if idx_same_bundle and end >= self.min_seq: \n                                # seeing multiple questions at a time, this list is not empty\n                                for j, idx in enumerate(idx_same_bundle):\n                                    index = f\"{user_id}_{i}_s_{j}\"\n                                    self.user_ids.append(index)\n                                    self.samples[index] = (np.r_[content_id_shift[:end-1], content_id_shift[idx]], \n                                                           np.r_[answered_correctly_shift[:end-1], answered_correctly_shift[idx]],\n                                                           np.r_[timestamp_shift[:end-1], timestamp_shift[idx]]\n                                                           )\n\n    def __len__(self):\n        return len(self.user_ids)\n\n    def __getitem__(self, index):\n        user_id = self.user_ids[index]\n        # content_id, answered_correctly = self.samples[user_id]\n        content_id, answered_correctly, timestamp = self.samples[user_id]\n        seq_len = len(content_id)\n        \n        content_id_seq = np.zeros(self.max_seq, dtype=int)\n        answered_correctly_seq = np.zeros(self.max_seq, dtype=int)\n        timestamp_seq = np.zeros(self.max_seq, dtype=int)\n\n        if seq_len >= self.max_seq:\n            content_id_seq[:] = content_id[-self.max_seq:]\n            answered_correctly_seq[:] = answered_correctly[-self.max_seq:]\n            timestamp_seq[:] = timestamp[-self.max_seq:]\n        else:\n            content_id_seq[-seq_len:] = content_id\n            answered_correctly_seq[-seq_len:] = answered_correctly\n            timestamp_seq[-seq_len:] = timestamp\n            \n        target_id = content_id_seq[1:] # question including the current one\n        label = answered_correctly_seq[1:] # answers including the current\n        timestamp = timestamp_seq[1:] # timestamp including the current\n\n        x = content_id_seq[:-1].copy() # question till the previous one\n        # encoded answers till the previous one as the past correctly answering seq\n        x += (answered_correctly_seq[:-1] == 1) * self.n_skill\n        \n        return x, target_id, label, timestamp\n    \n    \nclass SAKTValDataset(Dataset):\n    '''\n    Only for validation\n    '''\n    def __init__(self, group, n_skill, \n                        max_seq=MAX_SEQ, \n                        min_seq=ACCEPTED_USER_CONTENT_SIZE,\n                        recent_seq=None):\n        super(SAKTValDataset, self).__init__()\n        self.samples = {}\n        self.n_skill = n_skill\n        self.max_seq = max_seq\n        self.min_seq = min_seq\n        self.recent_seq = recent_seq\n        \n        self.user_ids = []\n        for i, user_id in enumerate(group.index):\n            try:\n                content_id, answered_correctly = group[user_id]\n            except:\n                content_id, answered_correctly,_ = group[user_id]\n            if len(content_id) >= self.min_seq:\n                if len(content_id) > self.max_seq:\n                    total_questions = len(content_id)\n                    last_pos = total_questions // self.max_seq\n                    for seq in range(last_pos):\n                        index = f\"{user_id}_{seq}\"\n                        self.user_ids.append(index)\n                        start = seq * self.max_seq\n                        end = (seq + 1) * self.max_seq\n                        self.samples[index] = (content_id[start:end], \n                                               answered_correctly[start:end])\n                    if len(content_id[end:]) >= self.min_seq:\n                        index = f\"{user_id}_{last_pos + 1}\"\n                        self.user_ids.append(index)\n                        self.samples[index] = (content_id[end:], \n                                               answered_correctly[end:])\n                else:\n                    index = f'{user_id}'\n                    self.user_ids.append(index)\n                    self.samples[index] = (content_id, answered_correctly)\n            '''\n            New: adding a shifted sequence\n            '''\n            if self.recent_seq is None: self.recent_seq = self.max_seq + 1\n            if len(content_id) >= 2*self.recent_seq: #\n                for i in range(1, self.recent_seq): # adding a shifted sequence\n                    '''\n                    generating much much more sequences by truncating\n                    '''\n                    content_id_truncated_end = content_id[:-i]\n                    answered_correctly_truncated_end = answered_correctly[:-i]\n                    if len(content_id_truncated_end) >= self.min_seq:\n                        if len(content_id_truncated_end) > self.max_seq:\n                            total_questions_2 = len(content_id_truncated_end)\n                            last_pos = total_questions_2 // self.max_seq\n                            for seq in range(last_pos):\n                                index = f\"{user_id}_{seq}_{i}_2\"\n                                self.user_ids.append(index)\n                                start = seq * self.max_seq\n                                end = (seq + 1) * self.max_seq\n                                self.samples[index] = (content_id_truncated_end[start:end], \n                                                    answered_correctly_truncated_end[start:end])\n                            if len(content_id_truncated_end[end:]) >= self.min_seq:\n                                index = f\"{user_id}_{last_pos + 1}_{i}_2\"\n                                self.user_ids.append(index)\n                                self.samples[index] = (content_id_truncated_end[end:], \n                                                    answered_correctly_truncated_end[end:])\n                        else:\n                            index = f'{user_id}_{i}_2'\n                            self.user_ids.append(index)\n                            self.samples[index] = (content_id_truncated_end, \n                                                   answered_correctly_truncated_end)\n\n\n    def __len__(self):\n        return len(self.user_ids)\n\n    def __getitem__(self, index):\n        user_id = self.user_ids[index]\n        content_id, answered_correctly = self.samples[user_id]\n        seq_len = len(content_id)\n        \n        content_id_seq = np.zeros(self.max_seq, dtype=int)\n        answered_correctly_seq = np.zeros(self.max_seq, dtype=int)\n\n        if seq_len >= self.max_seq:\n            content_id_seq[:] = content_id[-self.max_seq:]\n            answered_correctly_seq[:] = answered_correctly[-self.max_seq:]\n        else:\n            content_id_seq[-seq_len:] = content_id\n            answered_correctly_seq[-seq_len:] = answered_correctly\n            \n        target_id = content_id_seq[1:] # question including the current one\n        label = answered_correctly_seq[1:]\n        \n        x = content_id_seq[:-1].copy() # question till the previous one\n        # encoded answers till the previous one\n        x += (answered_correctly_seq[:-1] == 1) * self.n_skill\n        \n        return x, target_id, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain_dataset = SAKTDataset(group, \n                            n_skill=NUM_SKILLS, \n                            max_seq=MAX_SEQ,\n                            min_seq=ACCEPTED_USER_CONTENT_SIZE, \n                            recent_seq=None,\n                            gap=TIMESTAMP_GAP)\ntrain_dataloader = DataLoader(train_dataset, \n                        batch_size=BATCH_SIZE, \n                        num_workers=4,\n                        shuffle=True, \n                        drop_last=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nval_dataset = SAKTValDataset(val_group, n_skill=NUM_SKILLS, max_seq=MAX_SEQ)\nval_dataloader = DataLoader(val_dataset, \n                        batch_size=VAL_BATCH_SIZE, \n                            num_workers=4,\n                        shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(train_dataloader), len(val_dataloader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_batch = next(iter(train_dataloader))\nsample_batch[0].shape, sample_batch[1].shape, sample_batch[2].shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Define model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FFN(nn.Module):\n    def __init__(self, state_size = MAX_SEQ, \n                    forward_expansion = 1, \n                    bn_size=MAX_SEQ - 1, \n                    dropout=DROPOUT):\n        super(FFN, self).__init__()\n        self.state_size = state_size\n        \n        self.lr1 = nn.Linear(state_size, forward_expansion * state_size)\n        self.relu = nn.ReLU()\n        self.bn = nn.BatchNorm1d(bn_size)\n        self.lr2 = nn.Linear(forward_expansion * state_size, state_size)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x):\n        x = self.relu(self.lr1(x))\n        x = self.bn(x)\n        x = self.lr2(x)\n        return self.dropout(x)\n\ndef future_mask(seq_length):\n    future_mask = (np.triu(np.ones([seq_length, seq_length]), k = 1)).astype('bool')\n    return torch.from_numpy(future_mask)\n\n\n\nclass TransformerBlock(nn.Module):\n    def __init__(self, embed_dim, \n                    heads = NUM_HEADS, \n                    dropout = DROPOUT, \n                    forward_expansion = 1):\n        super(TransformerBlock, self).__init__()\n        self.multi_att = nn.MultiheadAttention(embed_dim=embed_dim, \n                        num_heads=heads, dropout=dropout)\n        self.dropout = nn.Dropout(dropout)\n        self.layer_normal = nn.LayerNorm(embed_dim)\n        self.ffn = FFN(embed_dim, \n                    forward_expansion = forward_expansion, \n                    dropout=dropout)\n        self.layer_normal_2 = nn.LayerNorm(embed_dim)\n        \n\n    def forward(self, value, key, query, att_mask):\n        att_output, att_weight = self.multi_att(value, key, query, attn_mask=att_mask)\n        att_output = self.dropout(self.layer_normal(att_output + value))\n        att_output = att_output.permute(1, 0, 2) \n        # att_output: [s_len, bs, embed] => [bs, s_len, embed]\n        x = self.ffn(att_output)\n        x = self.dropout(self.layer_normal_2(x + att_output))\n        return x.squeeze(-1), att_weight\n    \nclass Encoder(nn.Module):\n    def __init__(self, n_skill, max_seq=MAX_SEQ, \n                 embed_dim=NUM_EMBED, \n                 dropout = DROPOUT, \n                 forward_expansion = 1, \n                 num_layers=1, \n                 heads = NUM_HEADS):\n        super(Encoder, self).__init__()\n        self.n_skill, self.embed_dim = n_skill, embed_dim\n        self.embedding = nn.Embedding(2 * n_skill + 1, embed_dim)\n        self.pos_embedding = nn.Embedding(max_seq - 1, embed_dim)\n        self.e_embedding = nn.Embedding(n_skill+1, embed_dim)\n        self.layers = nn.ModuleList([TransformerBlock(embed_dim, heads=heads,\n                forward_expansion = forward_expansion) for _ in range(num_layers)])\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x, question_ids):\n        device = x.device\n        x = self.embedding(x)\n        pos_id = torch.arange(x.size(1)).unsqueeze(0).to(device)\n        pos_x = self.pos_embedding(pos_id)\n        x = self.dropout(x + pos_x)\n        x = x.permute(1, 0, 2) # x: [bs, s_len, embed] => [s_len, bs, embed]\n        e = self.e_embedding(question_ids)\n        e = e.permute(1, 0, 2)\n        for layer in self.layers:\n            att_mask = future_mask(e.size(0)).to(device)\n            x, att_weight = layer(e, x, x, att_mask=att_mask)\n            x = x.permute(1, 0, 2)\n        x = x.permute(1, 0, 2)\n        return x, att_weight\n\nclass SAKTModel(nn.Module):\n    def __init__(self, \n                n_skill, \n                max_seq=MAX_SEQ, \n                embed_dim=NUM_EMBED, \n                dropout = DROPOUT, \n                forward_expansion = 1, \n                enc_layers=1, \n                heads = NUM_HEADS):\n        super(SAKTModel, self).__init__()\n        self.encoder = Encoder(n_skill, \n                               max_seq, \n                               embed_dim, \n                               dropout, \n                               forward_expansion, \n                               num_layers=enc_layers,\n                               heads=heads)\n        self.pred = nn.Linear(embed_dim, 1)\n        \n    def forward(self, x, question_ids):\n        x, att_weight = self.encoder(x, question_ids)\n        x = self.pred(x)\n        return x.squeeze(-1), att_weight","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_num_params(model):\n    model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n    n_params = sum([np.prod(p.size()) for p in model_parameters])\n    return n_params\n\nmodel = SAKTModel(n_skill, \n                  max_seq=MAX_SEQ, \n                  embed_dim=NUM_EMBED, \n                  heads=NUM_HEADS, \n                  dropout=DROPOUT)\n\nn_params = get_num_params(model)\nprint(f\"Current model has {n_params} parameters.\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model(sample_batch[0], sample_batch[1])[0].size()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_PATH = '/kaggle/working/sakt.pth'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_from_item(item):\n    x = item[0].to(device).long()\n    target_id = item[1].to(device).long()\n    label = item[2].to(device).float()\n    target_mask = (target_id != 0)\n    return x, target_id, label, target_mask\n\ndef update_stats(train_loss, loss, \n                 output, label, num_corrects, num_total, \n                 labels, outs):\n    train_loss.append(loss.item())\n    pred = (torch.sigmoid(output) >= 0.5).long()\n    num_corrects += (pred == label).sum().item()\n    num_total += len(label)\n    labels.extend(label.view(-1).data.cpu().numpy())\n    outs.extend(output.view(-1).data.cpu().numpy())\n    return num_corrects, num_total\n\ndef train_epoch(model, dataloader, 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(dataloader)\n    for k, item in enumerate(tbar):\n        x, target_id, label, target_mask = load_from_item(item)\n        \n        optim.zero_grad()\n        output, _ = model(x, target_id)\n        \n        output = torch.masked_select(output, target_mask)\n        label = torch.masked_select(label, target_mask)\n        \n        loss = criterion(output, label)\n        \n        loss.backward()\n        optim.step()\n        scheduler.step()\n        train_loss.append(loss.item())\n        if k % 10 == 0:\n            tbar.set_description('Train loss - {:.4f}'.format(np.mean(train_loss)))\n            tbar.update(10)\n\ndef val_epoch(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        x, target_id, label, target_mask = load_from_item(item)\n\n        with torch.no_grad():\n            output, atten_weight = model(x, target_id)\n        \n        output = torch.masked_select(output, target_mask)\n        label = torch.masked_select(label, target_mask)\n\n        loss = criterion(output, label) \n        \n        num_corrects, num_total = update_stats(val_loss, \n                                               loss, \n                                               output, \n                                               label, \n                                               num_corrects, \n                                               num_total, \n                                               labels, \n                                               outs)\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\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def do_train(lr=1e-3, \n             epochs=10):\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    criterion = nn.BCEWithLogitsLoss()\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, \n                                                    max_lr=lr, \n                                                    steps_per_epoch=len(train_dataloader), \n                                                    epochs=epochs)\n    model.to(device)\n    criterion.to(device)\n    best_auc = 0.0\n    for epoch in range(epochs):\n        train_epoch(model, train_dataloader, optimizer, criterion, scheduler, device)\n        val_loss, val_acc, val_auc = val_epoch(model, val_dataloader, criterion, device)\n        print(f\"epoch - {epoch + 1} val_loss - {val_loss:.3f} acc - {val_acc:.3f} auc - {val_auc:.3f}\")\n        if best_auc < val_auc:\n            print(f'epoch - {epoch + 1} best model with val auc: {val_auc}')\n            best_auc = val_auc\n        torch.save(model.state_dict(), MODEL_PATH)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if TRAIN:\n\n    LR = 1e-3\n    EPOCHS = 3\n    do_train(lr=LR, epochs=EPOCHS)\n\n    LR = 1e-4\n    EPOCHS = 2\n    do_train(lr=LR, epochs=EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = SAKTModel(n_skill, \n                  max_seq=MAX_SEQ, \n                  embed_dim=NUM_EMBED, \n                  heads=NUM_HEADS, \n                  dropout=DROPOUT)\nMODEL_PATH = '../input/riiid-models/sakt_seq_180_auc_0.7836.pt'\nmodel.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))\nmodel.to(device)\nmodel.eval();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, samples, test_df, n_skill, max_seq=100):\n        super(TestDataset, self).__init__()\n        self.samples, self.user_ids, self.test_df = samples, [x for x in test_df[\"user_id\"].unique()], test_df\n        self.n_skill, self.max_seq = n_skill, 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        \n        user_id = test_info['user_id']\n        target_id = test_info['content_id']\n        \n        content_id_seq = np.zeros(self.max_seq, dtype=int)\n        answered_correctly_seq = np.zeros(self.max_seq, dtype=int)\n        \n        if user_id in self.samples.index:\n            content_id, answered_correctly = self.samples[user_id]\n            \n            seq_len = len(content_id)\n            \n            if seq_len >= self.max_seq:\n                content_id_seq = content_id[-self.max_seq:]\n                answered_correctly_seq = answered_correctly[-self.max_seq:]\n            else:\n                content_id_seq[-seq_len:] = content_id\n                answered_correctly_seq[-seq_len:] = answered_correctly\n                \n        x = content_id_seq[1:].copy()\n        x += (answered_correctly_seq[1:] == 1) * self.n_skill\n        \n        questions = np.append(content_id_seq[2:], [target_id])\n        \n        return x, questions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import riiideducation\n\nenv = riiideducation.make_env()\niter_test = env.iter_test()\n\nprev_test_df = None","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"In the example test there are 4 batches of test_df, aside from the placeholder user `275030867`, there are 3 other common users"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_1_user_id = [ 275030867,  554169193, 1720860329,  288641214, 1728340777,\n       1364159702, 1521618396, 1317245193, 1700555100,  998511398,\n       1422853669, 1096784725,  385471210, 1202386221, 2018567473]\nbatch_2_user_id = [ 275030867, 1233875513,  891955351, 1981166446, 1637273633,\n       2030979309,  319060572,  288641214,   98059812,  674533997,\n        555691277, 1317245193, 1202386221,  775113212, 1219481379]\nbatch_3_user_id = [ 275030867, 1521618396, 1148874033,  554169193, 1281335472,\n        998511398, 2002570769,  706626847, 1422853669, 1357500007,\n       2018567473, 1720860329,  674533997, 1202386221,  891955351,\n       1317245193,  385471210,  555691277,  288641214, 1364159702,\n       1599808246,   98059812, 1728340777]\nbatch_4_user_id = [ 275030867, 1305988022, 1310228392, 1637273633,  674533997,\n       2093197291, 1202386221, 1468996389,  555691277, 1838324752,\n       2103436554,  311890082, 1817433235,  998511398, 1422853669,\n        554169193, 1317245193, 1900527744,    7792299,  288641214,\n       2018567473]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"set(batch_1_user_id).intersection(batch_2_user_id).intersection(batch_3_user_id).intersection(batch_4_user_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## First iteration\n\nLet us look at the user `288641214`"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"test_df, sample_prediction_df =next(iter_test) \n    \nif prev_test_df is not None:\n    prev_test_df['answered_correctly'] = eval(test_df['prior_group_answers_correct'].iloc[0])\n    prev_test_df = prev_test_df[prev_test_df.content_type_id == False]\n    prev_group = prev_test_df[['user_id', 'content_id', 'answered_correctly']]\\\n    .groupby('user_id').apply(lambda r: (\n        r['content_id'].values,\n        r['answered_correctly'].values))\n    for prev_user_id in prev_group.index:\n        prev_group_content = prev_group[prev_user_id][0]\n        prev_group_answered_correctly = prev_group[prev_user_id][1]\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_answered_correctly))\n        else:\n            group[prev_user_id] = (prev_group_content, prev_group_answered_correctly)\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_answered_correctly = group[prev_user_id][1][-MAX_SEQ:]\n            group[prev_user_id] = (new_group_content, new_group_answered_correctly)\n\nprev_test_df = test_df.copy()\ntest_df = test_df[test_df.content_type_id == False]\ntest_df['answered_correctly'] = 0.5\nenv.predict(test_df.loc[test_df['content_type_id'] == 0, \n                        ['row_id', 'answered_correctly']])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The last 10 question sequence this user sees is `[12098,   901,  1240,   873, 12075,   592,   276,   429, 11912, 13262]`"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df[test_df.user_id==288641214]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## second iteration\nThe `timestamp` from the first `test_df` is 62798072960. In the second iteration, this user sequence is updated in the last bit. The `timestamp` is 62798100988, and there is a lag of 28028. So if the model wants to thrive, the train loader needs to able to generate this type of sequence. A new question is appended to the end: `[  901,  1240,   873, 12075,   592,   276,   429, 11912, 13262,  5418]`"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df, sample_prediction_df =next(iter_test) \n    \nif prev_test_df is not None:\n    prev_test_df['answered_correctly'] = eval(test_df['prior_group_answers_correct'].iloc[0])\n    prev_test_df = prev_test_df[prev_test_df.content_type_id == False]\n    prev_group = prev_test_df[['user_id', 'content_id', 'answered_correctly']]\\\n    .groupby('user_id').apply(lambda r: (\n        r['content_id'].values,\n        r['answered_correctly'].values))\n    for prev_user_id in prev_group.index:\n        prev_group_content = prev_group[prev_user_id][0]\n        prev_group_answered_correctly = prev_group[prev_user_id][1]\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_answered_correctly))\n        else:\n            group[prev_user_id] = (prev_group_content, prev_group_answered_correctly)\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_answered_correctly = group[prev_user_id][1][-MAX_SEQ:]\n            group[prev_user_id] = (new_group_content, new_group_answered_correctly)\n\nprev_test_df = test_df.copy()\ntest_df = test_df[test_df.content_type_id == False]\ntest_df['answered_correctly'] = 0.5\nenv.predict(test_df.loc[test_df['content_type_id'] == 0, \n                        ['row_id', 'answered_correctly']])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"prev_group[288641214] # prev group is updated to have the correct answer despite the first prediction was wrong.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df[test_df.user_id==288641214]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Third batch\nAgain the `timestamp` is now 34626 greater than the second `test_df` for this user. The question sequence is updated to `[ 1240,   873, 12075,   592,   276,   429, 11912, 13262,  5418,  5620]`, and in the fourth batch the question sequence is updated to `[  873, 12075,   592,   276,   429, 11912, 13262,  5418,  5620,  9077]`. We can generate more of these sequences by shifting the sequence a tiny bit as in the train loader I wrote.\n\nIn prediction, only the last entry is used, the loss function can have a decaying weight to weigh the later prediction more, however, I did not get any CV increase using this method."},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df, sample_prediction_df =next(iter_test) \n    \nif prev_test_df is not None:\n    prev_test_df['answered_correctly'] = eval(test_df['prior_group_answers_correct'].iloc[0])\n    prev_test_df = prev_test_df[prev_test_df.content_type_id == False]\n    prev_group = prev_test_df[['user_id', 'content_id', 'answered_correctly']]\\\n    .groupby('user_id').apply(lambda r: (\n        r['content_id'].values,\n        r['answered_correctly'].values))\n    for prev_user_id in prev_group.index:\n        prev_group_content = prev_group[prev_user_id][0]\n        prev_group_answered_correctly = prev_group[prev_user_id][1]\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_answered_correctly))\n        else:\n            group[prev_user_id] = (prev_group_content, prev_group_answered_correctly)\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_answered_correctly = group[prev_user_id][1][-MAX_SEQ:]\n            group[prev_user_id] = (new_group_content, new_group_answered_correctly)\n\nprev_test_df = test_df.copy()\ntest_df = test_df[test_df.content_type_id == False]\ntest_df['answered_correctly'] = 0.5\nenv.predict(test_df.loc[test_df['content_type_id'] == 0, \n                        ['row_id', 'answered_correctly']])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df[test_df.user_id==288641214]","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}