{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import gc\nimport time\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import roc_auc_score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Thanks for the SAKT Implementation workbook:\nhttps://www.kaggle.com/wangsg/a-self-attentive-model-for-knowledge-tracing/execution"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df=pd.read_pickle('../input/riiid-trainpkl/riiid_train.pkl.gzip')\nquestions_df=pd.read_csv('../input/riiid-test-answer-prediction/questions.csv')\nquestions_df.rename(columns={'question_id': 'content_id'}, inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df=train_df[train_df.content_type_id==False]\ntrain_df=train_df[['row_id','user_id', 'content_id', 'answered_correctly']].merge(questions_df[['content_id', 'part']], how='left')\ntrain_df=train_df.groupby('user_id').head(100)\n\nprint(train_df.shape)\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df=train_df.groupby('user_id').apply(lambda row: (\n        row.content_id.values, \n        row.part.values, \n        row.answered_correctly.values))\n\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Filter the users with <=5 questions answered"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df=train_df[train_df.apply(lambda x:x[0].shape[0])>5]\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"n_questions=len(questions_df)\nn_categories=1+questions_df.part.nunique()\nn_responses=3\n\nbatch_size=2048\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTDataset(torch.utils.data.Dataset):\n    def __init__(self, df, max_seq=100):\n        self.user_ids=[]\n        self.df=df\n        self.max_seq=max_seq\n        for user_id in df.index.values:\n            self.user_ids.append(user_id)\n    \n    def __len__(self):\n        return len(self.user_ids)\n    \n    def __getitem__(self, idx):\n        user_id=self.user_ids[idx]\n        (q_, c_, r_)=self.df[user_id]\n        seq_len=len(q_)\n        \n        q_=torch.as_tensor(q_, dtype=int)\n        c_=torch.as_tensor(c_, dtype=int)\n        r_=torch.as_tensor(r_, dtype=int)\n        \n        \n        q=torch.zeros(self.max_seq, dtype=int)\n        c=torch.zeros(self.max_seq, dtype=int)\n        r=torch.zeros(self.max_seq, dtype=int)\n        y=torch.zeros(self.max_seq, dtype=int)\n        label_mask=torch.zeros(self.max_seq, dtype=bool)\n        label_mask[: seq_len]=1\n        \n        \n        r[0]=2 #2-for the start of the sequence\n        if seq_len > self.max_seq:\n            q[:]=q_[:self.max_seq]\n            c[:]=c_[:self.max_seq]\n            r[1:]=r_[:self.max_seq-1]\n            y[:]=r_[:self.max_seq]\n        elif seq_len <= self.max_seq:\n            q[:seq_len]=q_\n            c[:seq_len]=c_\n            r[1:seq_len]=r_[:seq_len-1]\n            y[:seq_len]=r_\n        \n        \n        return (q, c, r, y, label_mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train, val=train_test_split(train_df, test_size=0.2)\n\ntrain_dataset=SAINTDataset(train)\ntrain_dataloader=torch.utils.data.DataLoader(train_dataset, \n                                             batch_size=batch_size,\n                                             shuffle=True,\n                                             num_workers=4,\n                                             pin_memory=True\n                                            )\nval_dataset=SAINTDataset(val)\nval_dataloader=torch.utils.data.DataLoader(val_dataset, \n                                           batch_size=batch_size,\n                                           num_workers=4,\n                                           shuffle=False,\n                                           pin_memory=True\n                                          )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FFN(nn.Module):\n    def __init__(self, units):\n        super().__init__()\n        self.fc1=nn.Linear(units, units)\n        self.relu1=nn.ReLU()\n        self.dropout1=nn.Dropout(0.2)\n        self.fc2=nn.Linear(units, units)\n        \n    def forward(self, x):\n        x=self.fc1(x)\n        x=self.relu1(x)\n        x=self.dropout1(x)\n        x=self.fc2(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTDecoder(nn.Module):\n    def __init__(self, \n                 n_responses,\n                 device='cpu',\n                 max_seq=100,\n                 num_heads=4,\n                 embedd_dim=128,\n                 model_dim=128\n                ):\n        super().__init__()\n        self.device=device\n        self.r_embedding=nn.Embedding(n_responses, embedd_dim) #Adding the start token\n        self.multihead_attn1=nn.MultiheadAttention(model_dim, num_heads, dropout=0.2)\n        self.multihead_attn2=nn.MultiheadAttention(model_dim, num_heads, dropout=0.2)\n        \n        self.ffn1=FFN(model_dim)\n        self.ffn2=FFN(model_dim)\n        \n        self.layer_norm11=nn.LayerNorm(model_dim)\n        self.layer_norm12=nn.LayerNorm(model_dim)\n        self.layer_norm21=nn.LayerNorm(model_dim)\n        self.layer_norm22=nn.LayerNorm(model_dim)\n        \n        self.droupout=nn.Dropout(0.2)\n    \n    def model_stack_output(self, q, k, v, multihead_attn, ffn, layer_norm1, layer_norm2):\n        seq_len=q.shape[0]\n        attn_mask=torch.tensor(np.triu(np.ones((seq_len, seq_len)), k=1)).to(self.device)\n        \n        attn_out, attn_out_weights=multihead_attn(q, k, v, attn_mask=attn_mask)\n        q=layer_norm1(attn_out+q)\n        \n        ffn_out=ffn(q)\n        q=layer_norm2(ffn_out+q)\n        return q\n    \n    def forward(self, pos_r, y_encoder, r):\n        r_embedd=self.r_embedding(r)\n        x=pos_r+r_embedd\n        x=x.permute(1, 0, 2)\n        y_encoder=y_encoder.permute(1, 0, 2)\n        \n        stack_out1=self.model_stack_output(x, x, x, self.multihead_attn1, self.ffn1, self.layer_norm11, self.layer_norm12)\n        stack_out2=self.model_stack_output(stack_out1, y_encoder, y_encoder, self.multihead_attn2, self.ffn2, self.layer_norm21, self.layer_norm22)\n        \n        stack_out2=stack_out2.permute(1, 0, 2)\n        out=self.droupout(stack_out2)\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTEncoder(nn.Module):\n    def __init__(self, \n                 n_questions, n_categories,\n                 device='cpu',\n                 max_seq=100,\n                 num_heads=4,\n                 embedd_dim=128,\n                 model_dim=128\n                ):\n        \n        super().__init__()\n        self.device=device\n        self.q_embedding=nn.Embedding(n_questions, embedd_dim)\n        self.c_embedding=nn.Embedding(n_categories, embedd_dim)\n        \n        self.multihead_attn1=nn.MultiheadAttention(model_dim, num_heads, dropout=0.2)\n        self.ffn1=FFN(model_dim)\n        \n        self.multihead_attn2=nn.MultiheadAttention(model_dim, num_heads, dropout=0.2)\n        self.ffn2=FFN(model_dim)\n        \n        self.layer_norm11=nn.LayerNorm(model_dim) # layer norm for the attention_layer1\n        self.layer_norm12=nn.LayerNorm(model_dim)\n        \n        self.layer_norm21=nn.LayerNorm(model_dim) # layer norm for attention_layer2\n        self.layer_norm22=nn.LayerNorm(model_dim)\n    \n    def model_stack_output(self, x, multihead_attn, ffn, layer_norm1, layer_norm2):\n        seq_len=x.shape[0]\n        attn_mask=torch.tensor(np.triu(np.ones((seq_len, seq_len)), k=1)).to(self.device)\n        attn_out, attn_out_weights=multihead_attn(x,x,x, attn_mask=attn_mask)\n        x=layer_norm1(attn_out+x)\n        \n        ffn_out=ffn(x)\n        x=layer_norm2(ffn_out+x)\n        return x\n        \n        \n    def forward(self, pos_q, q, c):\n        q_embedd=self.q_embedding(q)\n        c_embedd=self.c_embedding(c)\n        x=q_embedd+c_embedd+pos_q\n        \n        \n        x=x.permute(1, 0, 2) # [seq_len, batch_size, emb_dim]\n        # Stack-1\n        stack_out1=self.model_stack_output(x, \n                                           self.multihead_attn1, \n                                           self.ffn1,\n                                           self.layer_norm11, self.layer_norm12)\n        #Stack-2\n        stack_out2=self.model_stack_output(stack_out1, \n                                           self.multihead_attn2, \n                                           self.ffn2,\n                                           self.layer_norm21, self.layer_norm22)\n        stack_out2=stack_out2.permute(1, 0, 2)\n        return stack_out2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTModel(nn.Module):\n    def __init__(self, \n                 n_questions, n_categories, n_responses,\n                 device='cpu',\n                 max_seq=100, \n                 embedd_dim=128, \n                 encoder_dim=128,  \n                 decoder_dim=128,\n                 num_heads=4):\n        \n        super().__init__()\n        self.max_seq=max_seq\n        self.device=device\n        self.pos_embedding=nn.Embedding(max_seq, embedd_dim)\n        self.encoder=SAINTEncoder(n_questions,\n                                  n_categories,\n                                  device=self.device,\n                                  max_seq=max_seq, \n                                  num_heads=num_heads,\n                                  model_dim=encoder_dim)\n        \n        \n        \n        self.decoder=SAINTDecoder(n_responses, \n                                  device=self.device,\n                                  max_seq=max_seq,\n                                  num_heads=num_heads,\n                                  embedd_dim=embedd_dim,\n                                  model_dim=decoder_dim\n                                 )\n        self.out=nn.Linear(decoder_dim, 1)\n        \n    def forward(self, q, c, r):\n        pos_ids=torch.arange(self.max_seq).to(self.device)\n        pos_x=self.pos_embedding(pos_ids)\n        y_encoder=self.encoder(pos_x, q, c)\n        y_decoder=self.decoder(pos_x, y_encoder, r)\n        y_out=self.out(y_decoder)\n        return y_out.squeeze(-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Train Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"device='cuda' if torch.cuda.is_available() else 'cpu'\nmodel=SAINTModel(n_questions, n_categories, n_responses, device=device)\n\noptimizer=torch.optim.Adam(model.parameters())\ncriterion=nn.BCEWithLogitsLoss()\n\n\nmodel.to(device)\ncriterion.to(device)\n\nnum_epochs=3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch():\n    train_loss=[]\n    model.train()\n    \n    for (q, c, r, y, label_mask) in train_dataloader:\n        q=q.to(device)\n        c=c.to(device)\n        r=r.to(device)\n        y=y.to(device)\n        label_mask=label_mask.to(device)\n        \n        optimizer.zero_grad()\n        yout=model(q, c, r)\n        y=torch.masked_select(y, label_mask)\n        yout=torch.masked_select(yout, label_mask)\n        \n        yout=yout.float().to(device)\n        y=y.float().to(device)\n        \n        loss_=criterion(yout, y)\n        loss_.backward()\n        optimizer.step()\n        train_loss.append(loss_.item())\n        \n    return np.mean(train_loss)\n\ndef val_epoch():\n    val_loss=[]\n    model.eval()\n    \n    with torch.no_grad():\n        for (q, c, r, y, label_mask) in train_dataloader:\n            q=q.to(device)\n            c=c.to(device)\n            r=r.to(device)\n            yout=model(q, c, r).float()\n\n            y=torch.masked_select(y, label_mask)\n            yout=torch.masked_select(yout, label_mask)\n\n            yout=yout.float().to(device)\n            y=y.float().to(device)\n            \n            loss_=criterion(yout, y)    \n            val_loss.append(loss_.item())\n            \n    return np.mean(val_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"best_score=None\nfor i in range(num_epochs):\n    epoch_start=time.time()\n    \n    train_loss=train_epoch()\n    val_loss=val_epoch()\n    \n    epoch_end=time.time()\n    print('Time To Run Epoch:{}'.format( (epoch_end - epoch_start)/60) )\n    print(\"Epoch:{} | Train Loss: {:.4f} | Val Loss:{:.4f}\".format(i, train_loss, val_loss))\n    \n    if (best_score is None) or  (best_score>val_loss) :\n        best_score=val_loss\n        torch.save(model.state_dict(), \"saint{}.pth\".format(i))\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}