{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport gc\nimport riiideducation\nfrom collections import defaultdict\nfrom tqdm.notebook import tqdm\nimport pickle\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport random\nfrom sklearn.metrics import roc_auc_score\nimport matplotlib.pyplot as plt\n\nrandom.seed(34)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# Read questions\nq_lookup_df = pd.read_pickle(\"../input/riiid-fe-part-i/q_lookup_df.pickle\")\n#\nq_lookup_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SEQ_LENGTH = 100\nseq_prep = SEQ_LENGTH - 1\nn_question = q_lookup_df.content_id.max()\nn_bundle = q_lookup_df.bundle_id.max() + 1\nn_qet = q_lookup_df.q_encoded_tags.max() + 1\nn_part = q_lookup_df.part.nunique()\nn_q_attempt = 5 # max value\nlag_time_bins = 143\nelapsed_time_bins = 27\n#\nn_answer = 2\nprint(\"lag_time_bins, elapsed_time_bins\", lag_time_bins, elapsed_time_bins)\nprint(\"number of questions, bundles, encoded_tags, parts, n_q_attempt, n_answer\", n_question, n_bundle, n_qet, n_part, n_q_attempt, n_answer)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FCN(nn.Module):\n    def __init__(self, embed_dim):\n        super(FCN, self).__init__()\n        ###################################### Layers ######################################\n        self.linear1 = nn.Linear(embed_dim, embed_dim)\n        self.relu = nn.ReLU()\n        self.linear2 = nn.Linear(embed_dim, embed_dim)\n        self.dropout = nn.Dropout(0.2)\n    \n    def forward(self, x):\n        x = self.linear1(x)\n        x = self.relu(x)\n        x = self.linear2(x)\n        return self.dropout(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTEncoder(nn.Module):\n    def __init__(self, seq_len, embed_dim, num_heads, dropout, device=\"cpu\"):\n        super(SAINTEncoder, self).__init__()\n        self.seq_len = seq_len\n        self.embed_dim = embed_dim\n        self.device=device\n        ###################################### Layers ######################################\n        self.multi_att1 = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)\n        #\n        self.layer_norm1 = nn.LayerNorm(embed_dim)\n        self.layer_norm2 = nn.LayerNorm(embed_dim)\n        #\n        self.fcn = FCN(embed_dim)\n        \n    def forward(self, X, attn_mask):\n        # Permute operation must before and for multi head attention\n        X = X.permute(1, 0, 2) # [bs, s_len, embed] => [s_len, bs, embed]\n        norm_X = self.layer_norm1(X)\n        #\n        att_output, _ = self.multi_att1(norm_X, norm_X, norm_X, attn_mask=attn_mask)\n        #\n        M = (att_output + X).permute(1, 0, 2) # [s_len, bs, embed] => [bs, s_len, embed]\n        # Norm 2\n        norm_M = self.layer_norm2(M)\n        #\n        fcn_output = self.fcn(norm_M)\n        #\n        return fcn_output + M","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTDecoder(nn.Module):\n    def __init__(self, seq_len, embed_dim, num_heads, dropout, device=\"cpu\"):\n        super(SAINTDecoder, self).__init__()\n        self.seq_len = seq_len\n        self.embed_dim = embed_dim\n        self.device=device\n        ###################################### Layers ######################################\n        self.multi_att1 = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)\n        self.multi_att2 = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout)\n        #\n        self.layer_norm1 = nn.LayerNorm(embed_dim)\n        self.layer_norm2 = nn.LayerNorm(embed_dim)\n        self.layer_norm3 = nn.LayerNorm(embed_dim)\n        self.layer_norm4 = nn.LayerNorm(embed_dim)\n        #\n        self.fcn = FCN(embed_dim)\n        \n    def forward(self, X, encoder_output, attn_mask):\n        #\n        encoder_output = encoder_output.permute(1, 0, 2) # [bs, s_len, embed] => [s_len, bs, embed]\n        X = X.permute(1, 0, 2) # [bs, s_len, embed] => [s_len, bs, embed]\n        norm_X = self.layer_norm1(X)\n        #\n        att_output, _ = self.multi_att1(norm_X, norm_X, norm_X, attn_mask=attn_mask)\n        #\n        M1 = att_output + X\n        # Norm 2, 3\n        norm_M1 = self.layer_norm2(M1)\n        norm_encoder_output = self.layer_norm3(encoder_output)\n        #\n        att_output, _ = self.multi_att2(norm_M1, norm_encoder_output, norm_encoder_output, attn_mask=attn_mask)\n        #\n        M2 = (att_output + M1).permute(1, 0, 2) # [s_len, bs, embed] => [bs, s_len, embed]\n        # Norm 4\n        norm_M2 = self.layer_norm4(M2)\n        # Final output\n        fcn_output = self.fcn(norm_M2)\n        #\n        return fcn_output + M2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SAINTModel(nn.Module):\n    def __init__(self, seq_len=SEQ_LENGTH, n_question=n_question, n_answer=n_answer, \n                 n_bundle=n_bundle, n_part=n_part, n_q_attempt=n_q_attempt, n_qet=n_qet,\n                 lag_time_bins=lag_time_bins, elapsed_time_bins=elapsed_time_bins,\n                 embed_dim=256, n_head=8, dropout=0.1, device=\"cpu\"):\n        #\n        super(SAINTModel, self).__init__()\n        self.seq_len = seq_len\n        self.embed_dim = embed_dim\n        self.device = device\n        self.n_head = n_head\n        #\n        self.pos_embedding = nn.Embedding(seq_len + 1, embed_dim)\n        # known\n        self.question_embedding = nn.Embedding(n_question + 1, embed_dim)\n        self.part_embedding = nn.Embedding(n_part + 1, embed_dim)\n        self.bundle_embedding = nn.Embedding(n_bundle + 1, embed_dim)\n        self.q_attempt_embedding = nn.Embedding(n_q_attempt + 1, embed_dim)\n        self.lag_embedding = nn.Embedding(lag_time_bins + 1, embed_dim)\n        self.qet_embedding = nn.Embedding(n_qet + 1, embed_dim)\n        self.h_mean_embedding = nn.Linear(1, embed_dim, bias=False)\n        # future\n        # 0, 1 values, 2 for padding and mask\n        self.answer_embedding = nn.Embedding(n_answer + 1, embed_dim)\n        self.elapsed_embedding = nn.Embedding(elapsed_time_bins + 1, embed_dim)        \n        #\n        ###################################### Layers ######################################\n        self.layer_norm1 = nn.LayerNorm(embed_dim)\n        # Encoders\n        self.encoder1 = SAINTEncoder(seq_len, embed_dim, num_heads=self.n_head, dropout=dropout, device=device)\n        self.encoder2 = SAINTEncoder(seq_len, embed_dim, num_heads=self.n_head, dropout=dropout, device=device)\n        # Decoders\n        self.decoder1 = SAINTDecoder(seq_len, embed_dim, num_heads=self.n_head, dropout=dropout, device=device)\n        self.decoder2 = SAINTDecoder(seq_len, embed_dim, num_heads=self.n_head, dropout=dropout, device=device)\n        #\n        self.pred = nn.Linear(embed_dim, 1)\n    \n    def forward(self, x):\n        task_ids = x[0]\n        # Longs\n        question_ids = x[1]\n        bundle_ids = x[2]\n        parts = x[3]\n        q_attempts = x[4]\n        lag_times = x[5]\n        qets = x[6]\n        elapsed_times = x[7]\n        answers = x[8]\n        # Floats\n        h_means = x[9]\n        #\n        pos1 = torch.arange(1, self.seq_len+1, device=self.device).unsqueeze(0)\n        ab_pos1_em = self.pos_embedding(pos1)\n        #\n        q_em = self.question_embedding(question_ids)\n        bundle_em = self.bundle_embedding(bundle_ids)\n        part_em = self.part_embedding(parts)\n        q_att_em = self.q_attempt_embedding(q_attempts)\n        lag_em = self.lag_embedding(lag_times)\n        qet_em = self.qet_embedding(qets)\n        h_mean_em = self.h_mean_embedding(h_means.view(-1, self.seq_len, 1))        \n        known_feats = ab_pos1_em + q_em + bundle_em + part_em +\\\n                        q_att_em + lag_em + qet_em + h_mean_em\n        #\n        pos2 = torch.arange(0, self.seq_len, device=self.device).unsqueeze(0)\n        ab_pos2_em = self.pos_embedding(pos2)\n        #\n        answer_bundle_ids = torch.roll(bundle_ids.detach().clone(), 1, dims=1)\n        answer_bundle_ids[:, 0] = 0\n        answer_bundle_em = self.bundle_embedding(answer_bundle_ids)\n        aq_em = self.answer_embedding(answers)\n        elapsed_em = self.elapsed_embedding(elapsed_times)\n        future_feats = ab_pos2_em + aq_em + elapsed_em + answer_bundle_em\n        #\n        attn_mask = torch.triu(torch.ones((self.seq_len, self.seq_len), device=self.device), diagonal=1).bool()\n        #\n        encoder_output = self.encoder1(known_feats, attn_mask)\n        encoder_output = self.encoder2(encoder_output, attn_mask)\n        #\n        output = self.decoder1(future_feats, encoder_output, attn_mask)\n        output = self.decoder2(output, encoder_output, attn_mask)\n        #\n        output = self.layer_norm1(output)\n        output = self.pred(output)\n        return output.squeeze(-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = SAINTModel(device=device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)\ncriterion = nn.BCEWithLogitsLoss()\n\nmodel.to(device)\ncriterion.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"################################ INIT ################################\ntc = torch.zeros((3, SEQ_LENGTH), device=device)\nq = torch.zeros((3, SEQ_LENGTH), device=device)\nb = torch.zeros((3, SEQ_LENGTH), device=device)\np = torch.zeros((3, SEQ_LENGTH), device=device)\nqatt = torch.zeros((3, SEQ_LENGTH), device=device)\nlag = torch.zeros((3, SEQ_LENGTH), device=device)\nqet = torch.zeros((3, SEQ_LENGTH), device=device)\nelapsed = torch.zeros((3, SEQ_LENGTH), device=device)\naq = torch.zeros((3, SEQ_LENGTH), device=device)\nh_mean = torch.zeros((3, SEQ_LENGTH), device=device)\noutput = model([\n    # Longs\n    tc.long(), q.long(), b.long(), p.long(), \n    qatt.long(), lag.long(), qet.long(),\n    elapsed.long(), aq.long(),\n    # Floats\n    h_mean.float(),\n])\noutput.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch(model, train_iterator, optim, criterion, device=\"cpu\"):\n    model.train()\n    #\n    total_loss = []\n    labels = []\n    preds = []\n    result = 0\n    total = 0\n    #\n    progress = tqdm(train_iterator)\n    for item in progress:\n        inputs = [inp.to(device).long() for inp in item[0][:-1]] +\\\n                    [inp.to(device).float() for inp in item[0][-1:]]\n        label = item[1].to(device).float()\n        mask = item[2].to(device).bool()\n        #\n        optim.zero_grad()\n        output = model(inputs)\n        #\n        loss = criterion(output[mask], label[mask])\n        loss.backward()\n        optim.step()\n        total_loss.append(loss.detach().item())\n        #\n        output = output[:, -1]\n        label = label[:, -1]\n        #\n        pred = (torch.sigmoid(output) >= 0.5).long()\n        #\n        result += (pred == label).sum().item()\n        total += len(label)\n        #\n        labels.extend(label.detach().cpu().numpy())\n        preds.extend(output.detach().cpu().numpy())\n        #\n        progress.set_description('loss - {:.4f}, acc - {:.4f}'.format(loss, result / total))\n    #\n    acc = result / total\n    auc = roc_auc_score(labels, preds)\n    loss = np.mean(total_loss)\n    #\n    return loss, acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def valid_epoch(model, val_dataset, device=\"cpu\"):\n    model.eval()\n    #\n    batch_size = 1024\n    features, org_labels = val_dataset\n    total_epoch = int(np.ceil(len(features[0]) / batch_size))\n    #\n    labels = []\n    preds = []\n    result = 0\n    total = 0\n    #\n    progress = tqdm(range(total_epoch))\n    for i in progress:\n        s = i * batch_size\n        e = s + batch_size\n        #\n        inputs = [torch.from_numpy(inp[s:e]).to(device).long() for inp in features[:-1]] +\\\n                    [torch.from_numpy(inp[s:e]).to(device).float() for inp in features[-1:]]\n        label = torch.from_numpy(org_labels[s:e]).to(device).float()\n        #\n        with torch.no_grad():\n            output = model(inputs)\n        #\n        output = torch.sigmoid(output[:, -1])\n        #\n        pred = (output >= 0.5).long()\n        #\n        result += (pred == label).sum().item()\n        total += len(label)\n        #\n        labels.extend(label.detach().cpu().numpy())\n        preds.extend(output.detach().cpu().numpy())\n        #\n        progress.set_description('acc - {:.4f}'.format(result / total))\n    #\n    acc = result / total\n    auc = roc_auc_score(labels, preds)\n    return acc, auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_model(path, model, optimizer, scheduler, epoch, loss):\n    torch.save({\n        'epoch': epoch + 1, # beacuse last epoch already completed\n        'model_state_dict': model.state_dict(),\n        'optimizer_state_dict': optimizer.state_dict(),\n        'scheduler_state_dict': scheduler.state_dict(),\n        'loss': loss,\n    }, path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nepochs = 25\nfor epoch in range(epochs):\n    loss, acc, auc = train_epoch(model, dataloader, optimizer, criterion, device)\n    print(\"{}/{} train_loss - {:.3f} train_acc - {:.4f} train_auc - {:.4f}\".format(epoch+1, epochs, loss, acc, auc))\n    acc, auc = valid_epoch(model, val_dataset, device)\n    print(\"{}/{} val_acc - {:.4f} val_auc - {:.4f}\".format(epoch+1, epochs, acc, auc))\n    if epoch >= 15:\n        save_model(\"SAINT_model_{}.pt\".format(epoch), model, optimizer, scheduler, epoch, loss)\n    scheduler.step()\n    print(\"lr:\", scheduler.get_last_lr())\n\"\"\"","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}