{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Preliminaries","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport math\nimport pickle\nimport psutil\nimport random\n\nimport json\nimport numpy as np\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\n\nimport riiideducation","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-08-12T18:28:03.086655Z","iopub.execute_input":"2021-08-12T18:28:03.086964Z","iopub.status.idle":"2021-08-12T18:28:04.301251Z","shell.execute_reply.started":"2021-08-12T18:28:03.086894Z","shell.execute_reply":"2021-08-12T18:28:04.300455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed = 0\nrandom.seed(seed)\ntorch.random.manual_seed(seed)\n\nn_workers = os.cpu_count()\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ncfg_path = '/kaggle/input/riiid-mydata/cfg.json'\ntrain_path = '/kaggle/input/riiid-mydata/train.pkl'\ntag_path = '/kaggle/input/riiid-mydata/tags.csv'\nstates_path = '/kaggle/input/riiid-mydata/states.pickle'\nmodel_path = '/kaggle/input/riiid-mydata/SAKT_08-12_06-24.pt'\n\nB = 512\nMAX_LEN = 256\nD_MODEL = 256","metadata":{"execution":{"iopub.status.busy":"2021-08-12T18:28:04.302692Z","iopub.execute_input":"2021-08-12T18:28:04.303015Z","iopub.status.idle":"2021-08-12T18:28:04.385801Z","shell.execute_reply.started":"2021-08-12T18:28:04.302981Z","shell.execute_reply":"2021-08-12T18:28:04.385016Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## SAKT Model","metadata":{}},{"cell_type":"code","source":"class FFN(nn.Module): \n    def __init__(self, d_model, dropout=0.0): \n        super().__init__()\n        self.lr1 = nn.Linear(d_model, d_model)\n        self.relu = nn.ReLU()\n        self.lr2 = nn.Linear(d_model, d_model)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x): \n        x = self.lr1(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n        x = self.lr2(x)\n        return x\n\n\nclass SAKTModel(nn.Module): \n    def __init__(self, n_exercises, max_len=160, d_model=200, dropout=0.2): \n        super().__init__()\n        self._n_exercises = n_exercises\n        self.position_embd = nn.Embedding(max_len, d_model)\n        self.interaction_embd = nn.Embedding(2 * n_exercises + 1, d_model)\n        self.exercise_embd = nn.Embedding(n_exercises + 1, d_model)\n        self.attn = nn.MultiheadAttention(d_model, num_heads=8, dropout=dropout)\n        self.ffn = FFN(d_model, dropout=dropout)\n        self.predict = nn.Linear(d_model, 1)\n\n        self.norm1 = nn.LayerNorm(d_model)\n        self.norm2 = nn.LayerNorm(d_model)\n        self.dropout = nn.Dropout(dropout)\n\n    @property\n    def n_exercises(self): \n        return self._n_exercises\n\n    def forward(self, x, e, attn_mask=None): \n        seq_len = x.shape[1]\n        device = x.device\n        pos = torch.arange(seq_len, device=device).unsqueeze(0)\n        \n        x = self.interaction_embd(x) + self.position_embd(pos)\n        e = self.exercise_embd(e)\n\n        # (B, L, d) -> (L, B, d)\n        x = x.transpose(0, 1)\n        e = e.transpose(0, 1)\n        x, attn_weights = self.attn(e, x, x, attn_mask=attn_mask)\n        x = self.norm1(x + e)\n\n        # (L, B, d) => (B, L, d)\n        x = x.transpose(0, 1)\n        x = self.norm2(x + self.ffn(x))\n        x = self.predict(x)\n        return x.squeeze(-1), attn_weights","metadata":{"execution":{"iopub.status.busy":"2021-08-12T18:28:04.389012Z","iopub.execute_input":"2021-08-12T18:28:04.389265Z","iopub.status.idle":"2021-08-12T18:28:04.402749Z","shell.execute_reply.started":"2021-08-12T18:28:04.389240Z","shell.execute_reply":"2021-08-12T18:28:04.401961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Testing Phase","metadata":{}},{"cell_type":"code","source":"class KaggleOnlineDataset(Dataset): \n    def __init__(self, train_path, tag_path, states_path, n_exercises, cols, max_len): \n        super().__init__()\n        self.df = pd.read_pickle(train_path)\n        self.test_df = None\n        tag_df = pd.read_csv(tag_path, usecols=['exercise_id', 'bundle_id', 'part', 'correct_rate', 'frequency'])\n        assert np.all(tag_df['exercise_id'].values == np.arange(n_exercises))\n        self.parts, self.correct_rate, self.frequency = tag_df[['part', 'correct_rate', 'frequency']].values.T\n        self.lag_info = pickle.load(open(states_path, 'rb'))\n        self.n_exercises = n_exercises\n        self.cols = cols\n        self.max_len = max_len\n\n    def __len__(self): \n        assert self.test_df is not None, 'Please call update() first'\n        return len(self.test_df)\n\n    def __getitem__(self, idx): \n        new_observation = self.test_df.iloc[idx]\n        # 'correct' is set to 0 temporarily\n        user_id = new_observation['user_id']\n        new_data = {col: np.array([new_observation.get(col, 0)]) for col in self.cols}\n        # retrieve old observations\n        if user_id in self.df.index: \n            old_items = self.df[user_id]\n            old_len = min(len(old_items[0]), self.max_len - 1)\n            data = {key: np.append(old_item[-old_len:], new_data[key]) for key, old_item in zip(self.cols, old_items)}\n        else: \n            old_len = 0\n            data = new_data\n        seq_len = old_len + 1\n        # retrieve addtional features\n        data['part'] = self.parts[data['exercise_id']]\n        data['correct_rate'] = self.correct_rate[data['exercise_id']]\n        # pad to max_len and set dtype\n        dtype_map = {key: int for key in self.cols + ['part']}\n        dtype_map['correct_rate'] = float\n        data = KaggleOnlineDataset._postpad_and_asdtype(data, self.max_len - seq_len, dtype_map)\n        data['valid_len'] = np.array([seq_len], dtype=int)\n        return data\n    \n    @staticmethod\n    def _postpad_and_asdtype(data, pad, dtype_map): \n        return {\n            key: np.pad(item, [[0, pad]]).astype(dtype_map[key]) for key, item in data.items()\n        }\n    \n    def update(self, test_df): \n        if self.test_df is not None and psutil.virtual_memory().percent < 90: \n            # update df according to previous labels\n            prev_df = self.test_df\n            prev_df['correct'] = np.array(eval(test_df.iloc[0]['prior_group_answers_correct']))[self.was_exercise]\n            user_df = prev_df.groupby('user_id').apply(lambda udf: tuple(udf[col].values for col in self.cols))\n            for user_id, new_items in user_df.iteritems(): \n                if user_id in self.df.index: \n                    self.df[user_id] = tuple(map(\n                        lambda old_item, new_item: np.append(old_item, new_item)[-min(self.max_len, len(old_item) + 1):], \n                        self.df[user_id], \n                        new_items\n                    )) # truncate at max_len to prevent OOM\n                else: \n                    self.df[user_id] = tuple(new_item for new_item in new_items) # create a new row\n            # update correct rate\n            # self._update_correct_rate(prev_df)\n        # process test_df\n        is_exercise = (test_df['content_type_id'] == 0)\n        test_df = test_df[is_exercise]\n        test_df = test_df.rename(columns={'content_id': 'exercise_id', 'prior_question_elapsed_time': 'prior_elapsed'})\n        # compute lag and convert ms -> min\n        test_df['prior_elapsed'] = test_df['prior_elapsed'].fillna(0).astype(int)\n        lag = self._compute_new_lag(test_df)\n        test_df['lag'] = np.where(\n            np.logical_and(0 < lag, lag < 60 * 1000), 1, np.round(lag / (1000 * 60))\n        ).astype(int)\n        # as for prior_elapsed, convert ms -> s\n        prior_elapsed = test_df['prior_elapsed'].values\n        test_df['prior_elapsed'] = np.where(\n            np.logical_and(0 < prior_elapsed, prior_elapsed < 1000), 1, np.round(prior_elapsed / 1000)\n        ).astype(int)\n        test_df.reset_index(drop=True, inplace=True)\n        # save relevent information\n        self.test_df = test_df\n        self.was_exercise = is_exercise.values\n        return test_df\n        \n    def _update_correct_rate(self, prev_df): \n        exercise_df = prev_df.groupby('exercise_id').aggregate({'correct': [sum, len]})['correct']\n        n_correct = np.arange(self.n_exercises)\n        np.put(n_correct, exercise_df.index, exercise_df['sum'].values)\n        n_correct = n_correct + np.round(self.correct_rate * self.frequency)\n        more_frequency = np.arange(self.n_exercises)\n        np.put(more_frequency, exercise_df.index, exercise_df['len'].values)\n        self.frequency += more_frequency\n        correct_rate = n_correct / self.frequency\n        self.correct_rate = np.where(np.isfinite(correct_rate), correct_rate, 0.5)\n        \n    def _compute_new_lag(self, df): \n        last_states, exercise_id_to_bundle, bundle_id_to_size = self.lag_info\n        # compute_lag from the original implementation\n        lag = np.zeros(len(df))\n        for i, (user_id, curr_timestamp, curr_exercise_id, prior_elapsed) in enumerate(\n            df[['user_id', 'timestamp', 'exercise_id', 'prior_elapsed']].values\n        ): \n            curr_bundle_id = exercise_id_to_bundle[curr_exercise_id]\n            last_state = last_states.get(user_id, None)\n            if last_state is None: \n                last_states[user_id] = (curr_timestamp, curr_bundle_id)\n                lag[i] = 0\n            else: \n                last_timestamp, last_bundle_id = last_state\n                if curr_bundle_id == last_bundle_id: \n                    # same bundle, do not update last_states\n                    lag[i] = 0\n                else: \n                    last_states[user_id] = (curr_timestamp, curr_bundle_id)\n                    elapsed_offset = bundle_id_to_size[last_bundle_id] * prior_elapsed\n                    lag[i] = curr_timestamp - last_timestamp - elapsed_offset\n        lag = np.clip(lag, a_min=0, a_max=None)\n        self.lag_info = (last_states, exercise_id_to_bundle, bundle_id_to_size)\n        return lag","metadata":{"execution":{"iopub.status.busy":"2021-08-12T18:28:04.404041Z","iopub.execute_input":"2021-08-12T18:28:04.404715Z","iopub.status.idle":"2021-08-12T18:28:04.433772Z","shell.execute_reply.started":"2021-08-12T18:28:04.404599Z","shell.execute_reply":"2021-08-12T18:28:04.432907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def truncate_and_prepare_masks(items, valid_len, need_pad_mask=True, need_attn_mask=True): \n    max_len = valid_len.max()\n    device = max_len.device\n    # truncate at the max_len for each sample\n    out = [None if item is None else item[:, :max_len] for item in items]\n    # pad to the same length for batch-ification\n    pad_mask = torch.arange(max_len, device=device) >= valid_len if need_pad_mask else None\n    # assume q_len = k_len in attention\n    attn_mask = torch.triu(torch.ones(max_len, max_len), diagonal=1).to(device, torch.bool) if need_attn_mask else None\n    return out, pad_mask, attn_mask","metadata":{"execution":{"iopub.status.busy":"2021-08-12T18:28:04.435128Z","iopub.execute_input":"2021-08-12T18:28:04.435539Z","iopub.status.idle":"2021-08-12T18:28:04.445340Z","shell.execute_reply.started":"2021-08-12T18:28:04.435498Z","shell.execute_reply":"2021-08-12T18:28:04.444400Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = json.load(open(cfg_path, 'r'))\nmodel = SAKTModel(cfg['n_exercises'], max_len=MAX_LEN, d_model=D_MODEL).to(device)\nmodel.load_state_dict(torch.load(model_path))\nmodel.eval()\ntestset = KaggleOnlineDataset(train_path, tag_path, states_path, cfg['n_exercises'], cfg['cols'], MAX_LEN)\n\nenv = riiideducation.make_env()\niter_test = env.iter_test()\nfor test_df, _ in iter_test: \n    test_df = testset.update(test_df)\n    testloader = DataLoader(testset, batch_size=B, shuffle=False, num_workers=n_workers, drop_last=False)\n    outs = np.array([], dtype='float32')\n    for data in testloader: \n        valid_len = data['valid_len'].to(device, torch.long)\n        interactions = F.pad(data['exercise_id'][:, :-1], (1, 0)) + F.pad(data['correct'][:, :-1], (1, 0)) * cfg['n_exercises']\n        inputs, _, attn_mask = truncate_and_prepare_masks(\n            [\n                interactions.to(device, torch.long), \n                data['exercise_id'].to(device, torch.long)\n            ], \n            valid_len, \n            need_pad_mask=False\n        )\n        out, _ = model(*inputs, attn_mask=attn_mask)\n        out = torch.gather(out, 1, valid_len - 1).squeeze(-1)\n        outs = np.append(outs, torch.sigmoid(out).detach().cpu().numpy())\n    test_df['answered_correctly'] = outs\n    env.predict(test_df[['row_id', 'answered_correctly']])","metadata":{"execution":{"iopub.status.busy":"2021-08-12T18:28:04.446346Z","iopub.execute_input":"2021-08-12T18:28:04.446875Z","iopub.status.idle":"2021-08-12T18:28:29.760024Z","shell.execute_reply.started":"2021-08-12T18:28:04.446838Z","shell.execute_reply":"2021-08-12T18:28:29.758981Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}