{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport time\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\n\nfrom pytorch_lightning.metrics.functional.classification import auroc","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"## Preprocessing"},{"metadata":{"trusted":true},"cell_type":"code","source":"def add_question_data(data):\n    questions = pd.read_csv(\"../input/riiid-test-answer-prediction/questions.csv\")\n    q_part_ids_map = dict(zip(questions[\"question_id\"], questions[\"part\"]))\n    data[\"part\"] = data[\"content_id\"].map(q_part_ids_map).astype(np.int64)\n    return data\n\n\ndef preprocessing(data, n_sample=100_000_000):\n    data = data.tail(n_sample)\n    data = data[data[\"content_type_id\"] == 0]\n    data = data[[\"row_id\", \"user_id\", \"content_id\", \"answered_correctly\"]]\n    data = add_question_data(data)\n\n    data = data.groupby(\"user_id\").apply(\n        lambda row: (\n            row[\"content_id\"].values,\n            row[\"part\"].values,\n            row[\"answered_correctly\"].values,\n        )\n    )\n    # Drop <= 5 questions answered.\n    data = data[data.apply(lambda x: x[0].shape[0]) > 5]\n    return data","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_path = \"../input/cv-strategy-in-the-kaggle-environment/cv5_train.parquet\"\nvalid_path = \"../input/cv-strategy-in-the-kaggle-environment/cv5_valid.parquet\"\nuse_cols = [\"row_id\", \"user_id\", \"content_id\", \"content_type_id\", \"answered_correctly\"]\n\ntrain = pd.read_parquet(train_path, columns=use_cols)\ntrain = preprocessing(train, n_sample=10_000_000)\n\nvalid = pd.read_parquet(valid_path, columns=use_cols)\nvalid = preprocessing(valid, n_sample=2_500_000)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Train shape: \", train.shape[0])\ndisplay(train.head())\n\nprint(\"Train shape: \", valid.shape[0])\ndisplay(valid.head())","execution_count":null,"outputs":[]},{"metadata":{},"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        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\n        src_mask = torch.ones(self.max_seq, dtype=bool)\n        label_mask = torch.ones(self.max_seq, dtype=bool)\n\n        src_mask[-seq_len:] = False\n        label_mask[-seq_len:] = False\n\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            # 2-for the start of the sequence\n            r[-seq_len:] = torch.cat((torch.tensor([2]), r_[: seq_len - 1]))\n            y[-seq_len:] = r_\n\n        return (q, c, r, y, src_mask, label_mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 128\ntrain_dataset = SAINTDataset(train)\ntrain_dataloader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=True,\n    num_workers=4,\n    pin_memory=True,\n)\nval_dataset = SAINTDataset(valid)\nval_dataloader = torch.utils.data.DataLoader(\n    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":"## Define Model (SAINT)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class EncoderEmbedding(nn.Module):\n    def __init__(self, n_content, n_part, n_dims, seq_len, device):\n        super(EncoderEmbedding, self).__init__()\n        self.n_dims = n_dims\n        self.seq_len = seq_len\n        self.device = device\n\n        self.position_embed = nn.Embedding(seq_len, n_dims)\n        self.content_embed = nn.Embedding(n_content, n_dims)\n        self.part_embed = nn.Embedding(n_part, n_dims)\n        \n        torch.nn.init.normal_(self.position_embed.weight, mean=0.0, std=0.01)\n        torch.nn.init.normal_(self.content_embed.weight, mean=0.0, std=0.01)\n        torch.nn.init.normal_(self.part_embed.weight, mean=0.0, std=0.01)\n\n    def forward(self, content_id, part_id):\n        seq = torch.arange(self.seq_len, device=self.device).unsqueeze(0)\n        pos = self.position_embed(seq)\n\n        content = self.content_embed(content_id)\n        part = self.part_embed(part_id)\n        return pos + content + part\n\n\nclass DecoderEmbedding(nn.Module):\n    def __init__(self, n_response, n_dims, seq_len, device):\n        super(DecoderEmbedding, self).__init__()\n        self.n_dims = n_dims\n        self.seq_len = seq_len\n        self.device = device\n\n        self.position_embed = nn.Embedding(seq_len, n_dims)\n        self.response_embed = nn.Embedding(n_response, n_dims)\n        \n        torch.nn.init.normal_(self.position_embed.weight, mean=0.0, std=0.01)\n        torch.nn.init.normal_(self.response_embed.weight, mean=0.0, std=0.01)\n\n    def forward(self, response):\n        seq = torch.arange(self.seq_len, device=self.device).unsqueeze(0)\n        pos = self.position_embed(seq)\n\n        res = self.response_embed(response)\n        return pos + res\n\n\nclass SAINTModel(nn.Module):\n    def __init__(\n        self,\n        n_questions,\n        n_categories,\n        n_responses,\n        device=\"cpu\",\n        max_seq=100,\n        d_model=128,\n        encoder_dim=64,\n        decoder_dim=64,\n        num_heads=2,\n    ):\n        super().__init__()\n        self.device = device\n        self.mask = self.generate_square_subsequent_mask(max_seq)\n        self.encoder_embedding = EncoderEmbedding(\n            n_content=n_questions,\n            n_part=n_categories,\n            n_dims=d_model,\n            seq_len=max_seq,\n            device=device,\n        )\n        self.decoder_embedding = DecoderEmbedding(\n            n_response=n_responses,\n            n_dims=d_model,\n            seq_len=max_seq,\n            device=device,\n        )\n\n        self.transformer = nn.Transformer(\n            d_model=d_model,\n            nhead=num_heads,\n            num_encoder_layers=4,\n            num_decoder_layers=4,\n            dim_feedforward=1024,\n            dropout=0.1,\n            activation=\"relu\",\n        )\n        self.ffn = torch.nn.Sequential(\n            torch.nn.Linear(d_model, d_model),\n            torch.nn.ReLU(),\n            torch.nn.Linear(d_model, d_model),\n            torch.nn.Dropout(p=0.1),\n        )\n        self.layer_norm = torch.nn.LayerNorm(d_model)\n        self.fc1 = nn.Linear(in_features=d_model, out_features=1)\n\n    def forward(self, q, c, r, src_pad_mask, tgt_pad_mask):\n        mask = self.mask.to(q.device)\n        enc = self.encoder_embedding(\n            content_id=q,\n            part_id=c,\n        )\n        enc = enc.transpose(0, 1)\n        dec = self.decoder_embedding(\n            response=r,\n        )\n        dec = dec.transpose(0, 1)\n        x = self.transformer(\n            enc,\n            dec,\n            src_mask=mask,\n            tgt_mask=mask,\n            # src_key_padding_mask=src_pad_mask,\n            # tgt_key_padding_mask=tgt_pad_mask,\n        )\n        x = self.ffn(x) + x\n        x = self.fc1(self.layer_norm(x))\n        x = x.transpose(0, 1)\n        return x.squeeze(-1)\n    \n    def generate_square_subsequent_mask(self, sz):\n        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n        return mask\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# (q, c, r, y, src_mask, label_mask) = next(iter(train_dataloader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define paramteres.\nn_questions = 13523\nn_categories = 8\nn_responses = 3\n\n# Train model\ndevice = \"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\nmodel.to(device)\ncriterion.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_epoch():\n    preds = []\n    labels = []\n    train_loss = []\n    model.train()\n\n    for i, (q, c, r, y, src_mask, label_mask) in enumerate(train_dataloader):\n        q = q.to(device)\n        c = c.to(device)\n        r = r.to(device)\n        y = y.to(device)\n        src_mask = src_mask.to(device)\n        label_mask = label_mask.to(device)\n\n        optimizer.zero_grad()\n        yout = model(q, c, r, src_mask, label_mask)\n\n        yout = torch.masked_select(yout, torch.logical_not(label_mask))\n        y = torch.masked_select(y, torch.logical_not(label_mask))\n\n        yout = yout.float()\n        y = y.float()\n        \n        preds.append(yout.clone().view(-1))\n        labels.append(y.clone().view(-1))\n\n        loss_ = criterion(yout, y)\n        loss_.backward()\n        optimizer.step()\n        train_loss.append(loss_.item())\n        \n    pred = torch.cat(preds, dim=0)\n    label = torch.cat(labels, dim=0)\n    auc = auroc(pred, label)\n\n    return np.mean(train_loss), auc\n\ndef val_epoch():\n    preds = []\n    labels = []\n    val_loss = []\n    model.eval()\n\n    with torch.no_grad():\n        for (q, c, r, y, src_mask, label_mask) in val_dataloader:\n            q = q.to(device)\n            c = c.to(device)\n            r = r.to(device)\n            y = y.to(device)\n            src_mask = src_mask.to(device)\n            label_mask = label_mask.to(device)\n            yout = model(q, c, r, src_mask, label_mask)\n\n            yout = torch.masked_select(yout, torch.logical_not(label_mask))\n            y = torch.masked_select(y, torch.logical_not(label_mask))\n\n            yout = yout.float()\n            y = y.float()\n            \n            preds.append(yout.clone().view(-1))\n            labels.append(y.clone().view(-1))\n\n            loss_ = criterion(yout, y)\n            val_loss.append(loss_.item())\n            \n    pred = torch.cat(preds, dim=0)\n    label = torch.cat(labels, dim=0)\n    auc = auroc(pred, label)\n\n    return np.mean(val_loss), auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs = 50\nbest_score = None\nfor i in range(num_epochs):\n    epoch_start = time.time()\n\n    train_loss, train_auc = train_epoch()\n    val_loss, val_auc = val_epoch()\n\n    epoch_end = time.time()\n    # print(\"Time To Run Epoch:{}\".format((epoch_end - epoch_start) / 60))\n    print(f\"Epoch:{i} | Train Loss:{train_loss:.6f} | Train AUC:{train_auc:.6f} | Val Loss:{val_loss:.6f} | Val ACU:{val_auc:.6f}\")\n\n    if (best_score is None) or (best_score > val_loss):\n        best_score = val_loss\n        best_epoch = i\n        torch.save(model.state_dict(), \"saint{}.pth\".format(i))\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Evaluation"},{"metadata":{"trusted":true},"cell_type":"code","source":"eval_path = \"../input/cv-strategy-in-the-kaggle-environment/cv1_valid.parquet\"\nuse_cols = [\"row_id\", \"user_id\", \"content_id\", \"content_type_id\", \"answered_correctly\"]\n\neval_data = pd.read_parquet(eval_path, columns=use_cols)\neval_data = preprocessing(eval_data)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Train shape: \", eval_data.shape[0])\ndisplay(eval_data.head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 128\neval_dataset = SAINTDataset(eval_data)\neval_dataloader = torch.utils.data.DataLoader(\n    eval_dataset,\n    batch_size=batch_size,\n    num_workers=4,\n    shuffle=False,\n    pin_memory=True,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = SAINTModel(n_questions, n_categories, n_responses, device=device)\nmodel.load_state_dict(torch.load(f\"saint{best_epoch}.pth\"))\nmodel.to(device)\n\npreds = []\nlabels = []\n\nmodel.eval()\nwith torch.no_grad():\n    for (q, c, r, y, src_mask, label_mask) in eval_dataloader:\n        q = q.to(device)\n        c = c.to(device)\n        r = r.to(device)\n        y = y.to(device)\n        src_mask = src_mask.to(device)\n        label_mask = label_mask.to(device)\n        yout = model(q, c, r, src_mask, label_mask)\n\n        yout = torch.masked_select(yout, torch.logical_not(label_mask))\n        y = torch.masked_select(y, torch.logical_not(label_mask))\n\n        preds.append(yout.float().view(-1))\n        labels.append(y.float().view(-1))\n        \npred = torch.cat(preds, dim=0)\nlabel = torch.cat(labels, dim=0)\nauc = auroc(pred, label)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"Evaluation AUC is {auc}\")","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}