{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Transformer Encoder Only Model in PyTorch\n\n- Below is a draft from my ongoing implementation towards the SAINT and SAINT+ papers introduced by the organizers. It's a WIP and is based on my understanding. So for sure, there might be bugs in it and that's why you, being the reader should consider letting me know the same and be a responsible citizen of Kaggle's kingdom :)\n\n- Not added validation mechanism as i am little bit confused as of now, once that's cleared, i will update it for sure.\n\n- Host's papers (https://arxiv.org/abs/2002.07033, https://arxiv.org/abs/2010.12042)."},{"metadata":{"trusted":true},"cell_type":"code","source":"# import some typical packages and fix few random seed.\n# Plz fix the torch seed as well...\n# NB many imports are not used, but they are just there...\n\nimport numpy as np\nimport pandas as pd\nimport pickle\nimport gc\nimport time\nimport itertools\n\nfrom tqdm.notebook import tqdm\nfrom sklearn.metrics import roc_auc_score, auc, roc_curve\nfrom contextlib import contextmanager\nfrom pathlib import Path\nfrom time import time\nfrom collections import namedtuple\n\nRANDOM_SEED = 42\nnp.random.seed(RANDOM_SEED)\ngc.enable()\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\n\nfrom torch.nn.modules.dropout import Dropout\nfrom torch.nn.modules.activation import MultiheadAttention\nfrom torch.nn.modules.normalization import LayerNorm\n\nTRAIN_DTYPES = {\n    # 'row_id': np.uint32,\n    'timestamp': np.uint64,\n    'user_id': np.uint32,\n    'content_id': np.uint16,\n    'content_type_id': np.uint8,\n    'task_container_id': np.uint16,\n    'user_answer': np.int8,\n    'answered_correctly': np.int8,\n    'prior_question_elapsed_time': np.float32,\n    'prior_question_had_explanation': 'boolean'\n}\n\nDATA_DIR = Path('../input/riiid-test-answer-prediction')\nTRAIN_PATH = DATA_DIR / 'train.csv'\nQUESTIONS_PATH = DATA_DIR / 'questions.csv'\nLECTURES_PATH = DATA_DIR / 'lectures.csv'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nLAST_N = 100 # this parameter denotes how many last seen content_ids I am going to consider <aka the max_seq_len or the window size>.\n\ndf_questions = pd.read_csv(QUESTIONS_PATH)\ndf_train = pd.read_csv(TRAIN_PATH, nrows=40_00_000, dtype=TRAIN_DTYPES, usecols=TRAIN_DTYPES.keys())\n\ndf_train['prior_question_had_explanation'] = df_train['prior_question_had_explanation'].astype(np.float16).fillna(-1).astype(np.int8)\ndf_train = df_train[df_train.content_type_id == 0]\n\npart_ids_map = dict(zip(df_questions.question_id, df_questions.part))\ndf_train['part_id'] = df_train['content_id'].map(part_ids_map)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import Counter\ndf_train[\"prior_question_elapsed_time\"].fillna(26000, inplace=True) # some random value fill in\ndf_train[\"prior_question_elapsed_time\"] = df_train[\"prior_question_elapsed_time\"] // 1000","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Preparing the Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nfrom collections import deque\n\n# we will be using a deque as it automatically limits the max_size as per the Data Strucutre's defination itself\n# so we don't need to manage that...\n\nd = {}\nuser_id_to_idx = {}\n\nPAD = 0\n\ngrp = df_train.groupby(\"user_id\").tail(LAST_N) # Select last_n rows of each user.\n\nfor idx, row in tqdm(grp.groupby(\"user_id\").agg({\n    \"content_id\":list, \"answered_correctly\":list, \"task_container_id\":list, \n    \"part_id\":list, \"prior_question_elapsed_time\":list\n}).reset_index().iterrows()):\n    # here we make a split whether a user has more than equal to 100 entries or less than that\n    # if it's less than LAST_N, then we need to PAD it using the PAD token defined as 0 by me in this cell block\n    # also, padded will be True where we have done padding obviously, rest places it's False.\n    if len(row[\"content_id\"]) >= 100:\n        d[idx] = {\n            \"user_id\": row[\"user_id\"],\n            \"content_id\" : deque(row[\"content_id\"], maxlen=LAST_N),\n            \"answered_correctly\" : deque(row[\"answered_correctly\"], maxlen=LAST_N),\n            \"task_container_id\" : deque(row[\"task_container_id\"], maxlen=LAST_N),\n            \"prior_question_elapsed_time\" : deque(row[\"prior_question_elapsed_time\"], maxlen=LAST_N),\n            \"part_id\": deque(row[\"part_id\"], maxlen=LAST_N),\n            \"padded\" : deque([False]*100, maxlen=LAST_N)\n        }\n    else:\n        # we have to pad...\n        # (max_batch_len - len(seq))\n        d[idx] = {\n            \"user_id\": row[\"user_id\"],\n            \"content_id\" : deque(row[\"content_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"answered_correctly\" : deque(row[\"answered_correctly\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"task_container_id\" : deque(row[\"task_container_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"prior_question_elapsed_time\" : deque(row[\"prior_question_elapsed_time\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"part_id\": deque(row[\"part_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"padded\" : deque([False]*len(row[\"content_id\"]) + [True]*(100-len(row[\"content_id\"])), maxlen=LAST_N)\n        }\n    user_id_to_idx[row[\"user_id\"]] = idx\n    # if in future a new user comes, we will just increase the counts as of now... <WIP>","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(d[0]) # user_id 115; I encourage you to match the same with the dataframes.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# this user has been known to me as 0th user basically. If a user comes which is unseen in past, \n# simply hash it to the len of the dict and increment that value; keping this just for reference and debugging.\nuser_id_to_idx[115]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# A Minimal Transformer Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn import TransformerEncoder, TransformerEncoderLayer\n\nclass TransformerModel(nn.Module):\n\n    def __init__(self, ninp:int=32, nhead:int=2, nhid:int=64, nlayers:int=2, dropout:float=0.3):\n        '''\n        nhead -> number of heads in the transformer multi attention thing.\n        nhid -> the number of hidden dimension neurons in the model.\n        nlayers -> how many layers we want to stack.\n        '''\n        super(TransformerModel, self).__init__()\n        self.src_mask = None\n        encoder_layers = TransformerEncoderLayer(d_model=ninp, nhead=nhead, dim_feedforward=nhid, dropout=dropout, activation='relu')\n        self.transformer_encoder = TransformerEncoder(encoder_layer=encoder_layers, num_layers=nlayers)\n        self.exercise_embeddings = nn.Embedding(num_embeddings=13523, embedding_dim=ninp) # exercise_id\n        self.pos_embedding = nn.Embedding(ninp, ninp) # positional embeddings\n        self.part_embeddings = nn.Embedding(num_embeddings=7+1, embedding_dim=ninp) # part_id_embeddings\n        self.prior_question_elapsed_time = nn.Embedding(num_embeddings=301, embedding_dim=ninp) # prior_question_elapsed_time\n        self.device = \"cpu\"\n        self.ninp = ninp\n        self.decoder = nn.Linear(ninp, 2)\n        self.init_weights()\n\n    def init_weights(self):\n        initrange = 0.1\n        # init embeddings\n        self.exercise_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.part_embeddings.weight.data.uniform_(-initrange, initrange)\n        self.prior_question_elapsed_time.weight.data.uniform_(-initrange, initrange)\n        self.decoder.bias.data.zero_()\n        self.decoder.weight.data.uniform_(-initrange, initrange)\n\n    def forward(self, content_id, part_id, prior_question_elapsed_time=None, mask_src=None):\n        '''\n        S is the sequence length, N the batch size and E the Embedding Dimension (number of features).\n        src: (S, N, E)\n        src_mask: (S, S)\n        src_key_padding_mask: (N, S)\n        padding mask is (N, S) with boolean True/False.\n        SRC_MASK is (S, S) with float(’-inf’) and float(0.0).\n        '''\n\n        embedded_src = self.exercise_embeddings(content_id) + \\\n        self.pos_embedding(torch.arange(0, content_id.shape[1]).to(self.device).unsqueeze(0).repeat(content_id.shape[0], 1)) + \\\n        self.part_embeddings(part_id) + self.prior_question_elapsed_time(prior_question_elapsed_time) # (N, S, E)\n        embedded_src = embedded_src.transpose(0, 1) # (S, N, E)\n        \n        _src = embedded_src * np.sqrt(self.ninp)\n        \n        output = self.transformer_encoder(src=_src, src_key_padding_mask=mask_src)\n        output = self.decoder(output)\n        output = output.transpose(1, 0)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from typing import List\n\ndef pad_seq(seq: List[int], max_batch_len: int = LAST_N, pad_value: int = True) -> List[int]:\n    return seq + (max_batch_len - len(seq)) * [pad_value]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# pytorch dataset class\n\nclass Riiid(torch.utils.data.Dataset):\n    \n    def __init__(self, d):\n        self.d = d\n    \n    def __len__(self):\n        return len(self.d)\n    \n    def __getitem__(self, idx):\n        # you can return a dict of these as well etc etc...\n        # remember the order\n        return idx, self.d[idx][\"content_id\"], self.d[idx][\"task_container_id\"], \\\n    self.d[idx][\"part_id\"], self.d[idx][\"prior_question_elapsed_time\"], self.d[idx][\"padded\"], \\\n    self.d[idx][\"answered_correctly\"]\n\ndef collate_fn(batch):\n    _, content_id, task_id, part_id, prior_question_elapsed_time, padded, labels = zip(*batch)\n    content_id = torch.Tensor(content_id).long()\n    task_id = torch.Tensor(task_id).long()\n    part_id = torch.Tensor(part_id).long()\n    prior_question_elapsed_time = torch.Tensor(prior_question_elapsed_time).long()\n    padded = torch.Tensor(padded).bool()\n    labels = torch.Tensor(labels)\n    # remember the order\n    return content_id, task_id, part_id, prior_question_elapsed_time, padded, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = Riiid(d=d)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":false,"trusted":true},"cell_type":"code","source":"print(dataset[0]) # sample dataset","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"next(iter(torch.utils.data.DataLoader(dataset=dataset, batch_size=1, collate_fn=collate_fn, num_workers=8))) # dummy check","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# createing the mdoel\nmodel = TransformerModel(ninp=LAST_N, nhead=4, nhid=128, nlayers=3, dropout=0.3)\nmodel # look into it!","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# dummy stuff\n\nBATCH_SIZE = 32\nlosses = []\ncriterion = nn.BCEWithLogitsLoss()\nlr = 1e-3 # learning rate\noptimizer = torch.optim.Adam(model.parameters(), lr=lr)\nmodel.train();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# here we train, NB I have doubta as to what i am doing, so please take it lightly.\nfor idx,batch in tqdm(enumerate(torch.utils.data.DataLoader(dataset=dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn, num_workers=8))):\n    content_id, _, part_id, prior_question_elapsed_time, mask, labels = batch\n    optimizer.zero_grad()\n    with torch.set_grad_enabled(mode=True):\n        output = model(content_id, part_id, prior_question_elapsed_time, mask)\n        # output is (N,S,2) # i am working on it\n        loss = criterion(output[:,:,1], labels)\n        loss.backward()\n        losses.append(loss.detach().data.numpy())\n        optimizer.step()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.Series(losses).astype(np.float32).plot(kind=\"line\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **** NB I AM NOT SURE ABOUT THIS VALIDATION THING AS OF NOW ****"},{"metadata":{"trusted":true},"cell_type":"code","source":"df_valid = pd.read_csv(TRAIN_PATH, nrows=20_00_000, dtype=TRAIN_DTYPES, \n                       usecols=TRAIN_DTYPES.keys(), skiprows=range(1, 40_00_000)\n                      )\ndf_valid = df_valid[df_valid.content_type_id == 0]\n\ndf_valid[\"prior_question_elapsed_time\"].fillna(26000, inplace=True)\ndf_valid[\"prior_question_elapsed_time\"] = df_valid[\"prior_question_elapsed_time\"] // 1000\n\npart_ids_map = dict(zip(df_questions.question_id, df_questions.part))\ndf_valid['part_id'] = df_valid['content_id'].map(part_ids_map)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import deque\nd = {}\nuser_id_to_idx = {}\nPAD = 0\n\ngrp = df_valid.groupby(\"user_id\").tail(LAST_N)\nfor idx, row in tqdm(grp.groupby(\"user_id\").agg({\n    \"content_id\":list, \"answered_correctly\":list, \"task_container_id\":list, \n    \"part_id\":list, \"prior_question_elapsed_time\":list\n}).reset_index().iterrows()):\n    if len(row[\"content_id\"]) >= 100:\n        d[idx] = {\n            \"user_id\": row[\"user_id\"],\n            \"content_id\" : deque(row[\"content_id\"], maxlen=LAST_N),\n            \"answered_correctly\" : deque(row[\"answered_correctly\"], maxlen=LAST_N),\n            \"task_container_id\" : deque(row[\"task_container_id\"], maxlen=LAST_N),\n            \"prior_question_elapsed_time\" : deque(row[\"prior_question_elapsed_time\"], maxlen=LAST_N),\n            \"part_id\": deque(row[\"part_id\"], maxlen=LAST_N),\n            \"padded\" : deque([False]*100, maxlen=LAST_N)\n        }\n    else:\n        # we have to pad...\n        # (max_batch_len - len(seq))\n        d[idx] = {\n            \"user_id\": row[\"user_id\"],\n            \"content_id\" : deque(row[\"content_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"answered_correctly\" : deque(row[\"answered_correctly\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"task_container_id\" : deque(row[\"task_container_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"prior_question_elapsed_time\" : deque(row[\"prior_question_elapsed_time\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"part_id\": deque(row[\"part_id\"] + [PAD]*(100-len(row[\"content_id\"])), maxlen=LAST_N),\n            \"padded\" : deque([False]*len(row[\"content_id\"]) + [True]*(100-len(row[\"content_id\"])), maxlen=LAST_N)\n        }\n    user_id_to_idx[row[\"user_id\"]] = idx","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def roc_auc_compute_fn(y_targets, y_preds):\n    try:\n        from sklearn.metrics import roc_auc_score\n    except ImportError:\n        raise RuntimeError(\"This contrib module requires sklearn to be installed.\")\n\n    y_true = y_targets.numpy()\n    y_pred = y_preds.numpy()\n    return roc_auc_score(y_true, y_pred)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# **** NB I AM NOT SURE ABOUT THIS VALIDATION THING AS OF NOW ****\n\nfrom sklearn.metrics import accuracy_score\n\nmodel.eval();\ndataset = Riiid(d=d)\nscores = []\n\nfor idx,batch in tqdm(enumerate(torch.utils.data.DataLoader(dataset=dataset, batch_size=32, collate_fn=collate_fn, drop_last=True))):\n    content_id, _, part_id, prior_question_elapsed_time, mask, labels = batch\n    with torch.set_grad_enabled(mode=False):\n            output = model(content_id, part_id, prior_question_elapsed_time, mask)\n            output_prob = output[:,:,1]\n            pred = output_prob >= 0.50\n            # print(output.shape, labels.shape) # torch.Size([N, S, 2]) torch.Size([N, S])\n            _, predicted = torch.max(output[:,:,].data, 1)\n            try:\n                score = roc_auc_compute_fn(labels, pred)\n                scores.append(score)\n            except:\n                print(\"Metric calculation failed for\", idx)\n                continue\nnp.mean(scores)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"That's it for now, stay tuned for updates! Plz let me know my bugs!"}],"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}