{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Sakt baseline model\n\nThis is the SAKT model baseline inference template following modifications from @mpware and @leadbest. Every irrelevant imports and lines are deleted to save memory (if CUDA is OOM your submission will fail).\n\nI updated some utils functions to make the selection more automatic, if your model is named following this pattern:\n> f'head_{n_head}_embed_{embed_dim}_seq_{max_seq}_auc_{val_auc}.pt' \n\nthen simple plugging in your model file on Kaggle and you are good to go.\n\nBefore submission it also verifies the auc score on the cv3 file generated by @marisakamozz, which is an improvement over @its7171's great kernel.\n\nReference:\n* https://www.kaggle.com/leadbest/sakt-with-randomization-state-updates\n* https://www.kaggle.com/wangsg/a-self-attentive-model-for-knowledge-tracing\n* https://www.kaggle.com/leadbest/sakt-self-attentive-knowledge-tracing-submitter\n* https://www.kaggle.com/its7171/cv-strategy\n* https://www.kaggle.com/mpware/sakt-fork\n"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import gc\nimport psutil\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport random\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.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Insert your model file below, it can be a kernel output or a dataset\nchange `VALID` to false if you want to skip the validation phase"},{"metadata":{"trusted":true},"cell_type":"code","source":"# some global settings\n\nrandom.seed(42)\nMAX_SEQ = 150\nWORKERS = 4\nVAL_BATCH_SIZE = 2048\nn_skill = 13523\nVALID = True\n\n# model_file = '../input/riiid-models/sakt_head_8_embed_128_seq_150_auc_0.7584.pt'\nmodel_file = '../input/riiid-models/sakt_layer_1_head_8_embed_128_seq_150_auc_0.7605.pt'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load data and preprocess"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nTRAIN_DTYPES = {'timestamp':'int64', \n         'user_id':'int32' ,\n         'content_id':'int16',\n         'content_type_id':'int8',\n         'answered_correctly':'int8'}\nTRAIN_COLS = TRAIN_DTYPES.keys()\n\ntrain_df = pd.read_csv('/kaggle/input/riiid-test-answer-prediction/train.csv', \n                       usecols=[1, 2, 3, 4, 7], dtype=TRAIN_DTYPES)\n# train_df = pd.read_parquet('../input/cv-strategy-in-the-kaggle-environment/cv3_train.parquet')\ntrain_df = train_df[TRAIN_COLS].astype(TRAIN_DTYPES)\n\ntrain_df = train_df[train_df[\"content_type_id\"] == False]\n\ntrain_df = train_df.sort_values(['timestamp'], ascending=True).reset_index(drop = True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"group = train_df[['user_id', 'content_id', 'answered_correctly']].groupby('user_id').apply(lambda r: (\n            r['content_id'].values,\n            r['answered_correctly'].values))\n\ndel train_df\ngc.collect();","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load model\n\nChange: `num_heads` is now in the initialization."},{"metadata":{"trusted":true},"cell_type":"code","source":"class FFN(nn.Module):\n    def __init__(self, state_size=200):\n        super(FFN, self).__init__()\n        self.state_size = state_size\n\n        self.lr1 = nn.Linear(state_size, state_size)\n        self.relu = nn.ReLU()\n        self.lr2 = nn.Linear(state_size, state_size)\n        self.dropout = nn.Dropout(0.2)\n    \n    def forward(self, x):\n        x = self.lr1(x)\n        x = self.relu(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\nclass SAKTModel(nn.Module):\n    def __init__(self, n_skill, max_seq=MAX_SEQ, embed_dim=128, num_heads=8): \n        super(SAKTModel, self).__init__()\n        self.n_skill = n_skill\n        self.embed_dim = embed_dim\n\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\n        self.multi_att = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=0.2)\n\n        self.dropout = nn.Dropout(0.2)\n        self.layer_normal = nn.LayerNorm(embed_dim) \n\n        self.ffn = FFN(embed_dim)\n        self.pred = nn.Linear(embed_dim, 1)\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\n        pos_x = self.pos_embedding(pos_id)\n        x = x + pos_x\n\n        e = self.e_embedding(question_ids)\n\n        x = x.permute(1, 0, 2) # x: [bs, s_len, embed] => [s_len, bs, embed]\n        e = e.permute(1, 0, 2)\n        att_mask = future_mask(x.size(0)).to(device)\n        att_output, att_weight = self.multi_att(e, x, x, attn_mask=att_mask)\n        att_output = self.layer_normal(att_output + e)\n        att_output = att_output.permute(1, 0, 2) # att_output: [s_len, bs, embed] => [bs, s_len, embed]\n\n        x = self.ffn(att_output)\n        x = self.layer_normal(x+att_output)\n        x = self.pred(x)\n\n        return x.squeeze(-1), att_weight","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_sakt_model(model_file, device='cuda'):\n    # creating the model and load the weights\n    configs = []\n    model_file_lst = model_file.split('_')\n    for c in ['head', 'embed', 'seq']:\n        idx = model_file_lst.index(c) + 1\n        configs.append(int(model_file_lst[idx]))\n\n    # configs.append(int(model_file[model_file.rfind('head')+5]))\n    # configs.append(int(model_file[model_file.rfind('embed')+6:model_file.rfind('embed')+9]))\n    # configs.append(int(model_file[model_file.rfind('seq')+4:model_file.rfind('seq')+7]))\n    conf_dict = dict(n_skill=n_skill,\n                     num_heads=configs[0],\n                     embed_dim=configs[1], \n                     max_seq=configs[2], \n                     )\n\n    model = SAKTModel(**conf_dict)\n        \n    model = model.to(device)\n    model.load_state_dict(torch.load(model_file, map_location=device))\n\n    return model\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = load_sakt_model(model_file, device=device)\n\nmodel.to(device)\nprint(model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Validation"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAKTDataset(Dataset):\n    def __init__(self, group, n_skill, subset=\"train\", max_seq=MAX_SEQ):\n        super(SAKTDataset, self).__init__()\n        self.max_seq = max_seq\n        self.n_skill = n_skill # 13523\n        self.samples = 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            q, qa = group[user_id] \n            if len(q) < 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        q_, qa_ = self.samples[user_id] # Pick full sequence for user\n        seq_len = len(q_)\n\n        q = np.zeros(self.max_seq, dtype=int)\n        qa = np.zeros(self.max_seq, dtype=int)\n\n        if seq_len >= self.max_seq:\n            if self.subset == \"train\":\n#                 if seq_len > self.max_seq:\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                    q[:] = q_[random_start_index:end_index] \n                    qa[:] = qa_[random_start_index:end_index] \n                else:\n                    q[:] = q_[-self.max_seq:]\n                    qa[:] = qa_[-self.max_seq:]\n            else:\n                q[:] = q_[-self.max_seq:] # Pick last 100 questions\n                qa[:] = qa_[-self.max_seq:] # Pick last 100 answers\n        else:\n            if random.random()>0.1:\n                seq_len = random.randint(2,seq_len)\n                q[-seq_len:] = q_[:seq_len]\n                qa[-seq_len:] = qa_[:seq_len]\n            else:\n                q[-seq_len:] = q_ # Pick last N question with zero padding\n                qa[-seq_len:] = qa_ # Pick last N answers with zero padding\n                \n        target_id = q[1:] # Ignore first item 1 to 99\n        label = qa[1:] # Ignore first item 1 to 99\n\n        # x = np.zeros(self.max_seq-1, dtype=int)\n        x = q[:-1].copy() # 0 to 98\n        x += (qa[:-1] == 1) * self.n_skill # y = et + rt x E\n\n        return x, target_id,  label\n    \n    \ndef valid_epoch(model, valid_iterator, criterion, device=\"cuda\"):\n    model.eval()\n\n    valid_loss = []\n    num_corrects = 0\n    num_total = 0\n    labels = []\n    outs = []\n    len_dataset = len(valid_iterator)\n    \n    with tqdm(total=len_dataset) as pbar:\n        for idx, item in enumerate(valid_iterator): \n            x = item[0].to(device).long()\n            target_id = item[1].to(device).long()\n            label = item[2].to(device).float()\n\n            with torch.no_grad():\n                output, _ = model(x, target_id)\n            loss = criterion(output, label)\n            valid_loss.append(loss.item())\n\n            output = output[:, -1] # (BS, 1)\n            output = torch.sigmoid(output)\n            label = label[:, -1] \n            pred = (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            pbar.set_description(f'val loss batch - {valid_loss[-1]:.4f}')\n            pbar.update(1)\n\n    acc = num_corrects / num_total\n    auc = roc_auc_score(labels, outs)\n    loss = np.mean(valid_loss)\n\n    return loss, acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if VALID:\n    valid_df = pd.read_parquet('../input/cv-strategy-in-the-kaggle-environment/cv3_valid.parquet')\n    valid_df = valid_df[TRAIN_COLS]\n\n    valid_df = valid_df[valid_df[\"content_type_id\"] == False]\n\n    valid_group = valid_df[['user_id', 'content_id', 'answered_correctly']].groupby('user_id').apply(lambda r: (\n            r['content_id'].values,\n            r['answered_correctly'].values))\n\n    del valid_df\n    gc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if VALID:\n    valid_dataset = SAKTDataset(valid_group, n_skill, subset=\"valid\")\n    val_loader = DataLoader(valid_dataset, \n                            batch_size=VAL_BATCH_SIZE, \n                            shuffle=False)\n    criterion = nn.BCEWithLogitsLoss()\n\n    criterion.to(device)\n    val_loss, val_acc, val_auc = valid_epoch(model, val_loader, criterion, device=device)\n    print(f\"Valid: loss - {val_loss:.2f} acc - {val_acc:.4f} auc - {val_auc:.4f}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"## Test\n\n`skill` is not passed to this dataset anymore as only the number of embeddings are needed."},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, samples, test_df, n_skill, max_seq=MAX_SEQ): \n        super(TestDataset, self).__init__()\n        self.samples = samples\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\n        user_id = test_info[\"user_id\"]\n        target_id = test_info[\"content_id\"]\n\n        q = np.zeros(self.max_seq, dtype=int)\n        qa = np.zeros(self.max_seq, dtype=int)\n\n        if user_id in self.samples.index:\n            q_, qa_ = self.samples[user_id]\n            \n            seq_len = len(q_)\n\n            if seq_len >= self.max_seq:\n                q = q_[-self.max_seq:]\n                qa = qa_[-self.max_seq:]\n            else:\n                q[-seq_len:] = q_\n                qa[-seq_len:] = qa_          \n        \n        x = np.zeros(self.max_seq-1, dtype=int)\n        x = q[1:].copy()\n        x += (qa[1:] == 1) * self.n_skill\n        \n        questions = np.append(q[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()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nprev_test_df = None\n\nfor (test_df, sample_prediction_df) in iter_test:\n    if (prev_test_df is not None) & (psutil.virtual_memory().percent<95):\n        print(psutil.virtual_memory().percent)\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']].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_ac = 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_ac))\n \n            else:\n                group[prev_user_id] = (prev_group_content,prev_group_ac)\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                group[prev_user_id] = (new_group_content,new_group_ac)\n\n    prev_test_df = test_df.copy()\n    \n    test_df = test_df[test_df.content_type_id == False]\n                \n    test_dataset = TestDataset(group, test_df, n_skill)\n    test_dataloader = DataLoader(test_dataset, batch_size=51200, shuffle=False)\n    \n    outs = []\n\n    for item in test_dataloader:\n        x = item[0].to(device).long()\n        target_id = item[1].to(device).long()\n\n        with torch.no_grad():\n            output, _ = model(x, target_id)\n            \n        output = torch.sigmoid(output)\n        output = output[:, -1]\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":{},"cell_type":"markdown","source":"## Sanity check\nIf the average of probability is not skewed toward 1, then probably your model is problematic."},{"metadata":{"trusted":true},"cell_type":"code","source":"sns.set()\nsub = pd.read_csv('../working/submission.csv')\nsub['answered_correctly'].hist();","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}