{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport random\nimport numpy as np \nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import roc_auc_score, accuracy_score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This notebook is an attempt to include Saint Embeddings into a single transformer block.\n\nthanks to \n1.  https://www.kaggle.com/wangsg/a-self-attentive-model-for-knowledge-tracing/execution\n2. https://www.kaggle.com/leadbest/sakt-with-randomization-state-updates"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain=pd.read_pickle('../input/riiid-trainpkl/riiid_train.pkl.gzip')\ntrain=train[train.content_type_id==False][['user_id', 'content_id', 'answered_correctly']].copy()\nquestions_df=pd.read_csv('../input/riiid-test-answer-prediction/questions.csv')\n\n\nquestions_df.rename(columns={'question_id': 'content_id'}, inplace=True)\ntrain=train.merge(questions_df[['content_id', 'part']])\n\ndel questions_df\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"questions_df=pd.read_csv('../input/riiid-test-answer-prediction/questions.csv')\nlen(questions_df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_questions=train.content_id.nunique()\nn_parts=train.part.nunique()\nn_responses=3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"group=train.groupby('user_id').apply(lambda row: (row.content_id.values, \n                                                  row.part.values,\n                                                  row.answered_correctly.values))\n\ngroup=group[group.apply(lambda x:len(x[0]) > 10) ]\ndel train\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAKTDataset(torch.utils.data.Dataset):\n    def __init__(self,\n                 group,\n                 max_seq=100):\n        \n        super(SAKTDataset, self).__init__()\n        self.max_seq=max_seq\n        self.group=group\n        self.user_ids=[]\n        for user_id in self.group.index:\n            self.user_ids.append(user_id)\n        \n    def __len__(self):\n        return len(self.group)\n    \n    def __getitem__(self, idx):\n        user_id=self.user_ids[idx]\n        (q_, p_, r_ )=self.group[user_id]\n        \n        seq_len=len(q_)\n        q_=torch.as_tensor(q_, dtype=int)\n        p_=torch.as_tensor(p_, dtype=int)\n        r_=torch.as_tensor(r_, dtype=int)\n        \n        q=torch.zeros(self.max_seq, dtype=int)\n        p=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        \n        \n        label_mask=torch.zeros(self.max_seq, dtype=bool)\n        label_mask[:seq_len]=True\n        \n        \n        if seq_len <= self.max_seq:\n            q[:seq_len]=q_\n            p[:seq_len]=p_\n            y[:seq_len]=r_\n            r[0]=2\n            r[1:seq_len]=r_[:seq_len-1]\n        else:\n            \n            if random.rand() >= 0.1:\n                start=random.randint(0, (seq_len-self.max_seq))\n                end=start+self.max_seq\n                \n                q[:]=q_[start:end]\n                p[:]=p_[start:end]\n                y[:]=r_[start:end]\n                \n                if start == 0:\n                    r[0]=2\n                    r[1:]=r_[:self.max_seq-1]\n                else:\n                    r[:]=r_[start-1: end-1]\n                \n            else:\n                q[:]=q_[:self.max_seq]\n                p[:]=p_[:self.max_seq]\n                y[:]=r_[:self.max_seq]\n                r[0]=2\n                r[1:]=r_[: self.max_seq-1]\n            \n        \n        return (q, p, r, y, label_mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train, val=train_test_split(group, test_size=0.2)\nprint(train.shape, val.shape)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"train_dataset=SAKTDataset(train)\ntrain_dataloader=torch.utils.data.DataLoader(train_dataset, \n                                             batch_size=1024,\n                                             shuffle=True, \n                                             pin_memory=True,\n                                             num_workers=8)\n\n\nval_dataset=SAKTDataset(val)\nval_dataloader=torch.utils.data.DataLoader(val_dataset, \n                                           batch_size=1024, \n                                           shuffle=False,\n                                           pin_memory=True,\n                                           num_workers=8\n                                          )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device= 'cuda' if torch.cuda.is_available() else 'cpu'\nprint('Device:', device)","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, emb_dim):\n        super(FFN, self).__init__()\n        \n        self.fc1=nn.Linear(emb_dim, emb_dim)\n        self.relu1=nn.PReLU()\n        self.dropout1=nn.Dropout(0.2)\n        self.fc2=nn.Linear(emb_dim, emb_dim)\n        self.dropout2=nn.Dropout(0.2)\n        \n    def forward(self, x):\n        x=self.fc1(x)\n        x=self.relu1(x)\n        x=self.dropout1(x)\n        \n        x=self.fc2(x)\n        x=self.dropout2(x)\n        return x\n\n\nclass SAKT(nn.Module):\n    def __init__(self, \n                 n_questions,\n                 n_parts,\n                 n_responses,\n                 device='cpu',\n                 emb_dim=128,\n                 model_dim=128,\n                 num_heads=8,\n                 max_seq=100):\n        \n        super(SAKT, self).__init__()\n        \n        self.pos_idx=torch.arange(max_seq).to(device)\n        self.n_questions=n_questions\n        self.n_parts=n_parts\n        self.n_responses=n_responses\n        self.max_seq=max_seq\n        self.device=device\n        \n        self.emb_dim=emb_dim\n        self.model_dim=model_dim\n        \n        self.pos_embedding=nn.Embedding(max_seq, emb_dim)\n        self.q_embedding=nn.Embedding(n_questions, emb_dim)\n        self.p_embedding=nn.Embedding(n_parts+1, emb_dim)\n        self.r_embedding=nn.Embedding(n_responses, emb_dim)\n        \n        \n        self.multihead_attn=nn.MultiheadAttention(model_dim, num_heads=num_heads, dropout=0.2)\n        self.layernorm1=nn.LayerNorm(model_dim)\n        \n        self.dropout1=nn.Dropout(0.2)\n\n        self.ffn=FFN(model_dim)\n        self.layernorm2=nn.LayerNorm(model_dim)\n        \n        self.dropout2=nn.Dropout(0.2)\n        self.out = nn.Linear(model_dim, 1)\n    \n    def get_attention_mask(self, s):\n        attn_mask=torch.tensor(np.triu(np.ones((s, s)), k=1).astype('bool'))\n        attn_mask=attn_mask.to(self.device)\n        return attn_mask\n    \n    def forward(self, q, p, r):\n        pos_embedd=self.pos_embedding(self.pos_idx)\n        q_embedd=self.q_embedding(q)\n        p_embedd=self.p_embedding(p)\n        r_embedd=self.r_embedding(r)\n        \n        query=q_embedd+p_embedd\n        x=pos_embedd+q_embedd+p_embedd+r_embedd\n        attn_mask=self.get_attention_mask(q.size(1))\n        \n        query=query.permute(1, 0, 2)\n        x=x.permute(1, 0, 2)\n        \n        attn_output, attn_weights=self.multihead_attn(query, x, x, attn_mask=attn_mask)\n        attn_output=self.layernorm1(query+attn_output)\n        \n        ffn_out=self.ffn(attn_output)\n        y=self.layernorm2(attn_output+ffn_out)\n        \n        y=y.permute(1, 0, 2)\n        yout=self.out(y).squeeze(-1)\n        return yout","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train the model"},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs=35\nearly_stop=4\n\n\nmodel=SAKT(n_questions,n_parts,n_responses,device=device).to(device)\noptim=torch.optim.Adam(model.parameters())\ncriterion=torch.nn.BCEWithLogitsLoss().to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(model, optim, criterion):\n    train_loss=[]\n    model.train()\n    for (q, p, r, y, label_mask) in train_dataloader:\n        q=q.to(device)\n        p=p.to(device)\n        r=r.to(device)\n        y=y.to(device)\n        label_mask=label_mask.to(device)\n        \n        optim.zero_grad()\n        yout=model(q, p, r)\n        \n        y=torch.masked_select(y, label_mask).type(torch.cuda.FloatTensor)\n        yout=torch.masked_select(yout, label_mask).to(device)\n        \n        loss_=criterion(yout, y)\n        loss_.backward()\n        optim.step()\n        train_loss.append(loss_.item())\n    return np.mean(train_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def val_epoch(model, criterion):\n    val_loss=[]\n    model.eval()\n    for (q, p, r, y, label_mask) in val_dataloader:\n        q=q.to(device)\n        p=p.to(device)\n        r=r.to(device)\n        y=y.to(device)\n        \n        with torch.no_grad():\n            yout=model(q, p, r)\n            y=torch.masked_select(y, label_mask).type(torch.cuda.FloatTensor)\n            yout=torch.masked_select(yout, label_mask).to(device)\n            loss_=criterion(yout, y)\n            val_loss.append(loss_.item())\n    return np.mean(val_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain_loss=[]\nval_loss=[]\n\nbest_loss=None\n\nfor i in range(num_epochs):\n    start_time=time.time()\n    \n    if early_stop==0:\n        print('Early Stopping on the epoch:{}'.format(i))\n        break\n        \n    train_loss_=train_epoch(model, optim, criterion)\n    val_loss_=val_epoch(model, criterion)\n    \n    train_loss.append(train_loss_)\n    val_loss.append(val_loss)\n    \n    if best_loss==None or best_loss>val_loss_:\n        early_stop=4\n        best_loss=val_loss_\n        torch.save(model.state_dict(), 'sakt.pth')\n    elif best_loss!=None and best_loss<=val_loss_:\n        early_stop-=1\n        \n    end_time=time.time()\n    gc.collect()\n    \n    print('Epoch Time:', (end_time-start_time))\n    print(\"Epoch:{} | Train Loss:{:.4f}\".format(i, train_loss_))\n    print(\"Epoch:{} | Val Loss:{:4f}\".format(i, val_loss_))\n    print('-------')\n    print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.title('Train Loss')\nplt.plot(train_loss)\nplt.show()\n\nplt.title('Val Loss')\nplt.plot(val_loss)\nplt.show()","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}