{"cells":[{"metadata":{},"cell_type":"markdown","source":" Please Go through the code and help me resolve crash issue. \n Fork and improve this code implementation."},{"metadata":{},"cell_type":"markdown","source":"Paper:   [SAINT+: Integrating Temporal Features for EdNet Correctness Prediction](https://arxiv.org/abs/2010.12042)"},{"metadata":{},"cell_type":"markdown","source":"# Import everything now..."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn\nimport pandas as pd\nimport numpy as np\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import train_test_split\nimport torch\nimport numpy as np \nfrom torch import nn \nimport copy \nimport pytorch_lightning as pl \nimport pandas as pd\nimport numpy as np\nfrom torch.utils.data import Dataset, DataLoader\nimport gc\nfrom sklearn.model_selection import train_test_split \nfrom sklearn.metrics import roc_auc_score\nimport torch.nn.functional as F\nimport gc ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Configure constants"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class config:\n        device = torch.device(\"cuda\") \n        MAX_SEQ = 100 \n        EMBED_DIMS = 512 \n        ENC_HEADS = DEC_HEADS = 8\n        NUM_ENCODER = NUM_DECODER = 4\n        BATCH_SIZE = 64\n        TRAIN_FILE = \"../input/riiid-test-answer-prediction/train.csv\"\n        TOTAL_EXE = 13523\n        TOTAL_CAT = 10000 ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DKTDataset(Dataset):\n  def __init__(self,samples,max_seq,start_token=0): \n    super().__init__()\n    self.samples = samples\n    self.max_seq = max_seq\n    self.start_token = start_token\n    self.data = []\n    for id in self.samples.index:\n      exe_ids,answers,ela_time,categories = self.samples[id]\n      if len(exe_ids)>max_seq:\n        for l in range((len(exe_ids)+max_seq-1)//max_seq):\n            self.data.append((exe_ids[l:l+max_seq],answers[l:l+max_seq],ela_time[l:l+max_seq],categories[l:l+max_seq]))\n      elif len(exe_ids)<self.max_seq and len(exe_ids)>10:\n            self.data.append((exe_ids,answers,ela_time,categories))\n      else :\n            continue\n\n  def __len__(self):\n    return len(self.data)\n  \n  def __getitem__(self,idx):\n    question_ids,answers,ela_time,exe_category = self.data[idx]\n    seq_len = len(question_ids)\n\n    exe_ids = np.zeros(self.max_seq,dtype=int)\n    ans = np.zeros(self.max_seq,dtype=int)\n    elapsed_time = np.zeros(self.max_seq,dtype=int)\n    exe_cat = np.zeros(self.max_seq,dtype=int)\n    if seq_len<self.max_seq:\n      exe_ids[-seq_len:] = question_ids\n      ans[-seq_len:] = answers\n      elapsed_time[-seq_len:] = ela_time \n      exe_cat[-seq_len:] = exe_category\n    else:\n      exe_ids[:] = question_ids[-self.max_seq:]\n      ans[:] = answers[-self.max_seq:]\n      elapsed_time[:] = ela_time[-self.max_seq:]\n      exe_cat[:] = exe_category[-self.max_seq:]\n\n    input_rtime = np.zeros(self.max_seq,dtype=int)\n    input_rtime = np.insert(elapsed_time,0,self.start_token)\n    input_rtime = np.delete(input_rtime,-1)\n\n    input = {\"input_ids\":exe_ids,\"input_rtime\":input_rtime.astype(np.int),\"input_cat\":exe_cat}\n    answers = np.append([0],ans[:-1]) #start token\n    assert ans.shape[0]==answers.shape[0] and answers.shape[0]==input_rtime.shape[0], \"both ans and label should be \\\n                                                                                            same len with start-token\"\n    return input,answers,ans\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SAINT+ model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FFN(nn.Module):\n  def __init__(self,in_feat):\n    super(FFN,self).__init__()\n    self.linear1 = nn.Linear(in_feat,in_feat)\n    self.linear2 = nn.Linear(in_feat,in_feat)\n    self.drop = nn.Dropout(0.2)\n  \n  def forward(self,x):\n    out = F.relu(self.drop(self.linear1(x)))\n    out = self.linear2(out)\n    return out \n\n\nclass EncoderEmbedding(nn.Module):\n  def __init__(self,n_exercises,n_categories,n_dims,seq_len):\n    super(EncoderEmbedding,self).__init__()\n    self.n_dims = n_dims\n    self.seq_len = seq_len\n    self.exercise_embed = nn.Embedding(n_exercises,n_dims)\n    self.category_embed = nn.Embedding(n_categories,n_dims)\n    self.position_embed = nn.Embedding(seq_len,n_dims)\n\n  def forward(self,exercises,categories):\n    e = self.exercise_embed(exercises)\n    c = self.category_embed(categories)\n    seq = torch.arange(self.seq_len,device=config.device).unsqueeze(0)\n    p = self.position_embed(seq)\n    return p + c + e\n\nclass DecoderEmbedding(nn.Module):\n  def __init__(self,n_responses,n_dims,seq_len):\n    super(DecoderEmbedding,self).__init__()\n    self.n_dims = n_dims\n    self.seq_len = seq_len\n    self.response_embed = nn.Embedding(n_responses,n_dims)\n    self.time_embed = nn.Linear(1,n_dims,bias=False)\n    self.position_embed = nn.Embedding(seq_len,n_dims)\n\n  def forward(self,responses):\n    e = self.response_embed(responses)\n    seq = torch.arange(self.seq_len,device=config.device).unsqueeze(0)\n    p = self.position_embed(seq)\n    return p + e \n\n\n# layers of encoders stacked onver, multiheads-block in each encoder is n.\n# Stacked N MultiheadAttentions \nclass StackedNMultiHeadAttention(nn.Module):\n  def __init__(self,n_stacks,n_dims,n_heads,seq_len,n_multihead=1,dropout=0.2):\n    super(StackedNMultiHeadAttention,self).__init__()\n    self.n_stacks = n_stacks\n    self.n_multihead = n_multihead\n    self.n_dims = n_dims \n    self.norm_layers = nn.LayerNorm(n_dims)\n    #n_stacks has n_multiheads each\n    self.multihead_layers = nn.ModuleList(n_stacks*[nn.ModuleList(n_multihead*[nn.MultiheadAttention(embed_dim = n_dims,\n                                                      num_heads = n_heads,\n                                                        dropout = dropout),]),])\n    self.ffn = nn.ModuleList(n_stacks*[FFN(n_dims)])\n    self.mask = torch.triu(torch.ones(seq_len,seq_len),diagonal=1).to(dtype=torch.bool)\n  \n  def forward(self,input_q,input_k,input_v,encoder_output=None,break_layer=None):\n    for stack in range(self.n_stacks):\n        for multihead in range(self.n_multihead):\n          norm_q = self.norm_layers(input_q)\n          norm_k = self.norm_layers(input_k)\n          norm_v = self.norm_layers(input_v) \n          heads_output,_ = self.multihead_layers[stack][multihead](query=norm_q.permute(1,0,2),\n                                                                    key=norm_k.permute(1,0,2),\n                                                                    value=norm_v.permute(1,0,2),\n                                                                    attn_mask=self.mask.to(config.device))\n          heads_output = heads_output.permute(1,0,2)\n          #assert encoder_output != None and break_layer is not None     \n          if encoder_output != None and multihead == break_layer:\n            assert break_layer <= multihead, \" break layer should be less than multihead layers and postive integer\"\n            input_k = input_v = encoder_output\n            input_q =input_q + heads_output\n          else:\n            input_q =input_q+ heads_output\n            input_k =input_k+ heads_output\n            input_v =input_v +heads_output\n        last_norm = self.norm_layers(heads_output)\n        ffn_output = self.ffn[stack](last_norm)\n        ffn_output =ffn_output+ heads_output\n    return ffn_output\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Final Model with Trainer"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Main model for training \nclass PlusSAINTModule(pl.LightningModule):\n  def __init__(self):\n    super(PlusSAINTModule,self).__init__()\n    self.loss = nn.BCEWithLogitsLoss()\n    self.encoder_layer = StackedNMultiHeadAttention(n_stacks=config.NUM_DECODER,\n                                                    n_dims=config.EMBED_DIMS,\n                                                    n_heads=config.DEC_HEADS,\n                                                    seq_len=config.MAX_SEQ,\n                                                    n_multihead=1,dropout=0.2)\n    self.decoder_layer = StackedNMultiHeadAttention(n_stacks=config.NUM_ENCODER,\n                                                    n_dims=config.EMBED_DIMS,\n                                                    n_heads=config.ENC_HEADS,\n                                                    seq_len=config.MAX_SEQ,\n                                                    n_multihead=2,dropout=0.2)\n    self.encoder_embedding = EncoderEmbedding(n_exercises=config.TOTAL_EXE,\n                                              n_categories=config.TOTAL_CAT,\n                                              n_dims=config.EMBED_DIMS,seq_len=config.MAX_SEQ)\n    self.decoder_embedding = DecoderEmbedding(n_responses=3,n_dims=config.EMBED_DIMS,seq_len=config.MAX_SEQ)\n    self.elapsed_time = nn.Linear(1,config.EMBED_DIMS)\n    self.fc = nn.Linear(config.EMBED_DIMS,1)\n\n  #TODO: implement embdding layer and its output\n  def forward(self,x,y): \n    enc = self.encoder_embedding(exercises=x[\"input_ids\"].long().to(config.device),categories=x['input_cat'].long().to(config.device))\n    dec = self.decoder_embedding(responses=y.long().to(config.device))\n    elapsed_time=x[\"input_rtime\"].unsqueeze(-1).float()\n    ela_time = self.elapsed_time(elapsed_time)\n    dec = dec + ela_time\n    # this encoder \n    encoder_output = self.encoder_layer(input_k=enc,\n                                        input_q=enc,\n                                        input_v=enc)\n    #this is decoder\n    decoder_output = self.decoder_layer(input_k=dec,\n                                        input_q=dec,\n                                        input_v=dec,\n                                        encoder_output = encoder_output,\n                                        break_layer=1)\n    #fully connected layer\n    out = self.fc(decoder_output)\n    return out.squeeze()\n\n  def configure_optimizers(self):\n    return torch.optim.Adam(self.parameters(),0.0001)\n  \n  def training_step(self,batch,batch_ids):\n    input,ans,labels = batch\n    target_mask = (input[\"input_ids\"]!=0)\n    out = self(input,ans)\n    loss = self.loss(out.view(-1).float(),labels.view(-1).float()) \n    out = torch.masked_select(out,target_mask)\n    out = torch.sigmoid(out) \n    labels = torch.masked_select(labels,target_mask)    \n    self.log(\"train_loss\",loss,on_step=True,prog_bar=True)\n    return {\"loss\":loss,\"outs\":out,\"labels\":labels}\n  \n  def validation_step(self,batch,batch_ids):\n    input,ans,labels = batch\n    target_mask = (input[\"input_ids\"]!=0)\n    out = self(input,ans)\n    loss = self.loss(out.view(-1).float(),labels.view(-1).float())\n    out = torch.masked_select(out,target_mask)\n    out = torch.sigmoid(out) \n    labels = torch.masked_select(labels,target_mask) \n    self.log(\"val_loss\",loss,on_step=True,prog_bar=True)\n    output = {\"outs\":out,\"labels\":labels}\n    return output\n  \n  def validation_epoch_end(self,validation_ouput): \n    out = torch.cat([i[\"outs\"] for i in validation_ouput]).view(-1) \n    labels = torch.cat([i[\"labels\"] for i in validation_ouput]).view(-1)\n    auc = roc_auc_score(labels.cpu().detach().numpy(),out.cpu().detach().numpy())\n    self.print(\"val auc\",auc)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_dataloaders():              \n    dtypes = {'timestamp': 'int64', 'user_id': 'int32' ,'content_id': 'int16',\n                'answered_correctly':'int8',\"content_type_id\":\"int8\",\n                  \"prior_question_elapsed_time\":\"float32\",\"task_container_id\":\"int16\"}\n    print(\"loading csv.....\")\n    train_df = pd.read_csv(config.TRAIN_FILE,usecols=[1,2,3,4,5,7,8],dtype=dtypes)\n    print(\"shape of dataframe :\",train_df.shape) \n\n    train_df = train_df[train_df.content_type_id==0] \n    train_df.prior_question_elapsed_time.fillna(0,inplace=True)\n    train_df.prior_question_elapsed_time /=3600 \n    #train_df.prior_question_elapsed_time.clip(lower=0,upper=300,inplace=True)\n    train_df.prior_question_elapsed_time = train_df.prior_question_elapsed_time.astype(np.int)\n    \n    train_df = train_df.sort_values([\"timestamp\"],ascending=True).reset_index(drop=True)\n    n_skills = train_df.content_id.nunique() \n    print(\"no. of skills :\",n_skills)\n    print(\"shape after exlusion:\",train_df.shape)\n\n    #grouping based on user_id to get the data supplu\n    print(\"Grouping users...\") \n    group = train_df[[\"user_id\",\"content_id\",\"answered_correctly\",\"prior_question_elapsed_time\",\"task_container_id\"]]\\\n                    .groupby(\"user_id\")\\\n                    .apply(lambda r: (r.content_id.values,r.answered_correctly.values,\\\n                                      r.prior_question_elapsed_time.values,r.task_container_id.values))\n    del train_df\n    gc.collect() \n    print(\"splitting\") \n    train,val = train_test_split(group,test_size=0.2) \n    print(\"train size: \",train.shape,\"validation size: \",val.shape)\n    train_dataset = DKTDataset(train,max_seq = config.MAX_SEQ)\n    val_dataset = DKTDataset(val,max_seq = config.MAX_SEQ)\n    train_loader = DataLoader(train_dataset,\n                          batch_size=config.BATCH_SIZE,\n                          num_workers=8,\n                          shuffle=True) \n    val_loader = DataLoader(val_dataset,\n                          batch_size=config.BATCH_SIZE,\n                          num_workers=8,\n                          shuffle=False)\n    del train_dataset,val_dataset \n    gc.collect() \n    return train_loader, val_loader \ntrain_loader, val_loader = get_dataloaders() ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"saint_plus = PlusSAINTModule()\ntrainer = pl.Trainer(gpus=-1,max_epochs=3,progress_bar_refresh_rate=21) \ntrainer.fit(model=saint_plus,\n            train_dataloader=train_loader, \n            val_dataloaders = [val_loader,]) \ntrainer.save_checkpoint(\"model.pt\") ","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}