riiid-trax-transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "riiid-trax-transformer.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lAula_PU9jqB"
      },
      "source": [
        "## Intro\r\n",
        "\r\n",
        "This notebook trains a transformer model on the [EdNet dataset](https://github.com/riiid/ednet) using the [google/trax library](https://github.com/google/trax). The EdNet dataset is large set of student responses to multiple choice questions related to English language learning. A recent Kaggle competition, [Riiid! Answer Correctness Prediction](https://www.kaggle.com/c/riiid-test-answer-prediction), provided as subset of this data, consisting of 100 million responses to 13 thousand questions from 300 thousand students.\r\n",
        "\r\n",
        "The state of the art result, detailed in [SAINT+: Integrating Temporal Features for EdNet Correctness Prediction](https://arxiv.org/abs/2010.12042), achieves an AUC ROC of 0.7914. The winning solution in the [Riiid! Answer Correctness Prediction](https://www.kaggle.com/c/riiid-test-answer-prediction) competition achieved an AUC ROC of 0.820. This notebook achieves an AUC ROC of 0.776 implementing an approach similar to the state of the art approach, training for 25,000 steps. It demonstrates several techniques that may be useful to those getting started with the [google/trax library](https://github.com/google/trax) or deep learning in general. This notebook demonstrates how to:\r\n",
        "\r\n",
        "* Use BigQuery to perform feature engineering\r\n",
        "* Create TFRecords with multiple sequences per record\r\n",
        "* Modify the trax Transformer model to accommodate a knowledge tracing dataset:\r\n",
        "    * Utilize multiple encoder and decoder embeddings - aggregated either by concatenation or sum\r\n",
        "    * Include a custom metric - AUC ROC\r\n",
        "    * Utilize a combined padding and future mask\r\n",
        "* Use trax's [gin-config](https://github.com/google/gin-config) integration to specify training parameters\r\n",
        "* Display training progress using trax's tensorboard integration\r\n",
        "\r\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/CalebEverett/riiid_transformer/blob/master/riiid-trax-transformer.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tuG_-VFcpxLc"
      },
      "source": [
        "# Choose a location for your storage bucket and BigQuery dataset to minimize data egress charges. Once you have\r\n",
        "# created them, if you restart your notebook you can run this to see where your colab is running\r\n",
        "# and factory reset until you get a location that is near your data.\r\n",
        "!curl ipinfo.io"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_SQN6SX89XNq"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vo5bzc9z7nw_"
      },
      "source": [
        "# <hide-output>\r\n",
        "!git clone https://github.com/google/trax.git\r\n",
        "!pip install ./trax\r\n",
        "!pip install -U pyarrow\r\n",
        "!pip install -U google-cloud-bigquery google-cloud-bigquery-storage"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0W7kto2g7Sfa"
      },
      "source": [
        "from functools import partial\r\n",
        "import json\r\n",
        "import math\r\n",
        "import os\r\n",
        "from pathlib import Path\r\n",
        "import subprocess\r\n",
        "import sys\r\n",
        "import time\r\n",
        "\r\n",
        "import gin\r\n",
        "from google.cloud import storage, bigquery\r\n",
        "from google.cloud.bigquery import LoadJobConfig, QueryJobConfig, \\\r\n",
        "    SchemaField, SourceFormat\r\n",
        "import jax\r\n",
        "from jax.config import config\r\n",
        "import pandas as pd\r\n",
        "import numpy as np\r\n",
        "import requests\r\n",
        "import sqlite3\r\n",
        "import trax\r\n",
        "from trax import fastmath\r\n",
        "from trax import layers as tl\r\n",
        "from trax.fastmath import numpy as tnp\r\n",
        "import tensorflow as tf\r\n",
        "from tqdm.notebook import tqdm\r\n",
        "import zipfile\r\n",
        "\r\n",
        "# Create google credentials and store in drive\r\n",
        "# https://colab.research.google.com/drive/1LWhrqE2zLXqz30T0a0JqXnDPKweqd8ET\r\n",
        "# \r\n",
        "# Create a config.json file with variables for:\r\n",
        "# \"BUCKET\": \"\",\r\n",
        "# \"BQ_DATASET\": \"\",\r\n",
        "# \"KAGGLE_USERNAME\": \"\",\r\n",
        "# \"KAGGLE_KEY\": \"\",\r\n",
        "# \"PROJECT\": \"\",\r\n",
        "# \"LOCATION\": \"\"\r\n",
        "from google.colab import drive\r\n",
        "\r\n",
        "DRIVE = Path('/content/drive/My Drive')\r\n",
        "PATH = 'riiid-transformer'\r\n",
        "\r\n",
        "if not DRIVE.exists():\r\n",
        "    drive.mount(str(DRIVE.parent))\r\n",
        "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = str(DRIVE/PATH/'google.json')\r\n",
        "\r\n",
        "with open(str(DRIVE/PATH/'config.json')) as f:\r\n",
        "    CONFIG = json.load(f)\r\n",
        "    os.environ = {**os.environ, **CONFIG}\r\n",
        "\r\n",
        "from kaggle.api.kaggle_api_extended import KaggleApi\r\n",
        "kaggle_api = KaggleApi()\r\n",
        "kaggle_api.authenticate()\r\n",
        "\r\n",
        "AUTO = tf.data.experimental.AUTOTUNE\r\n",
        "BUCKET = os.getenv('BUCKET', 'riiid-transformer')\r\n",
        "BQ_DATASET = os.getenv('BQ_DATASET', 'my_data')\r\n",
        "LOCATION = os.getenv('LOCATION', 'us-central1')\r\n",
        "PROJECT = os.getenv('PROJECT', 'fastai-caleb')\r\n",
        "\r\n",
        "bucket = storage.Client(project=PROJECT).get_bucket(BUCKET)\r\n",
        "dataset = bigquery.Dataset(f'{PROJECT}.{BQ_DATASET}')\r\n",
        "bq_client = bigquery.Client(project=PROJECT, location=LOCATION)\r\n",
        "\r\n",
        "%matplotlib inline\r\n",
        "from matplotlib import pyplot as plt\r\n",
        "\r\n",
        "%load_ext tensorboard\r\n",
        "\r\n",
        "gin.enter_interactive_mode()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vL0eRGAnyK9x"
      },
      "source": [
        "## Control Panel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YaAhPw-zv1la"
      },
      "source": [
        "These variables can be set to True to run the code in the sections described or False to skip over them after they have been run for the first time."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MNrBIpVPyPGX"
      },
      "source": [
        "USE_TPU = False\r\n",
        "DOWNLOAD_DATASET = False\r\n",
        "LOAD_DATA_TO_BQ = False\r\n",
        "PERFORM_FEATURE_ENGINEERING = False\r\n",
        "TEST_FEATURE_ENGNEERING = False\r\n",
        "CREATE_TFRECORDS = False\r\n",
        "TEST_TFRECORDS = False\r\n",
        "TRAIN_MODEL = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t8Jvva6lBRyI"
      },
      "source": [
        "## Initialize TPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PsczFYbe80ei"
      },
      "source": [
        "if USE_TPU:\r\n",
        "    if 'TPU_DRIVER_MODE' not in globals():\r\n",
        "        url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'\r\n",
        "        resp = requests.post(url)\r\n",
        "        TPU_DRIVER_MODE = 1\r\n",
        "\r\n",
        "    config.FLAGS.jax_xla_backend = \"tpu_driver\"\r\n",
        "    config.FLAGS.jax_backend_target = \"grpc://\" + os.environ['COLAB_TPU_ADDR']\r\n",
        "    print(config.FLAGS.jax_backend_target)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GXP1CnQXBtzd"
      },
      "source": [
        "## Download Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YSAnW-bzBzCE"
      },
      "source": [
        "if DOWNLOAD_DATASET:\r\n",
        "    kaggle_api.competition_download_cli('riiid-test-answer-prediction')\r\n",
        "    with zipfile.ZipFile('riiid-test-answer-prediction.zip', 'r') as zip_ref:\r\n",
        "        zip_ref.extractall()\r\n",
        "    for f in ['train.csv', 'questions.csv', 'lectures.csv']:\r\n",
        "        bucket.blob(f).upload_from_filename(f)\r\n",
        "\r\n",
        "if False:\r\n",
        "    for f in tqdm(['train.csv', 'questions.csv', 'lectures.csv']):\r\n",
        "        bucket.blob(f).download_to_filename(f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wM1_VVnm-61P"
      },
      "source": [
        "## Create BigQuery Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_Eo0iR8Y-5Sv"
      },
      "source": [
        "if False:\r\n",
        "    delete_contents=False\r\n",
        "    bq_client.delete_dataset(BQ_DATASET, delete_contents=delete_contents)\r\n",
        "    print(f'Dataset {dataset.dataset_id} deleted from project {dataset.project}.')\r\n",
        "\r\n",
        "try:\r\n",
        "    dataset = bq_client.get_dataset(dataset.dataset_id)\r\n",
        "    print(f'Dataset {dataset.dataset_id} already exists '\r\n",
        "          f'in location {dataset.location} in project {dataset.project}.')\r\n",
        "except:\r\n",
        "    dataset = bq_client.create_dataset(dataset)\r\n",
        "    print(f'Dataset {dataset.dataset_id} created '\r\n",
        "          f'in location {dataset.location} in project {dataset.project}.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7tZZN449eH-"
      },
      "source": [
        "## Dtypes"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qt70hdhk_j6p"
      },
      "source": [
        "dtypes_orig = {\r\n",
        "    'lectures': {\r\n",
        "        'lecture_id': 'uint16',\r\n",
        "        'tag': 'uint8',\r\n",
        "        'part': 'uint8',\r\n",
        "        'type_of': 'str',\r\n",
        "    },\r\n",
        "    'questions': {\r\n",
        "        'question_id': 'uint16',\r\n",
        "        'bundle_id': 'uint16',\r\n",
        "        'correct_answer': 'uint8',\r\n",
        "        'part': 'uint8',\r\n",
        "        'tags': 'str',\r\n",
        "        \r\n",
        "    },\r\n",
        "    'train': {\r\n",
        "        'row_id': 'int64',\r\n",
        "        'timestamp': 'int64',\r\n",
        "        'user_id': 'int32',\r\n",
        "        'content_id': 'int16',\r\n",
        "        'content_type_id': 'int8',\r\n",
        "        'task_container_id': 'int16',\r\n",
        "        'user_answer': 'int8',\r\n",
        "        'answered_correctly': 'int8',\r\n",
        "        'prior_question_elapsed_time': 'float32', \r\n",
        "        'prior_question_had_explanation': 'bool'\r\n",
        "    }\r\n",
        "    \r\n",
        "}\r\n",
        "\r\n",
        "dtypes_new = {\r\n",
        "    'lectures': {},\r\n",
        "    'questions': {\r\n",
        "        'tags_array': 'str'\r\n",
        "    },\r\n",
        "    'train': {\r\n",
        "        'task_container_id_q': 'int16',\r\n",
        "        'pqet_current': 'int32',\r\n",
        "        'ts_delta': 'int32'\r\n",
        "    }\r\n",
        "}\r\n",
        "\r\n",
        "dtypes = {}\r\n",
        "for table_id in dtypes_orig:\r\n",
        "    dtypes[table_id] = {\r\n",
        "        **dtypes_orig[table_id],\r\n",
        "        **dtypes_new[table_id]\r\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zYIOhHEoDw-v"
      },
      "source": [
        "### Big Query Table Schemas"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q1LEgqZfDulc"
      },
      "source": [
        "# <hide-input>\r\n",
        "type_map = {\r\n",
        "    'int64': 'INT64',\r\n",
        "    'int32': 'INT64',\r\n",
        "    'int16': 'INT64',\r\n",
        "    'int8': 'INT64',\r\n",
        "    'uint8': 'INT64',\r\n",
        "    'uint16': 'INT64',\r\n",
        "    'str': 'STRING',\r\n",
        "    'bool': 'BOOL',\r\n",
        "    'float32': 'FLOAT64'\r\n",
        "}\r\n",
        "\r\n",
        "schemas_orig = {table: [SchemaField(f, type_map[t]) for f, t in\r\n",
        "                   fields.items()] for table, fields in dtypes_orig.items()}\r\n",
        "\r\n",
        "schemas = {}\r\n",
        "for table_id, fields in dtypes_new.items():\r\n",
        "    new_fields = [SchemaField(f, type_map[t]) for\r\n",
        "                  f, t in fields.items() if 'array' not in f]\r\n",
        "    \r\n",
        "    new_array_feilds = [SchemaField(f, 'INT64', 'REPEATED') for\r\n",
        "                  f, t in fields.items() if 'array' in f]\r\n",
        "\r\n",
        "    new_fields += new_array_feilds\r\n",
        "\r\n",
        "    schemas[table_id] = schemas_orig[table_id] + new_fields"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sv7wPwp2EJpH"
      },
      "source": [
        "### Load Tables"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EtBgHrBvC_H3"
      },
      "source": [
        "def load_job_cb(future):\r\n",
        "    \"\"\"Prints update upon completion to output of last run cell.\"\"\"\r\n",
        "    \r\n",
        "    seconds = (future.ended - future.created).total_seconds()\r\n",
        "    print(f'Loaded {future.output_rows:,d} rows to table {future.job_id.split(\"_\")[0]} in '\r\n",
        "        f'{seconds:>4,.1f} sec, {int(future.output_rows / seconds):,d} per sec.')\r\n",
        "\r\n",
        "def load_csv_from_uri(table_id, schemas_orig):\r\n",
        "    full_table_id = f'{BQ_DATASET}.{table_id}'\r\n",
        "\r\n",
        "    job_config = LoadJobConfig(\r\n",
        "        schema=schemas_orig[table_id],\r\n",
        "        source_format=SourceFormat.CSV,\r\n",
        "        skip_leading_rows=1\r\n",
        "        )\r\n",
        "\r\n",
        "    uri = f'gs://{BUCKET}/{table_id}.csv'\r\n",
        "    load_job = bq_client.load_table_from_uri(uri, full_table_id,\r\n",
        "                                            job_config=job_config,\r\n",
        "                                            job_id_prefix=f'{table_id}_')\r\n",
        "    print(f'job {load_job.job_id} started')\r\n",
        "    load_job.add_done_callback(load_job_cb)\r\n",
        "    \r\n",
        "    return load_job"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L44_o0NYEOcC"
      },
      "source": [
        "if LOAD_DATA_TO_BQ:\r\n",
        "    for table_id in dtypes_orig:\r\n",
        "        lj = load_csv_from_uri(table_id, schemas_orig).result()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JUAg3Pz5ImSx"
      },
      "source": [
        "### Update BiqQuery Schemas"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ym62FdoNgU8t"
      },
      "source": [
        "Before performing feature engineering, we have to update the table schemas in Big Query to create columns for the new features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qByuVM7MIr8b"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    for table_id, schema in schemas.items():\r\n",
        "        table = bq_client.get_table(f'{BQ_DATASET}.{table_id}')\r\n",
        "        table.schema = schema\r\n",
        "        table = bq_client.update_table(table, ['schema'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oCHq9dJiFOPh"
      },
      "source": [
        "## Feature Engineering"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d6tN2qREdc9A"
      },
      "source": [
        "Using BigQuery for a dataset of 100 million rows is much faster than using local dataframes. In addition, you get to use the full power of SQL, including [window functions](https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts), which are especially useful for time series feature engineering.\r\n",
        "\r\n",
        "Feature engineering for this problem is fairly minimal and includes:\r\n",
        "* Replacing missing null values for `prior_question_elapsed_time` and `prior_question_had_explanation` in the train table\r\n",
        "* Replacing one missing tag value in the questions table\r\n",
        "* Recalcuating the `task_container_id` as `task_container_id_q` so that it excludes lecture records and increases monotonically with `timetamp` so that the calucations for elapsed time and time delta, which depend on values from the immediately prior and immediately succeeding records, are calculated correctly.\r\n",
        "* Calculating `pqet_current`, the time it took on average to answer the questions in the current `task_container_id_q`.\r\n",
        "* Calculating `ts_delta`, the elapsed time between the last `task_container_id_q` and the current one.\r\n",
        "* Creating `folds` table, in which users are assigned to one of 20 folds.\r\n",
        "* Creating a `tags_array` field in the questions table, that returns an array of six elements populated with the tags assigned to each questions, padded with zeros if there are less than six."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X2ynWZqPFnSj"
      },
      "source": [
        "def done_cb(future):\r\n",
        "    seconds = (future.ended - future.started).total_seconds()\r\n",
        "    print(f'Job {future.job_id} finished in {seconds} seconds.')\r\n",
        "\r\n",
        "def run_query(query, job_id_prefix=None, wait=True,\r\n",
        "                use_query_cache=True):\r\n",
        "\r\n",
        "    job_config = QueryJobConfig(\r\n",
        "        use_query_cache=use_query_cache)\r\n",
        "\r\n",
        "    query_job = bq_client.query(query, job_id_prefix=job_id_prefix,\r\n",
        "                                        job_config=job_config)\r\n",
        "    print(f'Job {query_job.job_id} started.')\r\n",
        "    query_job.add_done_callback(done_cb)\r\n",
        "    if wait:\r\n",
        "        query_job.result()\r\n",
        "    \r\n",
        "    return query_job"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Qo21D1ITicH"
      },
      "source": [
        "def get_df_query_bqs(query, dtypes=None, fillna=None):\r\n",
        "    qj = bq_client.query(query)\r\n",
        "    df = qj.to_dataframe(create_bqstorage_client=True, progress_bar_type='tqdm_notebook')\r\n",
        "    if fillna is not None:\r\n",
        "        df = df.fillna(fillna)\r\n",
        "    try:\r\n",
        "        df = df.astype({c: dtypes.get(c, 'int32') for c in df.columns})    \r\n",
        "    except:\r\n",
        "        print('dtypes not applied.')\r\n",
        "    finally:    \r\n",
        "        return df"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N71-o9uQFSzV"
      },
      "source": [
        "### Replace Missing Values"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rnBL1HXxFWKX"
      },
      "source": [
        "def update_missing_values(table_id='train', column_id=None, value=None):\r\n",
        "    return f\"\"\"\r\n",
        "        UPDATE {BQ_DATASET}.{table_id}\r\n",
        "        SET {column_id} = {value}\r\n",
        "        WHERE {column_id} is NULL;\r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e0qBG2XrGIMB"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    qj = run_query(*update_missing_values('train', 'prior_question_elapsed_time', '0'))\r\n",
        "    qj = run_query(*update_missing_values('train', 'prior_question_had_explanation', 'false'))\r\n",
        "    qj = run_query(*update_missing_values('questions', 'tags', '\"188\"'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "elZXRogqL-pr"
      },
      "source": [
        "### Recalculate Task Container Ids for Questions Only"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Li0UdfY2MeOm"
      },
      "source": [
        "def update_task_container_id(table_id='train',\r\n",
        "                                column_id='task_container_id',\r\n",
        "                                excl_lectures=True):\r\n",
        "    excl_lec = 'WHERE content_type_id = 0' if excl_lectures else ''\r\n",
        "    \r\n",
        "    return f\"\"\"\r\n",
        "        UPDATE {BQ_DATASET}.{table_id} t\r\n",
        "        SET {column_id} = target.calc\r\n",
        "        FROM (\r\n",
        "            SELECT row_id, DENSE_RANK()\r\n",
        "            OVER (\r\n",
        "                PARTITION BY user_id\r\n",
        "                ORDER BY timestamp\r\n",
        "            ) calc\r\n",
        "            FROM {BQ_DATASET}.{table_id}\r\n",
        "            {excl_lec}\r\n",
        "        ) target\r\n",
        "        WHERE target.row_id = t.row_id\r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FGFisFdpMtGy"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    q  = update_task_container_id(table_id='train',\r\n",
        "                                column_id='task_container_id_q ',\r\n",
        "                                excl_lectures=True)\r\n",
        "    qj = run_query(*q)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2HblfPhCG618"
      },
      "source": [
        "### Calculate Current Question Elapsed Time and Timestamp Delta"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "29-ajMgjHEUl"
      },
      "source": [
        "def update_pqet_current(table_id='train'):\r\n",
        "    return f\"\"\"\r\n",
        "        UPDATE {BQ_DATASET}.{table_id} t\r\n",
        "        SET t.pqet_current = CAST(p.pqet_current AS INT64)\r\n",
        "        FROM (\r\n",
        "            SELECT\r\n",
        "            row_id, LAST_VALUE(prior_question_elapsed_time) OVER (\r\n",
        "                PARTITION BY user_id ORDER BY task_container_id_q\r\n",
        "                RANGE BETWEEN 1 FOLLOWING AND 1 FOLLOWING) pqet_current\r\n",
        "            FROM {BQ_DATASET}.train            \r\n",
        "            WHERE content_type_id = 0\r\n",
        "        ) p\r\n",
        "        WHERE t.row_id = p.row_id;\r\n",
        "        \r\n",
        "        UPDATE {BQ_DATASET}.{table_id}\r\n",
        "        SET pqet_current = 0\r\n",
        "        WHERE pqet_current IS NULL;\r\n",
        "        \r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "juf9vDrzIF2W"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    qj = run_query(*update_pqet_current())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E9LnKsjVLgRk"
      },
      "source": [
        "def update_ts_delta(table_id='train'):\r\n",
        "    return f\"\"\"\r\n",
        "        UPDATE {BQ_DATASET}.{table_id} t\r\n",
        "        SET t.ts_delta = timestamp - p.ts_prior\r\n",
        "        FROM (\r\n",
        "            SELECT\r\n",
        "            row_id, LAST_VALUE(timestamp) OVER (\r\n",
        "                PARTITION BY user_id ORDER BY task_container_id_q\r\n",
        "                RANGE BETWEEN 1 PRECEDING AND 1 PRECEDING) ts_prior\r\n",
        "            FROM {BQ_DATASET}.train            \r\n",
        "            WHERE content_type_id = 0\r\n",
        "        ) p\r\n",
        "        WHERE t.row_id = p.row_id;\r\n",
        "        \r\n",
        "        UPDATE {BQ_DATASET}.{table_id}\r\n",
        "        SET ts_delta = 0\r\n",
        "        WHERE ts_delta IS NULL;\r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0-CEUJsoL1dC"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    qj = run_query(*update_ts_delta())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "99qz0H8Xb3i1"
      },
      "source": [
        "### Create Folds Table\r\n",
        "Assign users randomly to one of 20 folds. Store total records to facilitate filtering based on record count."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N7UnAHF8cesC"
      },
      "source": [
        "def create_table_folds(table_id='folds', n_folds=20):\r\n",
        "    return f\"\"\"\r\n",
        "        DECLARE f INT64;\r\n",
        "\r\n",
        "        CREATE OR REPLACE TABLE {BQ_DATASET}.{table_id} (\r\n",
        "            user_id INT64,\r\n",
        "            fold INT64,\r\n",
        "            record_count INT64\r\n",
        "        );\r\n",
        "\r\n",
        "        INSERT {BQ_DATASET}.{table_id} (user_id, fold, record_count)\r\n",
        "        SELECT f.user_id, CAST(FLOOR(RAND() * {n_folds}) AS INT64) fold, f.record_count\r\n",
        "        FROM (\r\n",
        "        SELECT user_id,\r\n",
        "            COUNT(row_id) record_count\r\n",
        "        FROM {BQ_DATASET}.train\r\n",
        "        WHERE content_type_id = 0\r\n",
        "        GROUP BY user_id\r\n",
        "        ) f\r\n",
        "        ORDER BY user_id;\r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SVPio880dPSe"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    qj = run_query(*create_table_folds())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "14dQwOnzdolg"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_folds = get_df_query_bqs(f\"\"\"\r\n",
        "        SELECT *\r\n",
        "        FROM {BQ_DATASET}.folds\r\n",
        "    \"\"\",\r\n",
        "    dtypes=dtypes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R9Y9Xhwee6f7"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_folds.groupby('fold').count().user_id.plot(kind='bar', title='Count of Users by Fold');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qciROEyIoowx"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_folds.groupby('fold').mean().record_count.plot(kind='bar', title='Average Records per User by Fold');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q5zS5bWenJaj"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_fold_ac = get_df_query_bqs(f\"\"\"\r\n",
        "        SELECT fold, SUM(answered_correctly) ac_sum, COUNT(answered_correctly) rec_count\r\n",
        "        FROM {BQ_DATASET}.train\r\n",
        "        JOIN {BQ_DATASET}.folds\r\n",
        "        ON train.user_id = folds.user_id\r\n",
        "        GROUP BY fold\r\n",
        "    \"\"\",\r\n",
        "    dtypes=dtypes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3TGelsfEn7xf"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_fold_ac.rec_count.plot(kind='bar', title='Count of Records by Fold');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a1kcjfOkoGU_"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    (df_fold_ac.ac_sum / df_fold_ac.rec_count).plot(kind='bar', title='Percent Answered Correctly by Fold');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nqg9tqEnOa7l"
      },
      "source": [
        "### Create Tags Array on Questions Table\r\n",
        "We need the tags as an array later when we create TFRecords. We also increment by one and pad with zeros to a fixed length of 6 so that they can be concatentated as a feature for modeling."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wr4peSSWPHiW"
      },
      "source": [
        "def update_tags_array(table_id='questions', column_id='tags_array'):\r\n",
        "    \r\n",
        "    return f\"\"\"\r\n",
        "        UPDATE {BQ_DATASET}.{table_id} q\r\n",
        "        SET {column_id} = tp.tags_fixed_len\r\n",
        "        FROM (\r\n",
        "            WITH tags_padded AS (\r\n",
        "                WITH tags_table AS (SELECT question_id, tags FROM {BQ_DATASET}.{table_id})\r\n",
        "                SELECT question_id, ARRAY_CONCAT(ARRAY_AGG(CAST(tag AS INT64) + 1), [0,0,0,0,0]) tags_array\r\n",
        "                FROM tags_table, UNNEST(SPLIT(tags, ' ')) as tag\r\n",
        "                GROUP BY question_id\r\n",
        "            )\r\n",
        "            SELECT question_id,\r\n",
        "                ARRAY(SELECT x FROM UNNEST(tags_array) AS x WITH OFFSET off WHERE off < 6 ORDER BY off) tags_fixed_len\r\n",
        "            FROM tags_padded\r\n",
        "        ) tp\r\n",
        "        WHERE tp.question_id = q.question_id\r\n",
        "    \"\"\", sys._getframe().f_code.co_name + '_'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "exf_kIXuRagG"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    qj = run_query(*update_tags_array())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pB2YEqhISnRz"
      },
      "source": [
        "if PERFORM_FEATURE_ENGINEERING:\r\n",
        "    df_q = get_df_query_bqs('select * from my_data.questions', dtypes=dtypes)\r\n",
        "    print(df_q.head())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BrxKtvosXTfw"
      },
      "source": [
        "## Feature Engineering Tests\r\n",
        "* Features come back out of Biq Query with the same values they went in with\r\n",
        "* `ts_delta` is equal to difference between timestamps on consecutive records\r\n",
        "* `pqet_current` is equal to `prior_question_elapsed_time` from next record\r\n",
        "* visually inspect distributions of `ts_delta` and `pqet_current`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ynTomTnKY7F3"
      },
      "source": [
        "### Load Sample from train.csv"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c9aNhgV7Yfzw"
      },
      "source": [
        "if TEST_FEATURE_ENGNEERING:\r\n",
        "    df_train_samp = pd.read_csv('train.csv', nrows=100000)\r\n",
        "    df_train_samp.prior_question_had_explanation = df_train_samp.prior_question_had_explanation.fillna(False).astype(bool)\r\n",
        "    df_train_samp.prior_question_elapsed_time = df_train_samp.prior_question_elapsed_time.fillna(0)\r\n",
        "    user_ids_samp = df_train_samp.user_id.unique()[:-1]\r\n",
        "    print(len(user_ids_samp))\r\n",
        "    df_train_samp = df_train_samp[df_train_samp.user_id.isin(user_ids_samp) & (df_train_samp.content_type_id == 0)].reset_index(drop=True)\r\n",
        "    print(len(df_train_samp))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "elJgSDpmY0q0"
      },
      "source": [
        "### Pull sample of corresponding user_ids from BigQuery"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GvQF1yIzZPHE"
      },
      "source": [
        "if TEST_FEATURE_ENGNEERING:\r\n",
        "    df_bq_samp = get_df_query_bqs(f\"\"\"\r\n",
        "        SELECT *\r\n",
        "        FROM {BQ_DATASET}.train\r\n",
        "        WHERE user_id IN ({(',').join(map(str, user_ids_samp))})\r\n",
        "        AND content_type_id = 0\r\n",
        "        ORDER BY user_id, timestamp, row_id\r\n",
        "    \"\"\",\r\n",
        "    dtypes=None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gP3mDb-phsSt"
      },
      "source": [
        "### Tests"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TukjFI4YatpD"
      },
      "source": [
        "if TEST_FEATURE_ENGNEERING:\r\n",
        "    # values in columns are the same between train.csv and bq\r\n",
        "    for c in df_train_samp.columns:\r\n",
        "        assert all(df_train_samp[c] == df_bq_samp[c]), f'{c} is not the same'\r\n",
        "\r\n",
        "    # pqet_current pulls prior_question_elapsed_time back one task_container_id for each user\r\n",
        "    df_bq_samp_tst = df_bq_samp[['user_id', 'task_container_id_q', 'prior_question_elapsed_time', 'pqet_current']].groupby(['user_id', 'task_container_id_q']).max()\r\n",
        "\r\n",
        "    for user_id in user_ids_samp:\r\n",
        "        assert all(df_bq_samp_tst.loc[user_id].pqet_current.shift(1).iloc[1:] == df_bq_samp_tst.loc[user_id].prior_question_elapsed_time.iloc[1:])\r\n",
        "\r\n",
        "    # ts_delta equal to timestamp from current task_container_id_q minus timestamp from prior task_container_id_q\r\n",
        "    df_bq_samp_tst = df_bq_samp[['user_id', 'task_container_id_q', 'timestamp', 'ts_delta']].groupby(['user_id', 'task_container_id_q']).max()\r\n",
        "\r\n",
        "    for user_id in user_ids_samp:\r\n",
        "        assert all((df_bq_samp_tst.loc[user_id].timestamp - df_bq_samp_tst.loc[user_id].timestamp.shift(1)).iloc[1:] == df_bq_samp_tst.loc[user_id].ts_delta.iloc[1:])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jZhQStC4bIa_"
      },
      "source": [
        "if TEST_FEATURE_ENGNEERING:\r\n",
        "    df_bq_samp.pqet_current.hist();"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4WrlU_qZbisC"
      },
      "source": [
        "if TEST_FEATURE_ENGNEERING:\r\n",
        "    df_bq_samp.ts_delta.hist();"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qR33_34dpD5R"
      },
      "source": [
        "## Create TFRecords"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjZunnJNqG7U"
      },
      "source": [
        "We are going to create a set of TFRecords with one user per record and one fold per file. We are going to include the following columns as features:\r\n",
        "* `user_id` - this won't get used as a feature, but is included to able to tie back to original data\r\n",
        "* `content_id` - incremented by one to reserve 0 for padding character\r\n",
        "* `answered_correctly` - incremented by one to reserve 0 for padding character\r\n",
        "* `part`\r\n",
        "* `pqet_curret`\r\n",
        "* `ts_delta`\r\n",
        "* `tags` - already incremented by one with zeros as padding\r\n",
        "* `task_container_id` - excluding lectures and already indexed to one\r\n",
        "* `timestamp`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e3IhNw5C3epG"
      },
      "source": [
        "def _int64_feature(value):\r\n",
        "        \r\n",
        "    if type(value) != type(list()):\r\n",
        "        value = [value]\r\n",
        "\r\n",
        "    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jBv7Qtjb3gr-"
      },
      "source": [
        "def serialize_example(user_id, features):\r\n",
        "    \r\n",
        "    feature_names = ['content_id', 'answered_correctly', 'part', 'pqet_current', 'ts_delta', 'tags',\r\n",
        "                     'task_container_id', 'timestamp']\r\n",
        "    \r\n",
        "    feature = {'user_id': _int64_feature(user_id)}\r\n",
        "    \r\n",
        "    for i, n in enumerate(feature_names):\r\n",
        "        feature[n] = _int64_feature(features[i])\r\n",
        "\r\n",
        "    return tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8sG6HK_p3imX"
      },
      "source": [
        "def parse_example(example):\r\n",
        "    \r\n",
        "    feature_names = {'content_id': tf.int32, 'answered_correctly': tf.int32, 'part': tf.int32,\r\n",
        "                     'pqet_current': tf.int32, 'ts_delta': tf.int64, 'tags': tf.int32,\r\n",
        "                     'task_container_id': tf.int32, 'timestamp': tf.int64}\r\n",
        "      \r\n",
        "    features = {'user_id': tf.io.FixedLenFeature([1], tf.int64)}\r\n",
        "    \r\n",
        "    for k, v in feature_names.items():\r\n",
        "        features[k] = tf.io.VarLenFeature(tf.int64)\r\n",
        "\r\n",
        "    example = tf.io.parse_single_example(example, features)\r\n",
        "\r\n",
        "    for k, v in feature_names.items():\r\n",
        "        example[k] = tf.cast(example[k].values, v)\r\n",
        "        \r\n",
        "    example['tags'] = tf.reshape(example['tags'], (tf.size(example['answered_correctly']), 6))\r\n",
        "\r\n",
        "    return example"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nh-44zO158H8"
      },
      "source": [
        "def get_ds_tfrec_raw(folds=[0]):\r\n",
        "    file_pat = 'gs://{BUCKET}/tfrec/{f:02d}-*.tfrec'\r\n",
        "    file_pats = [file_pat.format(BUCKET=BUCKET, f=f) for f in folds]\r\n",
        "    options = tf.data.Options()\r\n",
        "\r\n",
        "    ds = (tf.data.Dataset.list_files(file_pats)\r\n",
        "          .with_options(options)\r\n",
        "          .interleave(tf.data.TFRecordDataset, num_parallel_calls=AUTO)\r\n",
        "          .map(parse_example, num_parallel_calls=AUTO)\r\n",
        "         )\r\n",
        "    \r\n",
        "    return ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NEfOS7m2pKZp"
      },
      "source": [
        "def get_df_tfrec(folds):\r\n",
        "    df_tfrec = get_df_query_bqs(f\"\"\"\r\n",
        "        SELECT fold, train.user_id, content_id + 1 content_id,\r\n",
        "            answered_correctly + 1 answered_correctly, part, pqet_current, ts_delta,\r\n",
        "            tags_array tags, task_container_id_q task_container_id, timestamp\r\n",
        "        FROM {BQ_DATASET}.train\r\n",
        "        JOIN {BQ_DATASET}.folds\r\n",
        "        ON train.user_id = folds.user_id\r\n",
        "        JOIN {BQ_DATASET}.questions\r\n",
        "        ON train.content_id = questions.question_id\r\n",
        "        WHERE fold IN ({(', ').join(map(str, folds))})\r\n",
        "        AND content_type_id = 0\r\n",
        "        ORDER BY user_id, timestamp, row_id\r\n",
        "    \"\"\",\r\n",
        "    dtypes=None)\r\n",
        "\r\n",
        "    return df_tfrec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HoYPoSKR3V5-"
      },
      "source": [
        "def write_tfrecords(folds):\r\n",
        "    \r\n",
        "    df_tfrec = get_df_tfrec(folds)\r\n",
        "    \r\n",
        "    for f in folds:\r\n",
        "        groups_dict = (df_tfrec[df_tfrec.fold == f]\r\n",
        "                       .groupby('user_id')\r\n",
        "                       .apply(lambda r: (list(r['content_id'].values),\r\n",
        "                                         list(r['answered_correctly'].values),\r\n",
        "                                         list(r['part'].values),\r\n",
        "                                         list(r['pqet_current'].values.astype(np.int64)),\r\n",
        "                                         list(r['ts_delta'].values.astype(np.int64)),\r\n",
        "                                         list(np.concatenate(r['tags'].values)),\r\n",
        "                                         list(r['task_container_id'].values.astype(np.int64)),\r\n",
        "                                         list(r['timestamp'].values.astype(np.int64)),\r\n",
        "                                         ))).to_dict()        \r\n",
        "        \r\n",
        "        out_path = f'gs://{BUCKET}/tfrec'\r\n",
        "        filename = f'{f:02d}-{len(groups_dict.keys())}.tfrec'\r\n",
        "        record_file = f'{out_path}/{filename}'\r\n",
        "\r\n",
        "        with tf.io.TFRecordWriter(record_file) as writer:\r\n",
        "            for user_id, features in tqdm(groups_dict.items(), desc=f'Fold {f:02d}'):\r\n",
        "                writer.write(serialize_example(user_id, features))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "48jjQ7L9g1_M"
      },
      "source": [
        "## Write TFRecords\r\n",
        "\r\n",
        "* Process in chunks to avoid running out of memory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mpw-Nb7Dg8xL"
      },
      "source": [
        "if CREATE_TFRECORDS:\r\n",
        "    fold_splits = np.array_split(np.arange(20), 10)\r\n",
        "    for folds in tqdm(fold_splits):\r\n",
        "        write_tfrecords(folds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Un9OMMVC4QQQ"
      },
      "source": [
        "## Test TFRecords\r\n",
        "\r\n",
        "* Same number of users and records as in `df_folds`\r\n",
        "* Values in tfrecords are the same as in original data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eIGmrsU28TJ4"
      },
      "source": [
        "def test_tfrecord_folds(folds_test, n_sample=100):\r\n",
        "    pbar = tqdm(total=n_sample)\r\n",
        "    ds = get_ds_tfrec_raw(folds_test)\r\n",
        "    df = get_df_tfrec(folds_test)\r\n",
        "\r\n",
        "    for b in ds.shuffle(10000).take(n_sample):\r\n",
        "        try:\r\n",
        "            for c in [c for c in df.columns if c not in  ['tags', 'fold', 'user_id']]:\r\n",
        "                try:\r\n",
        "                    assert all(df[df.user_id == b['user_id'].numpy()[0]][c] == b[c].numpy())\r\n",
        "                except:\r\n",
        "                    print(f\"Error for user {b['user_id'].numpy()[0]}\")\r\n",
        "            user_tags = np.concatenate(df[df.user_id == b['user_id'].numpy()[0]].tags.values)\r\n",
        "            assert all(user_tags == (b['tags'].numpy().flatten()))\r\n",
        "        except:\r\n",
        "            print(f\"Error for user {b['user_id'].numpy()[0]}\")\r\n",
        "        finally:\r\n",
        "            pbar.update()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ByaOx2_23BJJ"
      },
      "source": [
        "if TEST_TFRECORDS:\r\n",
        "    folds_test = list(range(20))\r\n",
        "    ds = get_ds_tfrec_raw(folds=folds_test)\r\n",
        "\r\n",
        "    df_folds = get_df_query_bqs(f\"\"\"\r\n",
        "        SELECT *\r\n",
        "        FROM {BQ_DATASET}.folds\r\n",
        "    \"\"\",\r\n",
        "    dtypes=dtypes)\r\n",
        "\r\n",
        "    user_ids = []\r\n",
        "    count = 0\r\n",
        "    for b in ds:\r\n",
        "        user_ids.append(b['user_id'].numpy()[0])\r\n",
        "        count += len(b['content_id'].numpy())\r\n",
        "\r\n",
        "    assert len(set(user_ids)) == len(df_folds)\r\n",
        "    assert df_folds.record_count.sum() == count\r\n",
        "\r\n",
        "    test_tfrecord_folds([10])\r\n",
        "\r\n",
        "    b = next(iter(ds))\r\n",
        "    print(b)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OMfgCXo159d2"
      },
      "source": [
        "## Dataset Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TcQ-dd1b6EKN"
      },
      "source": [
        "@gin.configurable\r\n",
        "def get_ds_tfrec(folds=None, max_len=None, min_len=None):\r\n",
        "    file_pat = 'gs://{BUCKET}/tfrec/{f:02d}-*.tfrec'\r\n",
        "    file_pats = [file_pat.format(BUCKET=BUCKET, f=f) for f in folds]\r\n",
        "    options = tf.data.Options()\r\n",
        "\r\n",
        "    ds = (tf.data.Dataset.list_files(file_pats, shuffle=True)\r\n",
        "          .with_options(options)\r\n",
        "          .interleave(tf.data.TFRecordDataset, num_parallel_calls=AUTO)\r\n",
        "          .shuffle(10000)\r\n",
        "          .map(parse_example, num_parallel_calls=AUTO)\r\n",
        "          .filter(partial(filter_min_len, min_len=min_len))\r\n",
        "          .map(example_to_tuple, num_parallel_calls=AUTO)\r\n",
        "          .map(partial(trunc_seq, max_len=max_len), num_parallel_calls=AUTO)\r\n",
        "          .map(con_to_cat, num_parallel_calls=AUTO)\r\n",
        "         )\r\n",
        "\r\n",
        "    ds = ds.repeat().prefetch(AUTO)\r\n",
        "    \r\n",
        "    def gen(generator=None):\r\n",
        "        del generator\r\n",
        "        for example in fastmath.dataset_as_numpy(ds):\r\n",
        "            yield example\r\n",
        "    \r\n",
        "    return gen"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l20hHtYADfGp"
      },
      "source": [
        "def filter_min_len(e, min_len):\r\n",
        "    return tf.size(e['content_id']) >= min_len"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2TcKGF8t70Gl"
      },
      "source": [
        "def example_to_tuple(example):\r\n",
        "    return (example['content_id'], example['part'], example['tags'], example['task_container_id'],\r\n",
        "            example['answered_correctly'], example['pqet_current'], example['ts_delta'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5PjKM5Q57uS5"
      },
      "source": [
        "def trunc_seq(*b, max_len=None):\r\n",
        "    \"\"\"Returns a sequence drawn randomly from available tokens with a max length\r\n",
        "        of max_len.\r\n",
        "    \"\"\"\r\n",
        "    \r\n",
        "    max_len = tf.constant(max_len)\r\n",
        "    seq_len = tf.size(b[0])\r\n",
        "    seq_end_min = tf.minimum(seq_len - 1, max_len)\r\n",
        "    seq_end = tf.maximum(max_len, tf.random.uniform((), seq_end_min, seq_len, dtype=tf.int32))\r\n",
        "    \r\n",
        "    def get_seq(m):\r\n",
        "        return m[seq_end-max_len:seq_end]\r\n",
        "    \r\n",
        "    return tuple(map(get_seq, b))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9aC7H12W7mRs"
      },
      "source": [
        "# SAINT+ Elapsed Time = prior_question_elapsed_time and Lag Time = time_stamp_1 - timestamp_0\r\n",
        "# Elapsed Time categorical - capped at 300 seconds, discrete value for each second\r\n",
        "# Lag Time - discretized to minutes 0, 1, 2, 3, 4, 5, 10, 20, 30 ... 1440. 150 discrete values.\r\n",
        "\r\n",
        "ts_delta_lookup = tf.concat([tf.range(6, dtype=tf.int32), tf.repeat(5, 5)], axis=0)\r\n",
        "\r\n",
        "cat = 10\r\n",
        "while cat < 1440:\r\n",
        "    ts_delta_lookup = tf.concat([ts_delta_lookup, tf.repeat(cat, 10)], axis=0)\r\n",
        "    cat += 10\r\n",
        "    \r\n",
        "ts_delta_lookup = tf.concat([ts_delta_lookup, [1440]], axis=0)\r\n",
        "\r\n",
        "def con_to_cat(*b):\r\n",
        "    \r\n",
        "    def pqet_cat(e, vocab_size=None, val_min=None, val_max=None):\r\n",
        "        e = tf.clip_by_value(e, val_min, val_max)\r\n",
        "        val_range = val_max - val_min\r\n",
        "        e = tf.cast((e - val_min) * (vocab_size - 1) / val_range, tf.int32)\r\n",
        "        return e\r\n",
        "    \r\n",
        "    def ts_delta_cat(e):\r\n",
        "        val_max = tf.cast(tf.reduce_max(ts_delta_lookup) * 60000, tf.float64)\r\n",
        "        e = tf.clip_by_value(tf.cast(e, tf.float64), 0, val_max)\r\n",
        "        e = tf.cast(e / 60000, tf.int32)\r\n",
        "        e = tf.gather(ts_delta_lookup, e)\r\n",
        "        return e\r\n",
        "    \r\n",
        "    pqet = pqet_cat(b[-2], vocab_size=300, val_min=0, val_max=300000)\r\n",
        "    ts_delta = ts_delta_cat(b[-1])\r\n",
        "    \r\n",
        "    return tuple((*b[:-2], pqet, ts_delta))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cWKX9-WHNIdJ"
      },
      "source": [
        "## Metrics Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BUC43hA69diL"
      },
      "source": [
        "def RocAucScore(num_thresholds=100, pos_label=2):\r\n",
        "    def f(y_score, y_true, weight):        \r\n",
        "        weight = tnp.expand_dims(tnp.ravel(weight), -1)\r\n",
        "        \r\n",
        "        softmax=tl.Softmax(axis=-1)\r\n",
        "        y_score = tnp.ravel(softmax(y_score)[:, :, -1])\r\n",
        "        y_score = tnp.expand_dims(y_score, -1)\r\n",
        "        y_true = tnp.expand_dims(tnp.ravel(y_true) == pos_label, -1).astype(tnp.float32)\r\n",
        "        \r\n",
        "        thresholds = tnp.expand_dims(tnp.linspace(1, 0, num_thresholds), 0)\r\n",
        "        \r\n",
        "        threshold_counts = y_score > thresholds\r\n",
        "        \r\n",
        "        tps = tnp.logical_and(threshold_counts, y_true)\r\n",
        "        fps = tnp.logical_and(threshold_counts, tnp.logical_not(y_true))\r\n",
        "        \r\n",
        "        tps = tnp.sum(tps * weight, axis=0)\r\n",
        "        fps = tnp.sum(fps * weight, axis=0)\r\n",
        "        \r\n",
        "        tpr = tps / tps[-1]\r\n",
        "        fpr = fps / fps[-1]\r\n",
        "        \r\n",
        "        return tnp.trapz(tpr, fpr)\r\n",
        "    \r\n",
        "    return tl.Fn('RocAucScore', f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LtbnwInF9fMf"
      },
      "source": [
        "metrics = {\r\n",
        "    'loss': tl.WeightedCategoryCrossEntropy(),\r\n",
        "    'accuracy': tl.WeightedCategoryAccuracy(),\r\n",
        "    'sequence_accuracy': tl.MaskedSequenceAccuracy(),\r\n",
        "    'auc_all': RocAucScore(),\r\n",
        "    'weights_per_batch_per_core': tl.Serial(tl.Drop(), tl.Drop(), tl.Sum())\r\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WA_-d9V_NN-h"
      },
      "source": [
        "## Model Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QK7Kf4AM9u9P"
      },
      "source": [
        "@gin.configurable\r\n",
        "@tl.assert_shape('bl->b1ll')\r\n",
        "def PaddingFutureMask(pad=0, block_self=False, tid=True, pad_end=False):\r\n",
        "    def f(x):\r\n",
        "        mask_pad = tnp.logical_not(tnp.equal(x, 0))[:, tnp.newaxis, tnp.newaxis, :]\r\n",
        "        \r\n",
        "        x_new = x\r\n",
        "        if pad_end:\r\n",
        "            x_new = tnp.where(tnp.equal(x, 0), tnp.max(x), x)\r\n",
        "        \r\n",
        "        if tid:\r\n",
        "            mask_future = x_new[:, :, tnp.newaxis] >= x_new[:, tnp.newaxis, :] + block_self\r\n",
        "            mask_future = mask_future[:, tnp.newaxis, :, :]\r\n",
        "        else:\r\n",
        "            mask_future = tnp.arange(x.shape[-1])[tnp.newaxis, tnp.newaxis, :, tnp.newaxis] \\\r\n",
        "                >= tnp.arange(x.shape[-1])[tnp.newaxis, :]\r\n",
        "        \r\n",
        "        return tnp.logical_and(mask_future, mask_pad)\r\n",
        "        \r\n",
        "    return tl.Fn(f'PaddingFutureMask({pad})', f)\r\n",
        "\r\n",
        "\r\n",
        "# the only thing different here is the shape assertions to accomodate the change\r\n",
        "# in mask shape from b11l to b1ll\r\n",
        "\r\n",
        "@tl.assert_shape('bld,b1ll->bld,b1ll')\r\n",
        "@gin.configurable\r\n",
        "def KTAttention(d_feature, n_heads=1, dropout=0.0, mode='train'):\r\n",
        "    return tl.Serial(\r\n",
        "        tl.Select([0, 0, 0]),\r\n",
        "        tl.AttentionQKV(\r\n",
        "            d_feature, n_heads=n_heads, dropout=dropout, mode=mode),\r\n",
        "    )\r\n",
        "\r\n",
        "def my_add_loss_weights(generator, id_to_mask=None):\r\n",
        "    for example in generator:\r\n",
        "        weights = (example[0] != id_to_mask).astype(tnp.float32)\r\n",
        "        yield (*example, weights)\r\n",
        "\r\n",
        "@gin.configurable\r\n",
        "def KTAddLossWeights(id_to_mask=0):  # pylint: disable=invalid-name\r\n",
        "    return lambda g: my_add_loss_weights(g, id_to_mask=id_to_mask)\r\n",
        "\r\n",
        "def trim_tags(generator):\r\n",
        "    for example in generator:\r\n",
        "        # content_id, part, tags, tid, ac, pqet, ts_delta\r\n",
        "        yield (example[0], example[1], example[2][:, :, :6], example[3], example[4], example[5], example[6])\r\n",
        "\r\n",
        "@gin.configurable\r\n",
        "def TrimTags():\r\n",
        "    return lambda g: trim_tags(g)\r\n",
        "\r\n",
        "@gin.configurable\r\n",
        "def KTPositionalEncoder(max_position=10000.0, d_model=512, tid=False):   \r\n",
        "    \"\"\"This is set up to perform standard positional encoding based on the\r\n",
        "    position in the sequence, but also to calculate position based on the\r\n",
        "    id of the task container to which the question belongs.\r\n",
        "    \"\"\"\r\n",
        "    def f(inputs):\r\n",
        "        # whether or not to use task_container_id or seq position\r\n",
        "        if tid:\r\n",
        "            position = tnp.expand_dims(inputs.astype(tnp.float32), -1)\r\n",
        "        else:\r\n",
        "            position = tnp.arange(inputs.shape[1])\r\n",
        "            \r\n",
        "            position = position.astype(tnp.float32)[tnp.newaxis, :, tnp.newaxis]\r\n",
        "\r\n",
        "        i = tnp.expand_dims(tnp.arange(d_model, dtype=tnp.float32), 0)\r\n",
        "\r\n",
        "        angles = 1 / tnp.power(max_position, (2 * (i // 2)) /\r\n",
        "                               tnp.array(d_model, dtype=tnp.float32))\r\n",
        "\r\n",
        "        angle_rads = position * angles\r\n",
        "\r\n",
        "        # apply sin to even index in the array\r\n",
        "        sines = tnp.sin(angle_rads[:, :, 0::2])\r\n",
        "        # apply cos to odd index in the array\r\n",
        "        cosines = tnp.cos(angle_rads[:, :, 1::2])\r\n",
        "\r\n",
        "        pos_encoding = tnp.concatenate([sines, cosines], axis=-1)\r\n",
        "\r\n",
        "        return pos_encoding\r\n",
        "\r\n",
        "    return tl.Fn('KTPositionalEncoder', f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pmo6yeiXkBAQ"
      },
      "source": [
        "@gin.configurable\r\n",
        "def KTTransformer(d_model,\r\n",
        "                  d_input,\r\n",
        "                  d_part,\r\n",
        "                  d_tags,\r\n",
        "                  d_out,\r\n",
        "                  d_pqet,\r\n",
        "                  d_ts_delta,\r\n",
        "                  d_tid,\r\n",
        "                  embed_concat=False,\r\n",
        "                  d_ff=2048,\r\n",
        "                  n_encoder_layers=6,\r\n",
        "                  n_decoder_layers=6,\r\n",
        "                  n_heads=8,\r\n",
        "                  max_len=2048,\r\n",
        "                  dropout=0.1,\r\n",
        "                  dropout_shared_axes=None,\r\n",
        "                  mode='train',\r\n",
        "                  ff_activation=tl.Relu):\r\n",
        "        \r\n",
        "    def Embedder(vocab_size, d_embed):  # tokens --> vectors\r\n",
        "        return [\r\n",
        "            tl.Embedding(vocab_size, d_embed),\r\n",
        "            tl.Dropout(\r\n",
        "                rate=dropout, shared_axes=dropout_shared_axes, mode=mode),\r\n",
        "        ]\r\n",
        "\r\n",
        "    # Encoder Embeddings\r\n",
        "    in_embedder = Embedder(*d_input)\r\n",
        "    part_embedder = Embedder(*d_part)\r\n",
        "    # Keeps the tags in the data batch tuple, but drops it if it\r\n",
        "    # isn't included in the embeddings.\r\n",
        "    if d_tags is not None:\r\n",
        "        tags_embedder = tl.Serial(Embedder(*d_tags), tl.Sum(axis=-2))\r\n",
        "    else:\r\n",
        "        tags_embedder = tl.Drop()\r\n",
        "    in_pos_encoder = KTPositionalEncoder(*d_tid)\r\n",
        "\r\n",
        "    # Decoder Embeddings\r\n",
        "    out_embedder = Embedder(*d_out)\r\n",
        "    pqet_embedder = Embedder(*d_pqet)\r\n",
        "    ts_delta_embedder = Embedder(*d_ts_delta)\r\n",
        "    out_pos_encoder = KTPositionalEncoder(*d_tid)\r\n",
        "\r\n",
        "    encoder_mode = 'eval' if mode == 'predict' else mode\r\n",
        "\r\n",
        "    in_encoder = [tl.Parallel(in_embedder, part_embedder, tags_embedder, in_pos_encoder)]\r\n",
        "    out_encoder = [tl.Parallel(out_embedder, pqet_embedder, ts_delta_embedder, out_pos_encoder)]\r\n",
        "    \r\n",
        "    if embed_concat:\r\n",
        "        if d_tags is not None:\r\n",
        "            in_encoder += [tl.Concatenate(n_items=3), tl.Add()]\r\n",
        "        else:\r\n",
        "            in_encoder += [tl.Concatenate(n_items=2), tl.Add()]\r\n",
        "        out_encoder += [tl.Concatenate(n_items=3), tl.Add()]\r\n",
        "    else:\r\n",
        "        if d_tags is not None:\r\n",
        "            in_encoder += [tl.Add(), tl.Add(), tl.Add()]\r\n",
        "        else:\r\n",
        "            in_encoder += [tl.Add(), tl.Add()]\r\n",
        "        out_encoder += [tl.Add(), tl.Add(), tl.Add()]\r\n",
        "\r\n",
        "    encoder_blocks = [\r\n",
        "        _KTEncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,\r\n",
        "                      mode, ff_activation)\r\n",
        "        for i in range(n_encoder_layers)]\r\n",
        "\r\n",
        "    encoder = tl.Serial(\r\n",
        "        in_encoder,\r\n",
        "        encoder_blocks,\r\n",
        "        tl.LayerNorm()\r\n",
        "    )\r\n",
        "\r\n",
        "    encoder_decoder_blocks = [\r\n",
        "        _KTEncoderDecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,\r\n",
        "                             mode, ff_activation)\r\n",
        "        for i in range(n_decoder_layers)]\r\n",
        "\r\n",
        "                                        # output tuple - leading number is max index    \r\n",
        "    return tl.Serial(                   # 7: 0:tok_e 1:tok_p 2:tok_t 3:tok_tid 4:tok_d 5:tok_pq, 6:tok_tsd 7:wts_l  \r\n",
        "        tl.Select([0, 1, 2, 3, 3, 3,    # 10: 0:tok_e 1:tok_p 2:tok_t 3:tok_tid 4:tok_tid 5: tok_tid\r\n",
        "                   4, 5, 6, 4]),        #     6:tok_d 7:tok`_pq, 8:tok_tsd 9:tok_d 10:wts_l\r\n",
        "\r\n",
        "        # Encode.\r\n",
        "        tl.Parallel(\r\n",
        "            tl.Select([0, 1, 2, 3]),\r\n",
        "            PaddingFutureMask(tid=True)\r\n",
        "        ),                              # 10: tok_e tok_p tok_t tok_tid mask_combined tok_tid tok_d tok_pq tok_tsd tok_d wts_l\r\n",
        "        encoder,                        # 7: vec_e mask_combined tok_tid tok_d tok_pq tok_tsd tok_d wts_l\r\n",
        "        # Decode.\r\n",
        "        tl.Select([3, 4, 5, 2, 2, 0]),  # 7: tok_d tok_pq tok_tsd tok_tid tok_tid vec_e tok_d wts_l\r\n",
        "        tl.Parallel(\r\n",
        "            tl.ShiftRight(mode=mode),\r\n",
        "            tl.ShiftRight(mode=mode),  \r\n",
        "            tl.ShiftRight(mode=mode),\r\n",
        "            tl.ShiftRight(mode=mode),\r\n",
        "            tl.Serial(tl.ShiftRight(),\r\n",
        "                      PaddingFutureMask(tid=False)),\r\n",
        "        ),                              # 7: tok_d tok_pq tok_tsd tok_tid mask_combined vec_e tok_d wts_l \r\n",
        "        out_encoder,                    # 4: vec_d mask_combined vec_e tok_d wts_l\r\n",
        "        encoder_decoder_blocks,         # 4: vec_d mask_combined vec_e tok_d wts_l\r\n",
        "        tl.LayerNorm(),                 # 4: vec_d mask_combined vec_e tok_d wts_l\r\n",
        "\r\n",
        "        # Map to output vocab.\r\n",
        "        tl.Select([0], n_in=3),         # 3: vec_d tok_d wts_l\r\n",
        "        tl.Dense(d_out[0]),             # vec_d .....\r\n",
        "    )\r\n",
        "\r\n",
        "\r\n",
        "def _KTEncoderBlock(d_model, d_ff, n_heads,\r\n",
        "                  dropout, dropout_shared_axes, mode, ff_activation):\r\n",
        "    \"\"\"Same as the default, but changes attention layer to KTAttention to \r\n",
        "    accept a combined padding and future mask.\r\n",
        "    \"\"\"\r\n",
        "    \r\n",
        "    attention = KTAttention(\r\n",
        "        d_model, n_heads=n_heads, dropout=dropout, mode=mode)\r\n",
        "\r\n",
        "    feed_forward = _KTFeedForwardBlock(\r\n",
        "        d_model, d_ff, dropout, dropout_shared_axes, mode, ff_activation)\r\n",
        "\r\n",
        "    dropout_ = tl.Dropout(\r\n",
        "        rate=dropout, shared_axes=dropout_shared_axes, mode=mode)\r\n",
        "\r\n",
        "    return [\r\n",
        "        tl.Residual(\r\n",
        "            tl.LayerNorm(),\r\n",
        "            attention,\r\n",
        "            dropout_,\r\n",
        "        ),\r\n",
        "        tl.Residual(\r\n",
        "            feed_forward\r\n",
        "        ),\r\n",
        "    ]\r\n",
        "\r\n",
        "def _KTEncoderDecoderBlock(d_model, d_ff, n_heads,\r\n",
        "                         dropout, dropout_shared_axes, mode, ff_activation):\r\n",
        "    \"\"\"Same as the default, but changes the first layer to KTAttention to \r\n",
        "    accept a combined padding and future mask.\r\n",
        "    \"\"\"\r\n",
        "    def _Dropout():\r\n",
        "        return tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)\r\n",
        "\r\n",
        "    attention = KTAttention(\r\n",
        "        d_model, n_heads=n_heads, dropout=dropout, mode=mode)\r\n",
        "\r\n",
        "    attention_qkv = tl.AttentionQKV(\r\n",
        "        d_model, n_heads=n_heads, dropout=dropout, mode=mode)\r\n",
        "\r\n",
        "    feed_forward = _KTFeedForwardBlock(\r\n",
        "        d_model, d_ff, dropout, dropout_shared_axes, mode, ff_activation)\r\n",
        "\r\n",
        "    return [                             # vec_d masks vec_e\r\n",
        "        tl.Residual(\r\n",
        "            tl.LayerNorm(),              # vec_d ..... .....\r\n",
        "            attention,                   # vec_d ..... .....\r\n",
        "            _Dropout(),                  # vec_d ..... .....\r\n",
        "        ),\r\n",
        "        tl.Residual(\r\n",
        "            tl.LayerNorm(),              # vec_d ..... .....\r\n",
        "            tl.Select([0, 2, 2, 1, 2]),  # vec_d vec_e vec_e masks vec_e\r\n",
        "            attention_qkv,               # vec_d masks vec_e\r\n",
        "            _Dropout(),                  # vec_d masks vec_e\r\n",
        "        ),\r\n",
        "        tl.Residual(\r\n",
        "            feed_forward                 # vec_d masks vec_e\r\n",
        "        ),\r\n",
        "    ]\r\n",
        "\r\n",
        "def _KTFeedForwardBlock(d_model, d_ff, dropout, dropout_shared_axes,\r\n",
        "                      mode, activation):\r\n",
        "    \"\"\"Same as default.\r\n",
        "    \"\"\"\r\n",
        "    dropout_middle = tl.Dropout(\r\n",
        "        rate=dropout, shared_axes=dropout_shared_axes, mode=mode)\r\n",
        "    dropout_final = tl.Dropout(\r\n",
        "        rate=dropout, shared_axes=dropout_shared_axes, mode=mode)\r\n",
        "\r\n",
        "    return [\r\n",
        "        tl.LayerNorm(),\r\n",
        "        tl.Dense(d_ff),\r\n",
        "        activation(),\r\n",
        "        dropout_middle,\r\n",
        "        tl.Dense(d_model),\r\n",
        "        dropout_final,\r\n",
        "    ]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFlp9RAINR4d"
      },
      "source": [
        "## Configuration"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5nj3MX9D97Nz"
      },
      "source": [
        "# Configure hyperparameters.\r\n",
        "\r\n",
        "total_steps = 10000\r\n",
        "\r\n",
        "gin.clear_config()\r\n",
        "gin.parse_config(f\"\"\"\r\n",
        "import trax.layers\r\n",
        "import trax.models\r\n",
        "import trax.optimizers\r\n",
        "import trax.data.inputs\r\n",
        "import trax.supervised.trainer_lib\r\n",
        "\r\n",
        "# Parameters that will vary between experiments:\r\n",
        "# ==============================================================================\r\n",
        "# min_len = 12\r\n",
        "# max_len = 64\r\n",
        "# d_model = 512 # need to make sure this works with concat embeddings\r\n",
        "# d_ff = 256\r\n",
        "# n_encoder_layers = 2\r\n",
        "# n_decoder_layers = 2\r\n",
        "# n_heads = 2\r\n",
        "# dropout = 0.0\r\n",
        "\r\n",
        "min_len = 12\r\n",
        "max_len = 256\r\n",
        "d_model = 512 # need to make sure this works with concat embeddings\r\n",
        "d_ff = 1024\r\n",
        "n_encoder_layers = 6\r\n",
        "n_decoder_layers = 6\r\n",
        "n_heads = 8\r\n",
        "dropout = 0.1\r\n",
        "\r\n",
        "# Set to True to aggregate embeddings by concatenation. If set\r\n",
        "# to False aggregation will be by sum.\r\n",
        "embed_concat = True\r\n",
        "\r\n",
        "# (Vocab, depth) Uncomment to use with aggregation by concatenation.\r\n",
        "d_input = (13500, 384)\r\n",
        "d_part = (8, 8)\r\n",
        "d_tags = (189, 120)\r\n",
        "\r\n",
        "# (Vocab, depth) Uncomment to use with aggregation by concatenation.\r\n",
        "d_out = (3, 384)\r\n",
        "d_pqet = (300, 64)\r\n",
        "d_ts_delta = (150, 64)\r\n",
        "\r\n",
        "# Used for positional encodings if not None. Positional encoding based\r\n",
        "# on sequence in batch if None.\r\n",
        "d_tid = (10000, %d_model)\r\n",
        "\r\n",
        "# d_input = (13500, %d_model)\r\n",
        "# d_part = (8, %d_model)\r\n",
        "# d_tags = (189, %d_model)\r\n",
        "# # d_tags = None\r\n",
        "# d_out = (3, %d_model)\r\n",
        "# d_pqet = (300, %d_model)\r\n",
        "# d_ts_delta = (150, %d_model)\r\n",
        "# d_tid = (10000, %d_model)\r\n",
        "\r\n",
        "total_steps = {total_steps}\r\n",
        "\r\n",
        "# Parameters for learning rate schedule:\r\n",
        "# ==============================================================================\r\n",
        "warmup_and_rsqrt_decay.n_warmup_steps = 3000\r\n",
        "warmup_and_rsqrt_decay.max_value = 0.001\r\n",
        "\r\n",
        "# multifactor.constant = 0.01\r\n",
        "# multifactor.factors = 'constant * linear_warmup * cosine_decay'\r\n",
        "# multifactor.warmup_steps = 4000\r\n",
        "# multifactor.steps_per_cycle = %total_steps\r\n",
        "# multifactor.minimum = .0001\r\n",
        "\r\n",
        "# Parameters for Adam:\r\n",
        "# ==============================================================================\r\n",
        "# Adam.weight_decay_rate=0.0\r\n",
        "Adam.b1 = 0.9\r\n",
        "Adam.b2 = 0.999\r\n",
        "Adam.eps = 1e-8\r\n",
        "\r\n",
        "# Parameters for input pipeline:\r\n",
        "# ==============================================================================\r\n",
        "get_ds_tfrec.min_len = %min_len\r\n",
        "get_ds_tfrec.max_len = %max_len\r\n",
        "train/get_ds_tfrec.folds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]\r\n",
        "eval/get_ds_tfrec.folds = [19]\r\n",
        "\r\n",
        "BucketByLength.boundaries =  [32, 64, 128]\r\n",
        "BucketByLength.batch_sizes = [512, 256, 128, 64]\r\n",
        "# BucketByLength.batch_sizes = [16, 8, 4,  2]\r\n",
        "\r\n",
        "BucketByLength.strict_pad_on_len = True\r\n",
        "\r\n",
        "KTAddLossWeights.id_to_mask = 0\r\n",
        "\r\n",
        "train/make_additional_stream.stream = [\r\n",
        "  @train/get_ds_tfrec(),\r\n",
        "  @BucketByLength(),\r\n",
        "  @TrimTags(),\r\n",
        "  @KTAddLossWeights()\r\n",
        "]\r\n",
        "\r\n",
        "eval/make_additional_stream.stream = [\r\n",
        "  @eval/get_ds_tfrec(),\r\n",
        "  @BucketByLength(),\r\n",
        "  @TrimTags(),\r\n",
        "  @KTAddLossWeights()\r\n",
        "]\r\n",
        "\r\n",
        "make_inputs.train_stream = @train/make_additional_stream()\r\n",
        "make_inputs.eval_stream = @eval/make_additional_stream()\r\n",
        "\r\n",
        "# Parameters for KTPositionalEncoder:\r\n",
        "# ==============================================================================\r\n",
        "KTPositionalEncoder.d_model = %d_model\r\n",
        "\r\n",
        "# Set to True to calculate positional encodings based on position in orginal\r\n",
        "# full length sequence, False to be based on position in batch sequence.\r\n",
        "KTPositionalEncoder.tid = False\r\n",
        "\r\n",
        "# Parameters for PaddingFutureMaske:\r\n",
        "# ==============================================================================\r\n",
        "PaddingFutureMask.pad_end = False\r\n",
        "\r\n",
        "# Set to True to calculate future mask based on task container id (questions\r\n",
        "# are delivered to users in groups identified by task_container id) or False\r\n",
        "# to be based next question only.\r\n",
        "PaddingFutureMask.tid = False\r\n",
        "\r\n",
        "# Parameters for KTTransformer:\r\n",
        "# ==============================================================================\r\n",
        "KTTransformer.d_model = %d_model\r\n",
        "KTTransformer.d_input = %d_input\r\n",
        "KTTransformer.d_part = %d_part\r\n",
        "KTTransformer.d_tags = %d_tags\r\n",
        "KTTransformer.d_out = %d_out\r\n",
        "KTTransformer.d_pqet = %d_pqet\r\n",
        "KTTransformer.d_ts_delta = %d_ts_delta\r\n",
        "KTTransformer.d_tid = %d_tid\r\n",
        "KTTransformer.embed_concat = %embed_concat\r\n",
        "KTTransformer.d_ff = %d_ff\r\n",
        "KTTransformer.n_encoder_layers = %n_encoder_layers\r\n",
        "KTTransformer.n_decoder_layers = %n_decoder_layers\r\n",
        "KTTransformer.n_heads = %n_heads\r\n",
        "KTTransformer.dropout = %dropout\r\n",
        "\r\n",
        "# Parameters for train:\r\n",
        "# ==============================================================================\r\n",
        "train.inputs = @make_inputs\r\n",
        "train.eval_frequency = 200\r\n",
        "train.eval_steps = 20\r\n",
        "train.checkpoints_at = {list(range(0,total_steps + 1, 2000))}\r\n",
        "train.optimizer = @trax.optimizers.Adam\r\n",
        "train.steps = %total_steps\r\n",
        "train.model = @KTTransformer\r\n",
        "train.lr_schedule_fn = @trax.supervised.lr_schedules.warmup_and_rsqrt_decay\r\n",
        "\"\"\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PP46-cEAB4i_"
      },
      "source": [
        "if False:\r\n",
        "    inputs = trax.data.inputs.make_inputs()\r\n",
        "    train_stream = inputs.train_stream(trax.fastmath.device_count())\r\n",
        "    train_eval_stream = inputs.train_eval_stream(trax.fastmath.device_count())\r\n",
        "    b = next(train_stream)\r\n",
        "    for i, m in enumerate(b):\r\n",
        "        print(i, m.shape)\r\n",
        "    b"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8OmbSWj5Cvt2"
      },
      "source": [
        "if False:\r\n",
        "    model = KTTransformer()\r\n",
        "    model.init(trax.shapes.signature(b))\r\n",
        "    outs = model(b)\r\n",
        "    for i, m in enumerate(outs):\r\n",
        "        print(i, m.shape)\r\n",
        "    outs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dcHfhEEFNXXJ"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XGIfCGL9GprB"
      },
      "source": [
        "run_no = 0\r\n",
        "prefix = f'model_runs/{run_no:02d}'\r\n",
        "output_dir = f'gs://{BUCKET}/{prefix}'\r\n",
        "log_dir = output_dir[:-3]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5rhH1YNVHEPO"
      },
      "source": [
        "%tensorboard --logdir $log_dir"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wQcrZYSTGMhX"
      },
      "source": [
        "if TRAIN_MODEL:\r\n",
        "    if False:\r\n",
        "        init_checkpoint = f'{output_dir}/model.pkl.gz'\r\n",
        "    else:\r\n",
        "        bucket.delete_blobs(list(bucket.list_blobs(prefix=prefix)))\r\n",
        "\r\n",
        "    loop = trax.supervised.trainer_lib.train(output_dir, metrics=metrics)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}