{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport numpy as np\nimport pandas as pd\nimport riiideducation\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":"# 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.ReLU()\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":{"trusted":true},"cell_type":"code","source":"%%time\ndf=pd.read_pickle('../input/riiid-trainpkl/riiid_train.pkl.gzip')\ndf=df[df.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)\ndf=df.merge(questions_df[['content_id', 'part']])\n\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_questions=df.content_id.nunique()\nn_parts=df.part.nunique()\nn_responses=3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device='cuda' if torch.cuda.is_available() else 'cpu'\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"group=df.groupby('user_id').apply(lambda row: (row.content_id.values[-100:],\n                                               row.part.values[-100:],\n                                               row.answered_correctly.values[-100:]))\n\ngroup.head()\ndel df\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model=SAKT(n_questions,n_parts,n_responses,device=device).to(device)\nmodel.load_state_dict(torch.load('../input/sakt-saint-randomstate-v1/sakt.pth'))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"env = riiideducation.make_env()\niter_test = env.iter_test()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(torch.utils.data.Dataset):\n    def __init__(self, test_df, max_seq=100):\n        self.test_df=test_df\n        self.max_seq=max_seq\n    def __len__(self):\n        return len(self.test_df)\n    def __getitem__(self, idx):\n        row=self.test_df.iloc[idx]\n        content_id=row.content_id\n        part=row.part\n        prev_seq=row.prev_seq\n        \n        (q_, p_, r_)=prev_seq\n        seq_len=q_.size\n        \n        q_=torch.tensor(q_, dtype=int)\n        p_=torch.tensor(p_, dtype=int)\n        r_=torch.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        \n        label_mask=0\n        \n        if seq_len == 0:\n            q[0]=content_id\n            p[0]=part\n            r[0]=2\n            label_mask=0\n        elif seq_len == self.max_seq:\n            q[:-1]=q_[1:]\n            p[:-1]=p_[1:]\n            r[:]=r_[:]\n            q[-1]=content_id\n            p[-1]=part\n            label_mask=seq_len-1\n        else:\n            q[:seq_len]=q_[:]\n            p[:seq_len]=p_[:]\n            r[1:seq_len+1]=r_[:]\n            \n            q[seq_len]=content_id\n            p[seq_len]=part\n            r[0]=2\n            label_mask=seq_len\n        return (q, p, r, label_mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def update_group(test_df, prev_test_df):\n    if prev_test_df is None:\n        return\n    \n    prev_answered_correctly=eval(test_df.prior_group_answers_correct.values[0])\n    prev_test_df['answered_correctly']=prev_answered_correctly\n    prev_test_df=prev_test_df[prev_test_df.content_type_id==0]\n    \n    \n    prev_group=prev_test_df.groupby('user_id').apply(lambda row: (row.content_id.values[-100:],\n                                                                  row.part.values[-100:],\n                                                                  row.answered_correctly.values[-100:]))\n    \n    for user_id in prev_group.index.values:\n        if user_id not in group.index:\n            group[user_id]=prev_group[user_id]\n        else:\n            (prev_q, prev_p, prev_r)=prev_group[user_id]\n            group[user_id]=(\n                np.append(group[user_id][0], prev_q),\n                np.append(group[user_id][1], prev_p),\n                np.append(group[user_id][2], prev_r),\n            )\n            \n        if len(group[user_id][0]) > 100:\n            new_q=group[user_id][0][-100:]\n            new_p=group[user_id][1][-100:]\n            new_r=group[user_id][2][-100:]\n            group[user_id]=(new_q, new_p, new_r)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"prev_test_df=None\n\nfor (test_df, sample_prediction_df) in iter_test:\n    test_df=test_df[['row_id', 'user_id', 'content_id', \n                     'content_type_id', 'prior_group_answers_correct']].merge(\n        questions_df[['content_id', 'part']],how='left',on='content_id')\n    test_df.part=test_df.part.fillna(5)\n    \n    update_group(test_df, prev_test_df)\n    prev_test_df=test_df.copy()\n\n\n    test_df=test_df[test_df.content_type_id==0]\n    test_df['prev_seq']=test_df.user_id.apply(lambda user_id: group[user_id] if user_id in group else (np.array([]), np.array([]), np.array([])))\n    test_dataset=TestDataset(test_df, max_seq=100)\n    test_dataloader=torch.utils.data.DataLoader(test_dataset,\n                                                batch_size=40960, shuffle=False, pin_memory=True, num_workers=4)\n\n\n    model.eval()\n    y_answered=[]\n    \n    with torch.no_grad():\n        for (q, p, r, label_mask) in test_dataloader:\n            q=q.to(device)\n            p=p.to(device)\n            r=r.to(device)\n            y=model(q, p, r)\n            \n            y_answered.extend([torch.sigmoid(y[idx][label_id]).cpu().item() for idx, label_id in enumerate(label_mask)])\n    test_df['answered_correctly']=y_answered\n    env.predict(test_df[['row_id', 'answered_correctly']])\n    del test_df\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"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}