submission_kernel_no_lec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submission Kernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code needs three kaggle datasets:\n",
    "- riiid-nn-data (contains npy files with parts and tags)\n",
    "- nn-data-no-lec (contains the .h5 file with all data organized by user/dataset)\n",
    "- riiid-nn-model (contains the model.ckpt file to load)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "from pytorch_lightning import seed_everything\n",
    "from pytorch_lightning.metrics.functional import accuracy\n",
    "from torch import nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, Dataset, random_split\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "pd.set_option(\"display.max_rows\", 100)\n",
    "\n",
    "\n",
    "SEED = 69\n",
    "seed_everything(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading tags and lectures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "folder_path = \"../input/riiid-nn-data\"\n",
    "\n",
    "questions_df = pd.read_csv(\"../input/riiid-test-answer-prediction/questions.csv\")\n",
    "lectures_df = pd.read_csv(\"../input/riiid-test-answer-prediction/lectures.csv\")\n",
    "\n",
    "print(\"Loading lectures arrays\")\n",
    "lectures_ids = np.load(f\"{folder_path}/lecture_ids.npy\")\n",
    "lectures_parts = np.load(f\"{folder_path}/lectures_parts.npy\")\n",
    "lectures_tags = lectures_df.tag.values\n",
    "\n",
    "print(\"Loading questions arrays\")\n",
    "questions_parts = np.load(f\"{folder_path}/questions_parts.npy\")\n",
    "questions_lectures_parts = np.concatenate([questions_parts, lectures_parts])\n",
    "\n",
    "\n",
    "# process tags\n",
    "def split_tags(t):\n",
    "    try:\n",
    "        return [int(i) for i in t.split(\" \")]\n",
    "    except AttributeError:\n",
    "        return list()\n",
    "\n",
    "\n",
    "# Get tags to be 2D array of shape (Q, T), where Q is question_idx, and T is the max number of tag possible (6)\n",
    "questions_df[\"tags\"] = questions_df.tags.apply(split_tags)\n",
    "questions_tags = pd.DataFrame(questions_df[\"tags\"].tolist(), index=questions_df.index)\n",
    "\n",
    "# map lecture id to new id\n",
    "lectures_mapping = dict(\n",
    "    zip(lectures_df.lecture_id.values, (lectures_df.index + 13523).values)\n",
    ")\n",
    "lectures_df.lecture_id = lectures_df.index + 13523\n",
    "lectures_tags = pd.DataFrame(\n",
    "    lectures_df.tag.values, index=lectures_df.lecture_id.values\n",
    ")\n",
    "\n",
    "questions_lectures_tags = pd.concat([questions_tags, lectures_tags])\n",
    "# pad with max tag + 1\n",
    "questions_lectures_tags = (\n",
    "    questions_lectures_tags.fillna(questions_lectures_tags.max().max() + 1)\n",
    "    .astype(np.int)\n",
    "    .values\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preproc code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_df(df):\n",
    "    \"\"\"\n",
    "    Converts the lecture ids to proper content_ids\n",
    "    Adds the answered_correctly column if not exists\n",
    "    \"\"\"\n",
    "    df.content_type_id = df.content_type_id.astype(bool)\n",
    "\n",
    "    # prior information\n",
    "    df.prior_question_had_explanation = df.prior_question_had_explanation.fillna(\n",
    "        0\n",
    "    ).astype(np.uint8)\n",
    "    df.prior_question_elapsed_time = (\n",
    "        df.prior_question_elapsed_time.fillna(0).clip(upper=300000) / 300000\n",
    "    )  # normalizes to 0-1\n",
    "\n",
    "    # map lecture ids to new content_ids\n",
    "    df.loc[df.content_type_id, \"content_id\"] = df[df.content_type_id].content_id.map(\n",
    "        lectures_mapping\n",
    "    )\n",
    "    # if not answered correctly then add column with\n",
    "    # y = 3 (padding) for all questions and y = 4 for lectures\n",
    "    if \"answered_correctly\" not in df.columns:\n",
    "        df[\"answered_correctly\"] = df.content_type_id.map({False: 3, True: 4})\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "eps = 0.0000001\n",
    "\n",
    "\n",
    "def get_time_elapsed_from_timestamp(arr):\n",
    "    arr_seconds = np.diff(arr, prepend=0) / 1000\n",
    "    return (np.log(arr_seconds + eps).astype(np.float32) - 3.5) / 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytorch Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning.core.decorators import auto_move_data\n",
    "from pytorch_lightning.metrics.functional.classification import auroc\n",
    "\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "        torch.nn.init.zeros_(m.bias)\n",
    "\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=1000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(\n",
    "            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)\n",
    "        )\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
    "        self.register_buffer(\"pe\", pe)\n",
    "\n",
    "    def forward(self, sequence_length):\n",
    "        # returns embeds (sequence_length, 1, d_model)\n",
    "        return self.pe[:sequence_length, :]\n",
    "\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "        torch.nn.init.zeros_(m.bias)\n",
    "\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=1000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(\n",
    "            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)\n",
    "        )\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
    "        self.register_buffer(\"pe\", pe)\n",
    "\n",
    "    def forward(self, sequence_length):\n",
    "        # returns embeds (sequence_length, 1, d_model)\n",
    "        return self.pe[:sequence_length, :]\n",
    "\n",
    "\n",
    "class RIIDDTransformerModel(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        learning_rate=0.001,\n",
    "        n_content_id=13943,  # number of different contents = 13942 + 1 (for padding)\n",
    "        n_part=8,  # number of different parts = 7 + 1 (for padding)\n",
    "        n_tags=189,  # number of different tags = 188 + 1 (for padding)\n",
    "        n_correct=5,  # 0,1 (false, true), 2 (start token), 3 (padding), 4 (lecture)\n",
    "        emb_dim=64,  # embedding dimension\n",
    "        dropout=0.1,\n",
    "        n_heads: int = 1,\n",
    "        n_encoder_layers: int = 2,\n",
    "        n_decoder_layers: int = 2,\n",
    "        dim_feedforward: int = 256,\n",
    "        activation: str = \"relu\",\n",
    "        max_window_size=100,\n",
    "        use_prior_q_times=False,\n",
    "        use_prior_q_explanation=False,\n",
    "    ):\n",
    "        super(RIIDDTransformerModel, self).__init__()\n",
    "        self.model_type = \"RiiidTransformer\"\n",
    "        self.learning_rate = learning_rate\n",
    "        self.max_window_size = max_window_size\n",
    "\n",
    "        self.use_prior_q_times = use_prior_q_times\n",
    "        self.use_prior_q_explanation = use_prior_q_explanation\n",
    "\n",
    "        # save params of models to yml\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        #### EXERCISE SEQUENCE\n",
    "        self.embed_content_id = nn.Embedding(n_content_id, emb_dim, padding_idx=13942)\n",
    "        self.embed_parts = nn.Embedding(n_part, emb_dim, padding_idx=0)\n",
    "        self.embed_tags = nn.Embedding(n_tags, emb_dim, padding_idx=188)\n",
    "\n",
    "        # exercise weights to weight the mean embeded excercise embeddings\n",
    "        self.exercise_weights = torch.nn.Parameter(torch.tensor([0.35, 0.55, 0.1]))\n",
    "\n",
    "        ### RESPONSE SEQUENCE (1st time stamp of sequence is useless)\n",
    "        self.embed_answered_correctly = nn.Embedding(\n",
    "            n_correct, emb_dim, padding_idx=3\n",
    "        )  # 2 + 1 for start token + 1 for padding_idn_inputs\n",
    "\n",
    "        self.embed_timestamps = nn.Linear(1, emb_dim)\n",
    "\n",
    "        # embed prior q time and q explanation\n",
    "        self.embed_prior_q_time = nn.Linear(1, emb_dim)\n",
    "        self.embed_prior_q_explanation = nn.Embedding(2, emb_dim)\n",
    "\n",
    "        # response weights to weight the mean embeded response embeddings\n",
    "        w = [0.5, 0.5]\n",
    "        if use_prior_q_times:\n",
    "            w.append(0.5)\n",
    "        if use_prior_q_explanation:\n",
    "            w.append(0.5)\n",
    "\n",
    "        self.response_weights = torch.nn.Parameter(torch.tensor([w]))\n",
    "\n",
    "        self.pos_encoder = PositionalEncoding(emb_dim)\n",
    "\n",
    "        self.transformer = nn.Transformer(\n",
    "            d_model=emb_dim,\n",
    "            nhead=n_heads,\n",
    "            num_encoder_layers=n_encoder_layers,\n",
    "            num_decoder_layers=n_decoder_layers,\n",
    "            dropout=dropout,\n",
    "            dim_feedforward=dim_feedforward,\n",
    "            activation=activation,\n",
    "        )\n",
    "\n",
    "        self.out_linear = nn.Linear(emb_dim, 2)\n",
    "        init_weights(self)\n",
    "\n",
    "    def generate_square_subsequent_mask(self, sz):\n",
    "        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n",
    "        mask = (\n",
    "            mask.float()\n",
    "            .masked_fill(mask == 0, float(\"-inf\"))\n",
    "            .masked_fill(mask == 1, float(0.0))\n",
    "        )\n",
    "        return mask\n",
    "\n",
    "    def forward(\n",
    "        self, content_ids, parts, answers, tags, timestamps, prior_q_times,\n",
    "    ):\n",
    "        # content_ids: (Source Sequence Length, Number of samples, Embedding)\n",
    "        # tgt: (Target Sequence Length,Number of samples, Embedding)\n",
    "\n",
    "        # if data is flat then expand to get Batch dim\n",
    "        if len(content_ids.shape) == 1:\n",
    "            content_ids = content_ids.unsqueeze(1)\n",
    "            parts = parts.unsqueeze(1)\n",
    "            answers = answers.unsqueeze(1)\n",
    "            tags = tags.unsqueeze(1)\n",
    "            timestamps = timestamps.unsqueeze(1)\n",
    "            prior_q_times = prior_q_times.unsqueeze(1)\n",
    "\n",
    "        sequence_length = content_ids.shape[0]\n",
    "\n",
    "        # sequence that will go into encoder\n",
    "        embeded_content = self.embed_content_id(content_ids)\n",
    "        embeded_parts = self.embed_parts(parts)\n",
    "        embeded_tags = self.embed_tags(tags).sum(dim=2)\n",
    "        e_w = F.softmax(self.exercise_weights, dim=0)\n",
    "\n",
    "        embeded_exercise_sequence = (\n",
    "            torch.stack([embeded_content, embeded_parts, embeded_tags], dim=3) * e_w\n",
    "        ).sum(dim=3)\n",
    "\n",
    "        # sequence that will go into decoder\n",
    "        embeded_answered_correctly = self.embed_answered_correctly(answers)\n",
    "        embeded_timestamps = self.embed_timestamps(timestamps.unsqueeze(2))\n",
    "        r_w = F.softmax(self.response_weights, dim=0)\n",
    "\n",
    "        exercise_sequence_components = [embeded_answered_correctly, embeded_timestamps]\n",
    "        if self.use_prior_q_times:\n",
    "            embeded_q_times = self.embed_prior_q_time(prior_q_times.unsqueeze(2))\n",
    "            # zero embedding - if start token\n",
    "            embeded_q_times[0, torch.where(answers[0, :] == 2)[0], :] = 0\n",
    "            exercise_sequence_components.append(embeded_q_times)\n",
    "\n",
    "        embeded_responses = (\n",
    "            torch.stack(exercise_sequence_components, dim=3) * r_w\n",
    "        ).sum(dim=3)\n",
    "\n",
    "        # adding positional vector\n",
    "        embedded_positions = self.pos_encoder(sequence_length + 1)\n",
    "\n",
    "        # add shifted position embedding ( start token is first position)\n",
    "        embeded_responses = embeded_responses + embedded_positions[:-1, :, :]\n",
    "\n",
    "        embeded_exercise_sequence = (\n",
    "            embeded_exercise_sequence + embedded_positions[1:, :, :]\n",
    "        )\n",
    "\n",
    "        # mask of shape S x S -> prevents attention looking forward\n",
    "        top_right_attention_mask = self.generate_square_subsequent_mask(\n",
    "            sequence_length\n",
    "        ).type_as(embeded_exercise_sequence)\n",
    "\n",
    "        output = self.transformer(\n",
    "            embeded_exercise_sequence,\n",
    "            embeded_responses,\n",
    "            tgt_mask=top_right_attention_mask,  # (S,S)\n",
    "            src_mask=top_right_attention_mask,  # (T,T)\n",
    "        )\n",
    "\n",
    "        output = self.out_linear(output)\n",
    "        return F.softmax(output, dim=2)[:, :, 1]\n",
    "\n",
    "    def process_batch_step(self, batch):\n",
    "        # return result\n",
    "        return self(\n",
    "            batch[\"content_ids\"],\n",
    "            batch[\"parts\"],\n",
    "            batch[\"answers\"],\n",
    "            batch[\"tags\"],\n",
    "            batch[\"timestamps\"],\n",
    "            batch[\"prior_q_times\"],\n",
    "        )\n",
    "\n",
    "    @auto_move_data\n",
    "    def predict_n_steps(self, batch, steps, return_all_preds=False):\n",
    "        \"\"\"\n",
    "        Predicts n steps for all items in batch and return predictions\n",
    "        only for those steps (flattened)\n",
    "        steps: tensor of length B where each item is the number of steps that need to be taken\n",
    "        \"\"\"\n",
    "        seq_length, n_users = batch[\"content_ids\"].shape\n",
    "        lengths = batch[\"length\"]\n",
    "\n",
    "        users = torch.arange(n_users)\n",
    "\n",
    "        user_indexes = []\n",
    "        sequence_indexes = []\n",
    "\n",
    "        for i in range(steps.max().int(), 0, -1):\n",
    "            preds = self.process_batch_step(batch)\n",
    "\n",
    "            sequence_indexes_at_i = lengths[steps >= i] - i\n",
    "            user_indexes_at_i = users[steps >= i]\n",
    "\n",
    "            # get index for which to update the answers\n",
    "            # since answers is shifted we want to map preds 0..98 -> answers 1:99\n",
    "            answers_idx = torch.where(sequence_indexes_at_i + 1 != seq_length)\n",
    "            a_seq_idx = sequence_indexes_at_i[answers_idx] + 1\n",
    "            u_seq_idx = user_indexes_at_i[answers_idx]\n",
    "\n",
    "            # set answer to either 0 or 1 if not lecture\n",
    "            batch[\"answers\"][a_seq_idx, u_seq_idx] = torch.where(\n",
    "                batch[\"answers\"][a_seq_idx, u_seq_idx] != 4,\n",
    "                (preds[sequence_indexes_at_i[answers_idx], u_seq_idx] > 0.5).long(),\n",
    "                batch[\"answers\"][a_seq_idx, u_seq_idx],\n",
    "            )\n",
    "\n",
    "            user_indexes.append(user_indexes_at_i)\n",
    "            sequence_indexes.append(sequence_indexes_at_i)\n",
    "\n",
    "        if return_all_preds:\n",
    "            return preds\n",
    "\n",
    "        user_indexes = torch.cat(user_indexes)\n",
    "        sequence_indexes = torch.cat(sequence_indexes)\n",
    "\n",
    "        return (\n",
    "            preds[sequence_indexes, user_indexes],\n",
    "            batch[\"row_ids\"][sequence_indexes, user_indexes],\n",
    "        )\n",
    "\n",
    "    @auto_move_data\n",
    "    def predict_fast_single_user(\n",
    "        self, content_ids, parts, answers, tags, timestamps, prior_q_times, n=1\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Predicts n steps for a single user in batch and return predictions\n",
    "        only for those steps (flattened)\n",
    "        \"\"\"\n",
    "        length = len(content_ids)\n",
    "        out_predictions = torch.zeros(n, device=self.device)\n",
    "        for i in range(n, 0, -1):\n",
    "            preds = self(content_ids, parts, answers, tags, timestamps, prior_q_times)\n",
    "            out_predictions[n - i] = preds[length - i, 0]\n",
    "\n",
    "            # answers are shifted (start token) so need + 1\n",
    "            answer_idx = length - i + 1\n",
    "            # don't update if at end of answers\n",
    "            if answer_idx + 1 < len(answers):\n",
    "                # don't update if true is lecture\n",
    "                if answers[answer_idx] != 4:\n",
    "                    answers[answer_idx] = (preds[length - i, 0] > 0.5).long()\n",
    "\n",
    "        return out_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RIIDDTransformerModel.load_from_checkpoint(\"../input/riiid-nn-model/model.ckpt\")\n",
    "model.freeze()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytorch Dataset and Collate Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InferenceRIIDDataset(Dataset):\n",
    "    \"\"\"RIID dataset.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, hdf5_file=\"feats.h5\", window_size=100,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            user_mapping (np.array): array of all unique user ids \n",
    "            hdf5_file (string): location of hf5 feats file\n",
    "        \"\"\"\n",
    "        self.hdf5_file = hdf5_file\n",
    "        self.max_window_size = window_size\n",
    "        self.cache = {}\n",
    "\n",
    "    def open_hdf5(self):\n",
    "        # opens the h5py file\n",
    "        self.f = h5py.File(self.hdf5_file, \"r\")\n",
    "\n",
    "    def __len__(self):\n",
    "        if not hasattr(self, \"f\"):\n",
    "            self.open_hdf5()\n",
    "        return len(self.f.keys())\n",
    "\n",
    "    def load_user_into_cache(self, user_id):\n",
    "        \"\"\"\n",
    "        add a user to self.cache\n",
    "        \"\"\"\n",
    "\n",
    "        if not hasattr(self, \"f\"):\n",
    "            self.open_hdf5()\n",
    "\n",
    "        if f\"{user_id}\" in self.f:\n",
    "            content_ids = np.array(self.f[f\"{user_id}/content_ids\"], dtype=\"int64\")\n",
    "            answers = np.array(self.f[f\"{user_id}/answered_correctly\"], dtype=\"int64\")\n",
    "            timestamps = np.array(self.f[f\"{user_id}/timestamps\"], dtype=\"float32\")\n",
    "            prior_q_times = np.array(\n",
    "                self.f[f\"{user_id}/prior_question_elapsed_time\"], dtype=\"float32\"\n",
    "            )\n",
    "\n",
    "            self.cache[user_id] = {\n",
    "                \"content_ids\": content_ids,\n",
    "                \"answers\": answers,\n",
    "                \"timestamps\": timestamps,\n",
    "                \"prior_q_times\": prior_q_times,\n",
    "                \"row_ids\": np.zeros(len(timestamps)),\n",
    "                \"steps\": 0,\n",
    "            }\n",
    "        else:\n",
    "            self.cache[user_id] = {\n",
    "                \"content_ids\": np.array([], dtype=\"int64\"),\n",
    "                \"timestamps\": np.array([], dtype=\"float32\"),\n",
    "                \"answers\": np.array([], dtype=\"int64\"),\n",
    "                \"prior_q_times\": np.array([], dtype=\"float32\"),\n",
    "                \"row_ids\": np.array([], dtype=\"int64\"),\n",
    "                \"steps\": 0,\n",
    "            }\n",
    "\n",
    "    def update_user_rows(self, user_rows):\n",
    "        if not hasattr(self, \"f\"):\n",
    "            self.open_hdf5()\n",
    "\n",
    "        user_id = user_rows.user_id.values[0]\n",
    "        num_rows = len(user_rows)\n",
    "        new_content_ids = user_rows.content_id.values\n",
    "        new_timestamps = user_rows.timestamp.values\n",
    "        new_answered_correctly = (\n",
    "            user_rows.answered_correctly.values\n",
    "        )  # should be 3 (3) for every question and 4 for lectures\n",
    "        new_prior_q_times = user_rows.prior_question_elapsed_time.values\n",
    "\n",
    "        new_row_ids = user_rows.row_id.values\n",
    "\n",
    "        if user_id not in self.cache:\n",
    "            self.load_user_into_cache(user_id)\n",
    "\n",
    "        self.cache[user_id][\"content_ids\"] = np.concatenate(\n",
    "            [self.cache[user_id][\"content_ids\"], new_content_ids]\n",
    "        )[-(self.max_window_size + 1) :]\n",
    "\n",
    "        self.cache[user_id][\"timestamps\"] = np.concatenate(\n",
    "            [self.cache[user_id][\"timestamps\"], new_timestamps]\n",
    "        )[-(self.max_window_size + 1) :]\n",
    "\n",
    "        self.cache[user_id][\"answers\"] = np.concatenate(\n",
    "            [self.cache[user_id][\"answers\"], new_answered_correctly]\n",
    "        )[-(self.max_window_size + 1) :]\n",
    "\n",
    "        self.cache[user_id][\"prior_q_times\"] = np.concatenate(\n",
    "            [self.cache[user_id][\"prior_q_times\"], new_prior_q_times]\n",
    "        )[-(self.max_window_size + 1) :]\n",
    "\n",
    "        self.cache[user_id][\"row_ids\"] = np.concatenate(\n",
    "            [self.cache[user_id][\"row_ids\"], new_row_ids]\n",
    "        )[-(self.max_window_size + 1) :]\n",
    "\n",
    "        self.cache[user_id][\"steps\"] = len(new_row_ids)\n",
    "\n",
    "    def update_answered_correctly(self, answered_correctly_rows):\n",
    "        user_id = answered_correctly_rows.name\n",
    "        new_answered_correctly = (\n",
    "            answered_correctly_rows.values\n",
    "        )  # this is only the answers\n",
    "\n",
    "        # update the correct answers for the questions (keep 4 - lectures)\n",
    "        n = self.cache[user_id][\"steps\"]\n",
    "        l = len(self.cache[user_id][\"answers\"])\n",
    "        idxs = np.where(self.cache[user_id][\"answers\"][-n:] != 4)[0] + l - n\n",
    "\n",
    "        self.cache[user_id][\"answers\"][idxs] = new_answered_correctly\n",
    "        self.cache[user_id][\"steps\"] = 0\n",
    "\n",
    "    def __getitem__(self, user_id):\n",
    "        if user_id not in self.cache:\n",
    "            self.load_user_into_cache(user_id)\n",
    "\n",
    "        length = len(self.cache[user_id][\"content_ids\"])\n",
    "        window_size = min(self.max_window_size, length)\n",
    "\n",
    "        content_ids = self.cache[user_id][\"content_ids\"][-window_size:]\n",
    "        answers = self.cache[user_id][\"answers\"][-window_size:]\n",
    "        timestamps = self.cache[user_id][\"timestamps\"][-window_size:]\n",
    "        prior_q_times = self.cache[user_id][\"prior_q_times\"][-window_size:]\n",
    "        row_ids = self.cache[user_id][\"row_ids\"][-window_size:]\n",
    "\n",
    "        # convert timestamps to time elapsed\n",
    "        time_elapsed_timestamps = get_time_elapsed_from_timestamp(timestamps)\n",
    "\n",
    "        # get tags\n",
    "        tags = questions_lectures_tags[content_ids, :].astype(np.int64)\n",
    "\n",
    "        # get parts\n",
    "        parts = questions_lectures_parts[content_ids].astype(np.int64)\n",
    "\n",
    "        # shift by one the answered_correctly sequence\n",
    "        answers = np.roll(answers, 1)\n",
    "\n",
    "        if length > self.max_window_size:\n",
    "            answers[0] = self.cache[user_id][\"answers\"][-window_size - 1]\n",
    "        else:\n",
    "            answers[0] = 2\n",
    "\n",
    "        return {\n",
    "            \"parts\": torch.from_numpy(parts).long(),\n",
    "            \"tags\": torch.from_numpy(tags).long(),\n",
    "            \"content_ids\": torch.from_numpy(content_ids).long(),\n",
    "            \"answers\": torch.from_numpy(answers).long(),\n",
    "            \"prior_q_times\": torch.from_numpy(prior_q_times).float(),\n",
    "            \"timestamps\": torch.from_numpy(time_elapsed_timestamps).float(),\n",
    "            \"length\": window_size,\n",
    "            \"row_ids\": torch.from_numpy(row_ids).int(),\n",
    "            \"steps\": self.cache[user_id][\"steps\"],\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "\n",
    "# The collate function is used to merge individual data samples into a batch\n",
    "# It handles the padding aspect\n",
    "def inference_collate_fn(batch):\n",
    "    # collate lenghts into 1D tensor\n",
    "    items = {\n",
    "        \"length\": torch.tensor([batch_item[\"length\"] for batch_item in batch]),\n",
    "        \"steps\": torch.tensor([batch_item[\"steps\"] for batch_item in batch]),\n",
    "    }\n",
    "\n",
    "    # find shape that the batch will have\n",
    "    max_length = items[\"length\"].max()\n",
    "    num_items = len(batch)\n",
    "    if num_items > 1:\n",
    "        # padding list\n",
    "        for (key, padding) in [\n",
    "            (\"parts\", 0),\n",
    "            (\"content_ids\", 13942),\n",
    "            (\"answers\", 3),\n",
    "            (\"timestamps\", 0.0),  # note timestamps isnt an embedding\n",
    "            (\"prior_q_times\", 0),\n",
    "            (\"tags\", 188),\n",
    "            (\"row_ids\", 0),\n",
    "        ]:\n",
    "            items[key] = pad_sequence(\n",
    "                [batch_item[key] for batch_item in batch],\n",
    "                batch_first=False,\n",
    "                padding_value=padding,\n",
    "            )\n",
    "    else:\n",
    "        for key in [\n",
    "            \"parts\",\n",
    "            \"content_ids\",\n",
    "            \"answers\",\n",
    "            \"timestamps\",\n",
    "            \"tags\",\n",
    "            \"row_ids\",\n",
    "            \"prior_q_times\",\n",
    "        ]:\n",
    "            items[key] = batch[0][key].unsqueeze(1)\n",
    "    return items"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WINDOW_SIZE = 100\n",
    "inference_dataset = InferenceRIIDDataset(\n",
    "    hdf5_file=\"../input/nn-data-no-lec/feats.h5\", window_size=WINDOW_SIZE\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_for_df(df):\n",
    "    # select questions only\n",
    "    unique_users_with_questions = df[~df.content_type_id].user_id.unique()\n",
    "    df[\"answered_correctly\"] = 0.5  # set useless value for column\n",
    "\n",
    "    # batch with only lectures\n",
    "    if len(unique_users_with_questions) < 1:\n",
    "        pass  # return empy\n",
    "\n",
    "    # case of single user with questions in batch\n",
    "    elif len(unique_users_with_questions) == 1:\n",
    "        item = inference_dataset[unique_users_with_questions[0]].copy()\n",
    "        predictions = model.predict_fast_single_user(\n",
    "            item[\"content_ids\"],\n",
    "            item[\"parts\"],\n",
    "            item[\"answers\"],\n",
    "            item[\"tags\"],\n",
    "            item[\"timestamps\"],\n",
    "            item[\"prior_q_times\"],\n",
    "            n=item[\"steps\"],\n",
    "        ).cpu()\n",
    "        df.loc[\n",
    "            df.user_id == unique_users_with_questions[0], \"answered_correctly\"\n",
    "        ] = predictions.numpy()\n",
    "    # case of multiple users with questions in batch\n",
    "    else:\n",
    "        batch = inference_collate_fn(\n",
    "            [inference_dataset[u].copy() for u in unique_users_with_questions]\n",
    "        )\n",
    "        predictions, row_ids = model.predict_n_steps(batch, batch[\"steps\"])\n",
    "        df[\"answered_correctly\"] = df[\"row_id\"].map(\n",
    "            dict(zip(row_ids.cpu().numpy(), predictions.cpu().numpy()))\n",
    "        )\n",
    "    return df[~df.content_type_id]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterate over test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import riiideducation\n",
    "env = riiideducation.make_env()\n",
    "iter_test = env.iter_test()\n",
    "set_predict = env.predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "previous_test_df = None\n",
    "for (current_test, sample_prediction_df) in iter_test:\n",
    "    if previous_test_df is not None:\n",
    "        previous_test_df[\"answered_correctly\"] = eval(current_test[\"prior_group_answers_correct\"].iloc[0])\n",
    "        previous_test_df[~previous_test_df.content_type_id].groupby(\n",
    "            \"user_id\"\n",
    "        ).answered_correctly.apply(inference_dataset.update_answered_correctly)\n",
    "        \n",
    "    # your feature extraction and model training code here\n",
    "    current_test = preprocess_df(current_test)\n",
    "    previous_test_df = current_test.copy()\n",
    "    current_test = current_test[~current_test.content_type_id]\n",
    "\n",
    "    # add current to cache\n",
    "    current_test[\n",
    "        [\"row_id\", \"user_id\", \"content_id\", \"timestamp\", \"answered_correctly\", \"prior_question_elapsed_time\"]\n",
    "    ].groupby(\"user_id\").apply(\n",
    "        lambda user_rows: inference_dataset.update_user_rows(user_rows)\n",
    "    )\n",
    "\n",
    "    # your prediction code here\n",
    "    current_test = predict_for_df(current_test)\n",
    "    set_predict(current_test[[\"row_id\", \"answered_correctly\"]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}