{"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-13T15:38:07.489292Z","iopub.execute_input":"2021-08-13T15:38:07.489893Z","iopub.status.idle":"2021-08-13T15:38:07.959423Z","shell.execute_reply.started":"2021-08-13T15:38:07.489817Z","shell.execute_reply":"2021-08-13T15:38:07.958560Z"},"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/DKVMN_08-13_21-47.pt'\n\nB = 512\nMAX_LEN = 128\nSEQ_LEN = 128\nN_SIZE = 64\nD_MODEL = 256","metadata":{"execution":{"iopub.status.busy":"2021-08-13T15:38:07.966262Z","iopub.execute_input":"2021-08-13T15:38:07.966602Z","iopub.status.idle":"2021-08-13T15:38:07.993288Z","shell.execute_reply.started":"2021-08-13T15:38:07.966567Z","shell.execute_reply":"2021-08-13T15:38:07.992456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## DKVMN Model","metadata":{}},{"cell_type":"code","source":"class DKVMNModel(nn.Module): \n    def __init__(self, n_exercises, n_size, d_model): \n        super().__init__()\n        self._n_exercises = n_exercises\n        self.mem_key = nn.Parameter(torch.randn(n_size, d_model))\n        self.mem_value = nn.Parameter(torch.randn(n_size, d_model))\n        self.exercise_embd = nn.Embedding(n_exercises, d_model)\n        self.interaction_embd = nn.Embedding(2 * n_exercises, d_model)\n        self.E = nn.Linear(d_model, d_model)\n        self.D = nn.Linear(d_model, d_model)\n        self.ffn = nn.Linear(2 * d_model, d_model)\n        self.predict = nn.Linear(d_model, 1)\n\n    @property\n    def n_exercises(self): \n        return self._n_exercises\n\n    def forward(self, e, x): \n        B, seq_len = e.shape\n        # embedding\n        query = self.exercise_embd(e)\n        response = self.interaction_embd(x)\n        # attention\n        attn_scores = torch.einsum('bld,nd->bln', (query, self.mem_key))\n        attn_weights = torch.softmax(attn_scores, dim=-1)\n        # iterate over the input sequence\n        value = self.mem_value.unsqueeze(0).expand(B, *self.mem_value.shape)\n        read = torch.empty_like(query)\n        for i in range(seq_len): \n            attn_weights_i = attn_weights[:, i, :]\n            # read\n            read[:, i, :] = torch.einsum('bn,bnd->bd', (attn_weights_i, value))\n            # write\n            erase = torch.sigmoid(self.E(response[:, i, :]))\n            add = torch.tanh(self.D(response[:, i, :]))\n            value = value * (1 - torch.einsum('bn,bd->bnd', (attn_weights_i, erase))) + \\\n                torch.einsum('bn,bd->bnd', (attn_weights_i, add))\n        # prediction\n        summary = self.ffn(torch.cat([read, query], dim=-1))\n        summary = torch.tanh(summary)\n        logit = self.predict(summary)\n        return logit.squeeze(-1), value","metadata":{"execution":{"iopub.status.busy":"2021-08-13T15:38:07.994850Z","iopub.execute_input":"2021-08-13T15:38:07.995160Z","iopub.status.idle":"2021-08-13T15:38:08.008909Z","shell.execute_reply.started":"2021-08-13T15:38:07.995132Z","shell.execute_reply":"2021-08-13T15:38:08.007699Z"},"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-13T15:38:08.010544Z","iopub.execute_input":"2021-08-13T15:38:08.010981Z","iopub.status.idle":"2021-08-13T15:38:08.041226Z","shell.execute_reply.started":"2021-08-13T15:38:08.010944Z","shell.execute_reply":"2021-08-13T15:38:08.040228Z"},"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-13T15:38:08.042616Z","iopub.execute_input":"2021-08-13T15:38:08.042963Z","iopub.status.idle":"2021-08-13T15:38:08.054192Z","shell.execute_reply.started":"2021-08-13T15:38:08.042927Z","shell.execute_reply":"2021-08-13T15:38:08.053282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = json.load(open(cfg_path, 'r'))\nmodel = DKVMNModel(cfg['n_exercises'], N_SIZE, 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 = data['exercise_id'] + data['correct'] * model.n_exercises\n        (*inputs, labels), pad_mask, _ = truncate_and_prepare_masks(\n            [\n                data['exercise_id'].to(device, torch.long), \n                interactions.to(device, torch.long), \n                data['correct'].to(device, torch.float)\n            ], \n            valid_len, \n            need_attn_mask=False\n        )\n        max_len = valid_len.max().item()\n        mem_x = None\n        out = torch.empty_like(labels)\n        for i, start in enumerate(range(0, max_len, SEQ_LEN)): \n            end = min(start + SEQ_LEN, max_len)\n            inputs_i = tuple(item[:, start:end] for item in inputs)\n            out_i, _ = model(*inputs_i)\n            out[:, start:end] = out_i.detach()\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-13T15:38:08.055489Z","iopub.execute_input":"2021-08-13T15:38:08.055943Z","iopub.status.idle":"2021-08-13T15:38:31.825566Z","shell.execute_reply.started":"2021-08-13T15:38:08.055904Z","shell.execute_reply":"2021-08-13T15:38:31.824545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}