RNN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "5a28efa6caa45fb2f3142586b9c59adc",
          "grade": false,
          "grade_id": "cell-52506fc51faeb1a1",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "# Recurent Neural Network\n",
        "\n",
        "## Overview\n",
        "\n",
        "Build a bi-directional RNN on diagnosis codes. The recurrent nature of RNN allows us to model the temporal relation of different visits of a patient. More specifically, we will still perform **Heart Failure Prediction**, but with different input formats."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "32e72084469253ba7b428e2d0bd46613",
          "grade": false,
          "grade_id": "cell-dcd6c662fba70926",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.621466Z",
          "start_time": "2021-12-10T02:49:07.849594Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "2db74f9db6e96fc42296ed510b0f4be1",
          "grade": false,
          "grade_id": "cell-4fe346254a16fed8",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "import pickle\n",
        "import random\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# set seed\n",
        "seed = 24\n",
        "random.seed(seed)\n",
        "np.random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "os.environ[\"PYTHONHASHSEED\"] = str(seed)\n",
        "\n",
        "\n",
        "# Define data path\n",
        "DATA_PATH = \"../data\" "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "212e434fd38be5ca223e82a1e1fddf5b",
          "grade": false,
          "grade_id": "cell-71f2f1fcbf0214c3",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "08063ba06edc72626d45a1f16564745e",
          "grade": false,
          "grade_id": "cell-f24c5a8a552afa64",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "## About Raw Data\n",
        "\n",
        "To get started, we will implement a naive RNN model for heart failure prediction using the diagnosis codes.\n",
        "\n",
        "We will use the same dataset synthesized from [MIMIC-III](https://mimic.physionet.org/gettingstarted/access/), but with different input formats.\n",
        "\n",
        "The data has been preprocessed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.630910Z",
          "start_time": "2021-12-10T02:49:08.623252Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "d0dd1f4063e22be64f2709deffde7a7b",
          "grade": false,
          "grade_id": "cell-0d031c45ba4a787e",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "pids = pickle.load(open(os.path.join(DATA_PATH,'train/pids.pkl'), 'rb'))\n",
        "vids = pickle.load(open(os.path.join(DATA_PATH,'train/vids.pkl'), 'rb'))\n",
        "hfs = pickle.load(open(os.path.join(DATA_PATH,'train/hfs.pkl'), 'rb'))\n",
        "seqs = pickle.load(open(os.path.join(DATA_PATH,'train/seqs.pkl'), 'rb'))\n",
        "types = pickle.load(open(os.path.join(DATA_PATH,'train/types.pkl'), 'rb'))\n",
        "rtypes = pickle.load(open(os.path.join(DATA_PATH,'train/rtypes.pkl'), 'rb'))\n",
        "\n",
        "assert len(pids) == len(vids) == len(hfs) == len(seqs) == 1000\n",
        "assert len(types) == 619"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "1427cf82d51752cd4e90e7d483141ffe",
          "grade": false,
          "grade_id": "cell-66a0abe057d9ca85",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "where\n",
        "\n",
        "- `pids`: contains the patient ids\n",
        "- `vids`: contains a list of visit ids for each patient\n",
        "- `hfs`: contains the heart failure label (0: normal, 1: heart failure) for each patient\n",
        "- `seqs`: contains a list of visit (in ICD9 codes) for each patient\n",
        "- `types`: contains the map from ICD9 codes to ICD-9 labels\n",
        "- `rtypes`: contains the map from ICD9 labels to ICD9 codes\n",
        "\n",
        "Let us take a patient as an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.636763Z",
          "start_time": "2021-12-10T02:49:08.632459Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "a0d3847b04bd6dadf5ff567bee973690",
          "grade": false,
          "grade_id": "cell-ae331190a6d48106",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "# take the 3rd patient as an example\n",
        "\n",
        "print(\"Patient ID:\", pids[3])\n",
        "print(\"Heart Failure:\", hfs[3])\n",
        "print(\"# of visits:\", len(vids[3]))\n",
        "for visit in range(len(vids[3])):\n",
        "    print(f\"\\t{visit}-th visit id:\", vids[3][visit])\n",
        "    print(f\"\\t{visit}-th visit diagnosis labels:\", seqs[3][visit])\n",
        "    print(f\"\\t{visit}-th visit diagnosis codes:\", [rtypes[label] for label in seqs[3][visit]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "f6484d87ff7b7dcc915b95cd6496a47f",
          "grade": false,
          "grade_id": "cell-945119717fb61cc7",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "Note that `seqs` is a list of list of list. That is, `seqs[i][j][k]` gives you the k-th diagnosis codes for the j-th visit for the i-th patient.\n",
        "\n",
        "And you can look up the meaning of the ICD9 code online. For example, `DIAG_276` represetns *disorders of fluid electrolyte and acid-base balance*.\n",
        "\n",
        "Further, let see number of heart failure patients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.642695Z",
          "start_time": "2021-12-10T02:49:08.639962Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "06d423d6893adfdc928488e362a86f3a",
          "grade": false,
          "grade_id": "cell-e6d339169f140694",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "print(\"number of heart failure patients:\", sum(hfs))\n",
        "print(\"ratio of heart failure patients: %.2f\" % (sum(hfs) / len(hfs)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2020-10-21T13:05:03.202250Z",
          "start_time": "2020-10-21T13:05:03.199011Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "c20a3d6188476868bb0477e6493dcc29",
          "grade": false,
          "grade_id": "cell-0a48d6dcc0f5b4ad",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "Now we have the data. Let us build the naive RNN."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "dfdb69398ff06e79f0ec1d8faeeb4d5c",
          "grade": false,
          "grade_id": "cell-308c526175fdb62e",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "## 1 Build the dataset\n",
        "\n",
        "### 1.1 CustomDataset\n",
        "\n",
        "First, implement a custom dataset using PyTorch class `Dataset`, which will characterize the key features of the dataset we want to generate.\n",
        "\n",
        "We will use the sequences of diagnosis codes `seqs` as input and heart failure `hfs` as output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.648717Z",
          "start_time": "2021-12-10T02:49:08.644248Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import Dataset\n",
        "\n",
        "\n",
        "class CustomDataset(Dataset):\n",
        "    \n",
        "    def __init__(self, seqs, hfs):\n",
        "        \n",
        "        \"\"\"\n",
        "        TODO: Store `seqs`. to `self.x` and `hfs` to `self.y`.\n",
        "        \n",
        "        Note that you DO NOT need to covert them to tensor as we will do this later.\n",
        "        Do NOT permute the data.\n",
        "        \"\"\"\n",
        "        \n",
        "        # your code here\n",
        "        raise NotImplementedError\n",
        "    \n",
        "    def __len__(self):\n",
        "        \n",
        "        \"\"\"\n",
        "        TODO: Return the number of samples (i.e. patients).\n",
        "        \"\"\"\n",
        "        \n",
        "        # your code here\n",
        "        raise NotImplementedError\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        \n",
        "        \"\"\"\n",
        "        TODO: Generates one sample of data.\n",
        "        \n",
        "        Note that you DO NOT need to covert them to tensor as we will do this later.\n",
        "        \"\"\"\n",
        "        \n",
        "        # your code here\n",
        "        raise NotImplementedError\n",
        "        \n",
        "\n",
        "dataset = CustomDataset(seqs, hfs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.658110Z",
          "start_time": "2021-12-10T02:49:08.650497Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "9698147ad965e9a4336317e632e30259",
          "grade": true,
          "grade_id": "cell-cc0baa6c9dadef8c",
          "locked": true,
          "points": 5,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "dataset = CustomDataset(seqs, hfs)\n",
        "\n",
        "assert len(dataset) == 1000\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "c4ef302b3eea1353a0c52f3110cec061",
          "grade": false,
          "grade_id": "cell-de0d816943d88377",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 1.2 Collate Function\n",
        "\n",
        "As you note that, we do not convert the data to tensor in the built `CustomDataset`. Instead, we will do this using a collate function `collate_fn()`. \n",
        "\n",
        "This collate function `collate_fn()` will be called by `DataLoader` after fetching a list of samples using the indices from `CustomDataset` to collate the list of samples into batches.\n",
        "\n",
        "For example, assume the `DataLoader` gets a list of two samples.\n",
        "\n",
        "```\n",
        "[ [ [0, 1, 2], [8, 0] ], \n",
        "  [ [12, 13, 6, 7], [12], [23, 11] ] ]\n",
        "```\n",
        "\n",
        "where the first sample has two visits `[0, 1, 2]` and `[8, 0]` and the second sample has three visits `[12, 13, 6, 7]`, `[12]`, and `[23, 11]`.\n",
        "\n",
        "The collate function `collate_fn()` is supposed to pad them into the same shape (3, 4), where 3 is the maximum number of visits and 4 is the maximum number of diagnosis codes.\n",
        "\n",
        "``` \n",
        "[ [ [0, 1, 2, *0*], [8, 0, *0*, *0*], [*0*, *0*, *0*, *0*]  ], \n",
        "  [ [12, 13, 6, 7], [12, *0*, *0*, *0*], [23, 11, *0*, *0*] ] ]\n",
        "```\n",
        "\n",
        "Further, the padding information will be stored in a mask with the same shape, where 1 indicates that the diagnosis code at this position is from the original input, and 0 indicates that the diagnosis code at this position is the padded value.\n",
        "\n",
        "```\n",
        "[ [ [1, 1, 1, 0], [1, 1, 0, 0], [0, 0, 0, 0] ], \n",
        "  [ [1, 1, 1, 1], [1, 0, 0, 0], [1, 1, 0, 0] ] ]\n",
        "```\n",
        "\n",
        "Lastly, we will have another diagnosis sequence in reversed time. This will be used in our RNN model for masking. Note that we only flip the true visits.\n",
        "\n",
        "``` \n",
        "[ [ [8, 0, *0*, *0*], [0, 1, 2, *0*], [*0*, *0*, *0*, *0*]  ], \n",
        "  [ [23, 11, *0*, *0*], [12, *0*, *0*, *0*], [12, 13, 6, 7] ] ]\n",
        "```\n",
        "\n",
        "And a reversed mask as well.\n",
        "\n",
        "```\n",
        "[ [ [1, 1, 0, 0], [1, 1, 1, 0], [0, 0, 0, 0] ], \n",
        "  [ [1, 1, 0, 0], [1, 0, 0, 0], [1, 1, 1, 1], ] ]\n",
        "```\n",
        "\n",
        "We need to pad the sequences into the same length so that we can do batch training on GPU. And we also need this mask so that when training, we can ignored the padded value as they actually do not contain any information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.667617Z",
          "start_time": "2021-12-10T02:49:08.659716Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "def collate_fn(data):\n",
        "    \"\"\"\n",
        "    TODO: Collate the the list of samples into batches. For each patient, you need to pad the diagnosis\n",
        "        sequences to the sample shape (max # visits, max # diagnosis codes). The padding infomation\n",
        "        is stored in `mask`.\n",
        "    \n",
        "    Arguments:\n",
        "        data: a list of samples fetched from `CustomDataset`\n",
        "        \n",
        "    Outputs:\n",
        "        x: a tensor of shape (# patiens, max # visits, max # diagnosis codes) of type torch.long\n",
        "        masks: a tensor of shape (# patiens, max # visits, max # diagnosis codes) of type torch.bool\n",
        "        rev_x: same as x but in reversed time. This will be used in our RNN model for masking \n",
        "        rev_masks: same as mask but in reversed time. This will be used in our RNN model for masking\n",
        "        y: a tensor of shape (# patiens) of type torch.float\n",
        "        \n",
        "    Note that you can obtains the list of diagnosis codes and the list of hf labels\n",
        "        using: `sequences, labels = zip(*data)`\n",
        "    \"\"\"\n",
        "\n",
        "    sequences, labels = zip(*data)\n",
        "\n",
        "    y = torch.tensor(labels, dtype=torch.float)\n",
        "    \n",
        "    num_patients = len(sequences)\n",
        "    num_visits = [len(patient) for patient in sequences]\n",
        "    num_codes = [len(visit) for patient in sequences for visit in patient]\n",
        "\n",
        "    max_num_visits = max(num_visits)\n",
        "    max_num_codes = max(num_codes)\n",
        "    \n",
        "    x = torch.zeros((num_patients, max_num_visits, max_num_codes), dtype=torch.long)\n",
        "    rev_x = torch.zeros((num_patients, max_num_visits, max_num_codes), dtype=torch.long)\n",
        "    masks = torch.zeros((num_patients, max_num_visits, max_num_codes), dtype=torch.bool)\n",
        "    rev_masks = torch.zeros((num_patients, max_num_visits, max_num_codes), dtype=torch.bool)\n",
        "    for i_patient, patient in enumerate(sequences):\n",
        "        for j_visit, visit in enumerate(patient):\n",
        "            \"\"\"\n",
        "            TODO: update `x`, `rev_x`, `masks`, and `rev_masks`\n",
        "            \"\"\"\n",
        "            # your code here\n",
        "            raise NotImplementedError\n",
        "    \n",
        "    return x, masks, rev_x, rev_masks, y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.682628Z",
          "start_time": "2021-12-10T02:49:08.669277Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "249be2d3863d3a27bb7f326e17a872f0",
          "grade": true,
          "grade_id": "cell-4b3472bbf5973793",
          "locked": true,
          "points": 20,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "loader = DataLoader(dataset, batch_size=10, collate_fn=collate_fn)\n",
        "loader_iter = iter(loader)\n",
        "x, masks, rev_x, rev_masks, y = next(loader_iter)\n",
        "\n",
        "assert x.dtype == rev_x.dtype == torch.long\n",
        "assert y.dtype == torch.float\n",
        "assert masks.dtype == rev_masks.dtype == torch.bool\n",
        "\n",
        "assert x.shape == rev_x.shape == masks.shape == rev_masks.shape == (10, 3, 24)\n",
        "assert y.shape == (10,)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "cf11addba53d9041094d45051435cb7e",
          "grade": false,
          "grade_id": "cell-125312ce2d90406a",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "Now we have `CustomDataset` and `collate_fn()`. Let us split the dataset into training and validation sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.688088Z",
          "start_time": "2021-12-10T02:49:08.684315Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "cc4a76fda00ecaef5a2e8857a49fb5f2",
          "grade": false,
          "grade_id": "cell-7f2e734b97c94232",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "from torch.utils.data.dataset import random_split\n",
        "\n",
        "split = int(len(dataset)*0.8)\n",
        "\n",
        "lengths = [split, len(dataset) - split]\n",
        "train_dataset, val_dataset = random_split(dataset, lengths)\n",
        "\n",
        "print(\"Length of train dataset:\", len(train_dataset))\n",
        "print(\"Length of val dataset:\", len(val_dataset))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "de98a7101e7388850706de0e357437ad",
          "grade": false,
          "grade_id": "cell-c9732f7be72cb6e4",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 1.3 DataLoader\n",
        "\n",
        "Now, we can load the dataset into the data loader."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.696240Z",
          "start_time": "2021-12-10T02:49:08.692166Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "def load_data(train_dataset, val_dataset, collate_fn):\n",
        "    \n",
        "    '''\n",
        "    TODO: Implement this function to return the data loader for  train and validation dataset. \n",
        "    Set batchsize to 32. Set `shuffle=True` only for train dataloader.\n",
        "    \n",
        "    Arguments:\n",
        "        train dataset: train dataset of type `CustomDataset`\n",
        "        val dataset: validation dataset of type `CustomDataset`\n",
        "        collate_fn: collate function\n",
        "        \n",
        "    Outputs:\n",
        "        train_loader, val_loader: train and validation dataloaders\n",
        "    \n",
        "    Note that you need to pass the collate function to the data loader `collate_fn()`.\n",
        "    '''\n",
        "    \n",
        "    batch_size = 32\n",
        "    # your code here\n",
        "    raise NotImplementedError\n",
        "    \n",
        "    return train_loader, val_loader\n",
        "\n",
        "\n",
        "train_loader, val_loader = load_data(train_dataset, val_dataset, collate_fn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.701374Z",
          "start_time": "2021-12-10T02:49:08.698130Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "d3e8781ea70db7a1afea7413cef8c3cf",
          "grade": true,
          "grade_id": "cell-0c30a49563819f13",
          "locked": true,
          "points": 5,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "train_loader, val_loader = load_data(train_dataset, val_dataset, collate_fn)\n",
        "\n",
        "assert len(train_loader) == 25, \"Length of train_loader should be 25, instead we got %d\"%(len(train_loader))\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "3c7771a13f558c5045dc0c4d2d2b39c8",
          "grade": false,
          "grade_id": "cell-9739d5ae7e1cafc5",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "## 2 Naive RNN\n",
        "\n",
        "Implement a naive bi-directional RNN model.\n",
        "\n",
        "<img src=\"img/bi-rnn.jpg\" width=\"600\"/>\n",
        "\n",
        "Remember from class that, first of all, we need to transform the diagnosis code for each visit of a patient to an embedding. To do this, we can use `nn.Embedding()`, where `num_embeddings` is the number of diagnosis codes and `embedding_dim` is the embedding dimension."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "cf997f094672a3be91cc5e296d32ba1d",
          "grade": false,
          "grade_id": "cell-7fa15685c339c1c7",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "Then, we can construct a simple RNN structure. Each input is this multi-hot vector. At the 0-th visit, this has $\\boldsymbol{X}_0$, and at t-th visit, this has $\\boldsymbol{X}_t$.\n",
        "\n",
        "Each one of the input will then map to a hidden state $\\boldsymbol{\\overleftrightarrow{h}}_t$. The forward hidden state $\\boldsymbol{\\overrightarrow{h}}_t$ can be determined by $\\boldsymbol{\\overrightarrow{h}}_{t-1}$ and the corresponding current input $\\boldsymbol{X}_t$.\n",
        "\n",
        "Similarly, we will have another RNN to process the sequence in the reverse order, so that the hidden state $\\boldsymbol{\\overleftarrow{h}}_t$ is determined by $\\boldsymbol{\\overleftarrow{h}}_{t+1}$ and $\\boldsymbol{X}_t$.\n",
        "\n",
        "Finally, once we have the $\\boldsymbol{\\overrightarrow{h}}_T$ and $\\boldsymbol{\\overleftarrow{h}}_{0}$, we will concatenate the two vectors as the feature vector and train a NN to perform the classification."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "70664a7944bd3f61aaf14e3a311e83d4",
          "grade": false,
          "grade_id": "cell-750f2f7f3226048f",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "Now, let us build this model. The forward steps will be:\n",
        "\n",
        "    1. Pass the sequence through the embedding layer;\n",
        "    2. Sum the embeddings for each diagnosis code up for a visit of a patient;\n",
        "    3. Pass the embeddings through the RNN layer;\n",
        "    4. Obtain the hidden state at the last visit;\n",
        "    5. Do 1-4 for both directions and concatenate the hidden states.\n",
        "    6. Pass the hidden state through the linear and activation layers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "2aaf3a348d4abc77607b8b288b56bb8b",
          "grade": false,
          "grade_id": "cell-a9cb7f4d8889ca27",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 2.1 Mask Selection\n",
        "\n",
        "Importantly, you need to use `masks` to mask out the paddings in before step 2 and before 4. So, let us first preform the mask selection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.705960Z",
          "start_time": "2021-12-10T02:49:08.703151Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "def sum_embeddings_with_mask(x, masks):\n",
        "    \"\"\"\n",
        "    TODO: mask select the embeddings for true visits (not padding visits) and then\n",
        "        sum the embeddings for each visit up.\n",
        "\n",
        "    Arguments:\n",
        "        x: the embeddings of diagnosis sequence of shape (batch_size, # visits, # diagnosis codes, embedding_dim)\n",
        "        masks: the padding masks of shape (batch_size, # visits, # diagnosis codes)\n",
        "\n",
        "    Outputs:\n",
        "        sum_embeddings: the sum of embeddings of shape (batch_size, # visits, embedding_dim)\n",
        "        \n",
        "    NOTE: Do NOT use for loop.\n",
        "\n",
        "    \"\"\"\n",
        "    \n",
        "    # your code here\n",
        "    raise NotImplementedError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.769541Z",
          "start_time": "2021-12-10T02:49:08.707555Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "bf2ca473060eda35326593899a7fb2ae",
          "grade": true,
          "grade_id": "cell-e2e6868f7bd913f8",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "import random\n",
        "import ast\n",
        "import inspect\n",
        "\n",
        "\n",
        "def uses_loop(function):\n",
        "    loop_statements = ast.For, ast.While, ast.AsyncFor\n",
        "\n",
        "    nodes = ast.walk(ast.parse(inspect.getsource(function)))\n",
        "    return any(isinstance(node, loop_statements) for node in nodes)\n",
        "\n",
        "def generate_random_mask(batch_size, max_num_visits , max_num_codes):\n",
        "    num_visits = [random.randint(1, max_num_visits) for _ in range(batch_size)]\n",
        "    num_codes = []\n",
        "    for n in num_visits:\n",
        "        num_codes_visit = [0] * max_num_visits\n",
        "        for i in range(n):\n",
        "            num_codes_visit[i] = (random.randint(1, max_num_codes))\n",
        "        num_codes.append(num_codes_visit)\n",
        "    masks = [torch.ones((l,), dtype=torch.bool) for num_codes_visit in num_codes for l in num_codes_visit]\n",
        "    masks = torch.stack([torch.cat([i, i.new_zeros(max_num_codes - i.size(0))], 0) for i in masks], 0)\n",
        "    masks = masks.view((batch_size, max_num_visits, max_num_codes)).bool()\n",
        "    return masks\n",
        "\n",
        "\n",
        "batch_size = 16\n",
        "max_num_visits = 10\n",
        "max_num_codes = 20\n",
        "embedding_dim = 100\n",
        "\n",
        "torch.random.manual_seed(7)\n",
        "x = torch.randn((batch_size, max_num_visits , max_num_codes, embedding_dim))\n",
        "masks = generate_random_mask(batch_size, max_num_visits , max_num_codes)\n",
        "out = sum_embeddings_with_mask(x, masks)\n",
        "\n",
        "assert uses_loop(sum_embeddings_with_mask) is False\n",
        "assert out.shape == (batch_size, max_num_visits, embedding_dim)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.774272Z",
          "start_time": "2021-12-10T02:49:08.770811Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "def get_last_visit(hidden_states, masks):\n",
        "    \"\"\"\n",
        "    TODO: obtain the hidden state for the last true visit (not padding visits)\n",
        "\n",
        "    Arguments:\n",
        "        hidden_states: the hidden states of each visit of shape (batch_size, # visits, embedding_dim)\n",
        "        masks: the padding masks of shape (batch_size, # visits, # diagnosis codes)\n",
        "\n",
        "    Outputs:\n",
        "        last_hidden_state: the hidden state for the last true visit of shape (batch_size, embedding_dim)\n",
        "        \n",
        "    NOTE: DO NOT use for loop.\n",
        "    \n",
        "    HINT: Consider using `torch.gather()`.\n",
        "    \"\"\"\n",
        "    \n",
        "    # your code here\n",
        "    raise NotImplementedError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.788183Z",
          "start_time": "2021-12-10T02:49:08.775986Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "a57f7f7d77bd03e1d4f0b3e53cca7ec7",
          "grade": true,
          "grade_id": "cell-611bb60b8cff5f77",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "assert uses_loop(get_last_visit) is False\n",
        "\n",
        "max_num_visits = 10\n",
        "batch_size = 16\n",
        "max_num_codes = 20\n",
        "embedding_dim = 100\n",
        "\n",
        "torch.random.manual_seed(7)\n",
        "hidden_states = torch.randn((batch_size, max_num_visits, embedding_dim))\n",
        "masks = generate_random_mask(batch_size, max_num_visits , max_num_codes)\n",
        "out = get_last_visit(hidden_states, masks)\n",
        "\n",
        "assert out.shape == (batch_size, embedding_dim)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "62bbd9209d97a0d0469cb4f9ff414d45",
          "grade": false,
          "grade_id": "cell-51a88c33b34e6827",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 2.2 Build NaiveRNN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.807198Z",
          "start_time": "2021-12-10T02:49:08.789522Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "class NaiveRNN(nn.Module):\n",
        "    \n",
        "    \"\"\"\n",
        "    TODO: implement the naive RNN model above.\n",
        "    \"\"\"\n",
        "    \n",
        "    def __init__(self, num_codes):\n",
        "        super().__init__()\n",
        "        \"\"\"\n",
        "        TODO: \n",
        "            1. Define the embedding layer using `nn.Embedding`. Set `embDimSize` to 128.\n",
        "            2. Define the RNN using `nn.GRU()`; Set `hidden_size` to 128. Set `batch_first` to True.\n",
        "            2. Define the RNN for the reverse direction using `nn.GRU()`;\n",
        "               Set `hidden_size` to 128. Set `batch_first` to True.\n",
        "            3. Define the linear layers using `nn.Linear()`; Set `in_features` to 256, and `out_features` to 1.\n",
        "            4. Define the final activation layer using `nn.Sigmoid().\n",
        "\n",
        "        Arguments:\n",
        "            num_codes: total number of diagnosis codes\n",
        "        \"\"\"\n",
        "        \n",
        "        self.embedding = None\n",
        "        self.rnn = None\n",
        "        self.rev_rnn = None\n",
        "        self.fc = None\n",
        "        self.sigmoid = None\n",
        "        \n",
        "        # your code here\n",
        "        raise NotImplementedError\n",
        "    \n",
        "    def forward(self, x, masks, rev_x, rev_masks):\n",
        "        \"\"\"\n",
        "        Arguments:\n",
        "            x: the diagnosis sequence of shape (batch_size, # visits, # diagnosis codes)\n",
        "            masks: the padding masks of shape (batch_size, # visits, # diagnosis codes)\n",
        "\n",
        "        Outputs:\n",
        "            probs: probabilities of shape (batch_size)\n",
        "        \"\"\"\n",
        "        \n",
        "        batch_size = x.shape[0]\n",
        "        \n",
        "        # 1. Pass the sequence through the embedding layer;\n",
        "        x = self.embedding(x)\n",
        "        # 2. Sum the embeddings for each diagnosis code up for a visit of a patient.\n",
        "        x = sum_embeddings_with_mask(x, masks)\n",
        "        \n",
        "        # 3. Pass the embegginds through the RNN layer;\n",
        "        output, _ = self.rnn(x)\n",
        "        # 4. Obtain the hidden state at the last visit.\n",
        "        true_h_n = get_last_visit(output, masks)\n",
        "        \n",
        "        \"\"\"\n",
        "        TODO:\n",
        "            5. Do the step 1-4 again for the reverse order, and concatenate the hidden\n",
        "               states for both directions;\n",
        "        \"\"\"\n",
        "        true_h_n_rev = None\n",
        "        # your code here\n",
        "        raise NotImplementedError\n",
        "        \n",
        "        # 6. Pass the hidden state through the linear and activation layers.\n",
        "        logits = self.fc(torch.cat([true_h_n, true_h_n_rev], 1))        \n",
        "        probs = self.sigmoid(logits)\n",
        "        return probs.view(batch_size)\n",
        "    \n",
        "\n",
        "# load the model here\n",
        "naive_rnn = NaiveRNN(num_codes = len(types))\n",
        "naive_rnn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.812597Z",
          "start_time": "2021-12-10T02:49:08.808592Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "1da92a4649f26a9fc60f4d6ded75c7ba",
          "grade": true,
          "grade_id": "cell-45de4813453d610f",
          "locked": true,
          "points": 5,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.822685Z",
          "start_time": "2021-12-10T02:49:08.814539Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "5c4bf456dd6172e11a5b2038801aedda",
          "grade": true,
          "grade_id": "cell-0c00e2f67834bd8d",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "f778ecb7179f34af1319d5e03d7ee599",
          "grade": false,
          "grade_id": "cell-3b34f300dcde3d1b",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "## 3 Model Training\n",
        "\n",
        "### 3.1 Loss and Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.827118Z",
          "start_time": "2021-12-10T02:49:08.823986Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "TODO: Specify Binary Cross Entropy as the loss function (`nn.BCELoss`) and assign it to `criterion`.\n",
        "      Spcify Adam as the optimizer (`torch.optim.Adam`)  with learning rate 0.001 and assign it to `optimizer`.\n",
        "\"\"\"\n",
        "\n",
        "criterion = None\n",
        "optimizer = None\n",
        "\n",
        "# your code here\n",
        "raise NotImplementedError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:08.831838Z",
          "start_time": "2021-12-10T02:49:08.828878Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "f642b0163a31499897a3b2a3e66657f6",
          "grade": true,
          "grade_id": "cell-40bee829e3b63b7d",
          "locked": true,
          "points": 5,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "eab2b782149c6d539961da3b7f25ab09",
          "grade": false,
          "grade_id": "cell-873df7380d762445",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 3.2 Evaluate\n",
        "\n",
        "Then, let us implement the `eval_model()` function first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:09.238600Z",
          "start_time": "2021-12-10T02:49:08.833460Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import precision_recall_fscore_support, roc_auc_score\n",
        "\n",
        "\n",
        "def eval_model(model, val_loader):\n",
        "    \n",
        "    \"\"\"\n",
        "    TODO: evaluate the model.\n",
        "    \n",
        "    Arguments:\n",
        "        model: the RNN model\n",
        "        val_loader: validation dataloader\n",
        "        \n",
        "    Outputs:\n",
        "        precision: overall precision score\n",
        "        recall: overall recall score\n",
        "        f1: overall f1 score\n",
        "        roc_auc: overall roc_auc score\n",
        "        \n",
        "    Note that please pass all four arguments to the model so that we can use this function for both \n",
        "    models. (Use `model(x, masks, rev_x, rev_masks)`.)\n",
        "        \n",
        "    HINT: checkout https://scikit-learn.org/stable/modules/classes.html#module-sklearn.metrics\n",
        "    \"\"\"\n",
        "    \n",
        "    model.eval()\n",
        "    y_pred = torch.LongTensor()\n",
        "    y_score = torch.Tensor()\n",
        "    y_true = torch.LongTensor()\n",
        "    model.eval()\n",
        "    for x, masks, rev_x, rev_masks, y in val_loader:\n",
        "        y_hat = model(x, masks, rev_x, rev_masks)\n",
        "        y_score = torch.cat((y_score,  y_hat.detach().to('cpu')), dim=0)\n",
        "        y_hat = (y_hat > 0.5).int()\n",
        "        y_pred = torch.cat((y_pred,  y_hat.detach().to('cpu')), dim=0)\n",
        "        y_true = torch.cat((y_true, y.detach().to('cpu')), dim=0)\n",
        "    \"\"\"\n",
        "    TODO:\n",
        "        Calculate precision, recall, f1, and roc auc scores.\n",
        "        Use `average='binary'` for calculating precision, recall, and fscore.\n",
        "    \"\"\"\n",
        "    p, r, f, roc_auc = None, None, None, None\n",
        "    # your code here\n",
        "    raise NotImplementedError\n",
        "    return p, r, f, roc_auc"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:09.409363Z",
          "start_time": "2021-12-10T02:49:09.240136Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "197e373f11419054f6005341d13867bb",
          "grade": true,
          "grade_id": "cell-764df4f66a8f01e4",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "\n",
        "p, r, f, roc_auc = eval_model(naive_rnn, val_loader)\n",
        "assert p.size == 1, \"Precision should be a scalar.\"\n",
        "assert r.size == 1, \"Recall should be a scalar.\"\n",
        "assert f.size == 1, \"F1 should be a scalar.\"\n",
        "assert roc_auc.size == 1, \"ROC-AUC should be a scalar.\"\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "markdown",
          "checksum": "4f2026101cfd198567edf0c5d2fe71b8",
          "grade": false,
          "grade_id": "cell-9b3672b70944a8d9",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "source": [
        "### 3.3 Training and evlauation\n",
        "\n",
        "Now let us implement the `train()` function. Note that `train()` should call `eval_model()` at the end of each training epoch to see the results on the validaion dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:09.415838Z",
          "start_time": "2021-12-10T02:49:09.410578Z"
        },
        "deletable": false
      },
      "outputs": [],
      "source": [
        "def train(model, train_loader, val_loader, n_epochs):\n",
        "    \"\"\"\n",
        "    TODO: train the model.\n",
        "    \n",
        "    Arguments:\n",
        "        model: the RNN model\n",
        "        train_loader: training dataloder\n",
        "        val_loader: validation dataloader\n",
        "        n_epochs: total number of epochs\n",
        "        \n",
        "    You need to call `eval_model()` at the end of each training epoch to see how well the model performs \n",
        "    on validation data.\n",
        "        \n",
        "    Note that please pass all four arguments to the model so that we can use this function for both \n",
        "    models. (Use `model(x, masks, rev_x, rev_masks)`.)\n",
        "    \"\"\"\n",
        "    \n",
        "    model.train()\n",
        "    for epoch in range(n_epochs):\n",
        "        train_loss = 0\n",
        "        for x, masks, rev_x, rev_masks, y in train_loader:\n",
        "            \"\"\"\n",
        "            TODO:\n",
        "                1. zero grad\n",
        "                2. model forward\n",
        "                3. calculate loss\n",
        "                4. loss backward\n",
        "                5. optimizer step\n",
        "            \"\"\"\n",
        "            loss = None\n",
        "            # your code here\n",
        "            raise NotImplementedError\n",
        "            train_loss += loss.item()\n",
        "        train_loss = train_loss / len(train_loader)\n",
        "        print('Epoch: {} \\t Training Loss: {:.6f}'.format(epoch+1, train_loss))\n",
        "        p, r, f, roc_auc = eval_model(model, val_loader)\n",
        "        print('Epoch: {} \\t Validation p: {:.2f}, r:{:.2f}, f: {:.2f}, roc_auc: {:.2f}'\n",
        "              .format(epoch+1, p, r, f, roc_auc))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:12.664804Z",
          "start_time": "2021-12-10T02:49:09.417438Z"
        }
      },
      "outputs": [],
      "source": [
        "# number of epochs to train the model\n",
        "n_epochs = 5\n",
        "train(naive_rnn, train_loader, val_loader, n_epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:12.792877Z",
          "start_time": "2021-12-10T02:49:12.666169Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "b2ab4e44c6b1125682a2d8ca02030673",
          "grade": true,
          "grade_id": "cell-8fc0a72d1a31aa34",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n",
        "p, r, f, roc_auc = eval_model(naive_rnn, val_loader)\n",
        "print(roc_auc)\n",
        "assert roc_auc > 0.7, \"ROC AUC is too low on the validation set (%f < 0.7)\"%(roc_auc)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "ExecuteTime": {
          "end_time": "2021-12-10T02:49:12.869415Z",
          "start_time": "2021-12-10T02:49:12.794442Z"
        },
        "deletable": false,
        "editable": false,
        "nbgrader": {
          "cell_type": "code",
          "checksum": "0ed19dc555eb7ae8c97e038425457774",
          "grade": true,
          "grade_id": "cell-8e9d1d7cb1c3a386",
          "locked": true,
          "points": 10,
          "schema_version": 3,
          "solution": false,
          "task": false
        }
      },
      "outputs": [],
      "source": [
        "'''\n",
        "AUTOGRADER CELL. DO NOT MODIFY THIS.\n",
        "'''\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.12"
    },
    "toc": {
      "base_numbering": 1,
      "nav_menu": {},
      "number_sections": false,
      "sideBar": true,
      "skip_h1_title": true,
      "title_cell": "Table of Contents",
      "title_sidebar": "Contents",
      "toc_cell": false,
      "toc_position": {
        "height": "calc(100% - 180px)",
        "left": "10px",
        "top": "150px",
        "width": "358.390625px"
      },
      "toc_section_display": true,
      "toc_window_display": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}