{"cells":[{"metadata":{"id":"UnkZKyp_11iC"},"cell_type":"markdown","source":"![](\"https://storage.googleapis.com/kaggle-competitions/kaggle/21651/logos/header.png?t=2020-09-09-03-03-31\")
\n
\nTPU - Track knowledge states of 1M+ students in the wild
\n
\n\n### In this notebook, I demonstrate a tensorflow pipeline for the competition [Riiid! Answer Correctness Prediction](#https://www.kaggle.com/c/riiid-test-answer-prediction). Other than the common knowledge about using TPU (which is presented in my another Kaggle notebook [Detailed guide to custom training with TPUs](#https://www.kaggle.com/yihdarshieh/detailed-guide-to-custom-training-with-tpus)), you will see some particular techniques, including:\n\n* Dealing with tf.RaggedTensor\n* Advanced tf.data.Dataset manipulation\n* Multi steps per TPU call\n* Accumulate the predictions from the validation dataset with TPU\n\nThe model definition is a modified copy of the [distilbert.py](#https://github.com/huggingface/transformers/tree/master/src/transformers/models/distilbert) file from [Hugging Face's transformer library](#https://github.com/huggingface/transformers). If you find the code quality and object naming in this part is not perfect, it is due to my personal modification which doesn't refelect the high quality work done by Hugging Face great work.\n\nTwo special contributions about model definition in this notebook:\n\n* implementation of decoder\n* implementation of auto-regressive prediction generation (only works on CPU / GPU).\n\nI tried to make this notebook working both on Kaggle and Google Colab. A few but minimal change is still required if you want to work on Google Colab."},{"metadata":{},"cell_type":"markdown","source":"# Table of Contents\n\n0. [Tips and clarifications](#tips)\n1. [Environment](#environment)\n2. [Packages](#packages)\n3. [TPU](#tpu)\n4. [Data](#data)\n5. [Configuration](#configuration)\n * [Model / Training](#model-training-settings)\n * [Running mode](#running-mode)\n * [Vocabulary](#vocabulary-settings)\n6. [Train Manager](#train-manager)\n * [Dataset](#dataset)\n - [TFRecord files](#tfrecord-files)\n - [Load TFRecord files - tf.io.RaggedFeature](#load-tfrecord-files)\n - [Train / Valid split](#split)\n - [Transformation - from tf.RaggedTensor to tf.Tensor](#transformation)\n * [Model inputs](#inputs)\n - [Special masks](#special-masks)\n - [Input tensors](#input-tensors)\n - [For validation dataset](#for-valid)\n * [Model definition](#model-def)\n * [Training Manager](#training-manager)\n7. [Train / Valid](#train-valid)\n8. [Conclusion](#conclusion) "},{"metadata":{},"cell_type":"markdown","source":"# List of tips and clarifications"},{"metadata":{},"cell_type":"markdown","source":"## TPU - Tips and remarks\n\n1. Since tensorflow 2.3 works better with the code in this notebook, I need to upgrade to TF 2.3 manually. However, the TPU version was still in version 2.2, use the line below to select the corresponding TPU version.\n \n ```\n Client().configure_tpu_version(tf.__version__, restart_type='ifNeeded')\n ```\n \n2. For loss calculation in training / validation, on each replica on TPU, we compute per-example loss, and divid the sum by the number of places where the target is not `NON_TARGET_ID` (which is `-100`) on the whole batch before distributing to each replica. This is because the gradients are synchronized over the replicas by summing them before the optimizer updates the model parameters. So the way we calculate the loss on each replica will give the final gradient corresponding to the average loss over the places with real targets in the whole batch distributed to each replica. See my another [Detecting contradiction and entailment in multilingual text using TPUs](#https://www.kaggle.com/yihdarshieh/masked-my-dear-watson-mlm-with-tpu#MLM-loss-calculation) for more details. The code looks like (where `train_batch` is a batch received by a replica, but `train_batch['nb_pred_places']` is the pre-computed number of places on the whole batch before being distributed to replicas):\n \n ```\n # Need to have the 1st dimension to make the losses not averaged\n losses = loss_obj(selected_targets[:, tf.newaxis], selected_logits[:, tf.newaxis])\n total_loss = tf.math.reduce_sum(losses)\n\n # `train_batch['nb_pred_places'][0]` is the total number of places used for calculating loss across replicas\n loss = total_loss / tf.cast(train_batch['nb_pred_places'][0], dtype=DTYPE) \n ```\n \n3. For training and validation, we want to perform multiple steps in a single TPU call to speed up the computation and avoid communication overhead between the local VM and XLA/TPU host. This is achieved by wrapping the dataset iteration in a `tf.range` for loop. See `dist_train_multi_steps` and ``dist_valid_multi_steps` in [Training Manager](#training-manager-main).\n\n4. For validation, the last call to TPU might contain fewer steps. However, using TPU requires `tf.range` has a compile time known number of steps. Therefore the following 2 methods are provided, despite they are almost identical. The later is used for the last call to TPU.\n\n * dist_valid_multi_steps\n * dist_valid_multi_steps_last_call\n \n5. During the validation, we want to collect the predictions and other information in order to save them to files to further investigation. However, the validation is run in graph model with TPU. In order to accumulate the restuls, we use `tf.TensorArray`. See the method `dist_valid_multi_steps` in [Training Manager](#training-manager-main).\n\n6. For validation, ideally we don't want to drop the last batch which might contain fewer examples. The code in this notebook, while running with Kaggle TPU, gives error if we don't drop the last batch. While running on Colab with TPU, no error occurs when we keep the last batch. So the following condition is used to make this notebook runs on both platform\n \n ```\n valid_ds.apply(\n tf.data.experimental.dense_to_ragged_batch(\n batch_size=batch_size, drop_remainder=(IS_KAGGLE and tpu is not None)\n )\n ) \n ```\n\n7. While training on Kaggle with TPU, `tf.train.CheckpointManager` can save the checkpoints locally and gives `File system scheme '[local]' not implemented` error. The following code solves this issue\n\n ```\n ckpt_manager.checkpoint.save(\n file_prefix=ckpt_manager.directory + 'ckpt',\n options=tf.train.CheckpointOptions(experimental_io_device=\"/job:localhost\")\n )\n ```\n \n8. Although I made (a lot of) effort to implement the auto-regression generation, unfortunately it can't run with TPU. It gives `XLA can't deduce compile time constant output shape for strided slice` error, which seems to me a limitation in XLA. However, it works while running on GPU. Therefore I need to comment out the block below in `valid_step` in [Training Manager](#training-manager-main), so XLA won't compile it at all. You can check in this version [GPU Running](#https://www.kaggle.com/yihdarshieh/r3id-tf-tpu/output?scriptVersionId=49027558) which shows it works with GPU.\n\n ```\n # if generative == 1:\n\n # start_pos = self.train_config.window_size - MAX_PRED_TIME_QUESTION_BUNDLE_LEN\n # _, logits = predictor.generate(\n # valid_batch, start_pos=start_pos,\n # window_size=self.train_config.window_size,\n # dim=self.config.dim,\n # c_mask=c_mask, r_mask=r_mask, r_c_mask=r_c_mask, c_r_mask=c_r_mask\n # )\n\n # else:\n ```"},{"metadata":{},"cell_type":"markdown","source":"## Model / Input - Terminology and explanation\n\n1. In `EdFormerConfig`, we have\n * model_type:\n - `'cr'`: encoder only\n - `'ed'` encoder decoder\n * generative (only applies to decoder during the inference time):\n - `True`: use auto-regressive generation, i.e. use the previous prediction result as input to the current prediction step. This is the usual way for decoder prediction.\n - `False`: pretend each of the questions in a bunlde (that we want to predict answer correctness) as a single question in a bundle, and only use the history before that bundle.\n\n1. In the method `add_input_ids_and_targets` in [Input tensors](#input-tensors), we compute (quite a lot) tensors that are used in the model. Among them:\n * c_input_ids: The (encoded) ids for questions and lectures. This is used for the model encoder. Questions and lectures (along 4 special token) form a single vocabulary and their encoded ids (i.e. `c_input_ids`) is fed to a single embedding layer.\n * r_input_ids: The (encoded) answer correction. They won't be `0` and `1` anymore - the `0` was used to say `answered incorrectly`, but the `0` in the encoded ids means padding See [Vocabulary settings](#vocabulary-settings) for the mapping between tokens and their encoded ids. `r_input_ids` is used only if the model is of type `encoder only`, which is the case if we pass `model_type='cr'` to `EdFormerConfig.__init__()`.\n * d_input_ids: This is a shifted version of `r_input_ids` to indicate the ids in the previous place in a user interaction history. This is used only if the model is of type `encoder decoder`, which could be specified in `EdFormerConfig.__init__()` by passing `model_type='ed'`.\n \n2. `r_input_ids` and `d_input_ids` will have different values, depending on if we are in the training or inference (including validation) time, because during inference, the answer correction information is not availabe (at the time we are predicting for a particular place). Basically, the values not available in inference time are replaced by `MASK_ID`.\n\n3. For `encoder decoder` type, the `d_input_ids` will have different values during the inference time depending on if `EdFormerConfig.generative` is `True` or `False`:\n\n * if `generative` is `False`, we copy the value of `d_input_ids` at the 1st question in the question bundle being predicted and use it as the values of `d_input_ids` for all the remaining places in the same bunlde. This trick also applies to the positional information and attention masks - the objective is to pretend each question in the question bundle under prediction as the single question in the bundle.\n * if `generative` is `True`, we leave it as it is - which contains `MASK_ID` for all places in the question bundle to predict except the 1st question in that bundle.\n \n4. Despite we have the `r_input_ids` and `d_input_ids`, the decoder self-attention mask and the decoder-to-encoder attention mask are called `r_mask` and `r_c_mask` rather than `d_mask` and `d_c_mask`.\n\n5. Considering the amount of options provided in this notebook, it is not easy to provide all the information in detail - in particular being in the competition and would like to improve my scores. I will try to answer questions if you have any."},{"metadata":{"id":"CKfIrA-M3JNc"},"cell_type":"markdown","source":"# Environment\n\nTo make this notebook both work on Kaggle and Google Colab.\n\nIn order to use this notebook on Google Colab, you have to specify your GCP bucket name and your project name. Also, copy the directories/files specifiedd in [Data section](#data) to your GCP bucket."},{"metadata":{},"cell_type":"markdown","source":"## Important for GPU"},{"metadata":{"trusted":true},"cell_type":"code","source":"# It is unclear why the model is not running on GPU when using tensorflow 2.3.1.\n# If GPU is enabled, it will avoid to upgrade to tensorflow 2.3.1 later.\n\ngpu_info = !nvidia-smi\nif 'command not found' in gpu_info[0]:\n USE_GPU = False\nelif 'GPU Name ' in str(gpu_info):\n USE_GPU = True\nelse:\n USE_GPU = False\n \nprint(f'USE_GPU: {USE_GPU}')","execution_count":null,"outputs":[]},{"metadata":{"id":"9if28vqJ3PFh","outputId":"10789cc9-5fa2-4ef2-d339-39cda6cf7d94","trusted":true},"cell_type":"code","source":"import os\n\nIS_KAGGLE = os.path.isdir('/kaggle/input')\n\nif IS_KAGGLE:\n BASE_DIR = '/kaggle/input'\nelse:\n BUCKET_DIR = 'gs://[YOUR_GCP_BUCKET_NAME]/r3id'\n BASE_DIR = '/content/drive/My Drive/r3id'\n BASE_DIR_QUOTED = '\"/content/drive/My Drive/r3id\"'\n \nif not IS_KAGGLE:\n\n # Access GCP Bucket\n from google.colab import auth\n auth.authenticate_user()\n project_id = '[YOUR_GCP_PROJECT_NAME]'\n !gcloud config set project {project_id}\n !gsutil ls {BUCKET_DIR}\n\n # Access Google Drive\n from google.colab import drive\n drive.mount('/content/drive')\n\n if not os.path.isdir(BASE_DIR):\n\n !mkdir {BASE_DIR_QUOTED}\n !gsutil -m cp -r 'gs://[YOUR_GCP_BUCKET_NAME]/r3id' \"/content/drive/My Drive\"\n\n!ls -l '{BASE_DIR}'","execution_count":null,"outputs":[]},{"metadata":{"id":"5VXzmicseiDu"},"cell_type":"markdown","source":"# Packages"},{"metadata":{"id":"EbTukhm4qeyt"},"cell_type":"markdown","source":"## Install"},{"metadata":{"trusted":true,"id":"jvhwIoMd11if","outputId":"7285ed3e-08c2-4f58-b6d2-31becaf3ed3f"},"cell_type":"code","source":"# Use a newer version of huggingface's `tokenizers` (?? -> 0.9.2) \n# Use a newer version of huggingface's `transformers` (3.0.2 -> 3.4.0) \n!pip uninstall -y datatable\n!pip uninstall -y tokenizers\n!pip uninstall -y transformers\n\nif IS_KAGGLE:\n !pip install '{BASE_DIR + '/r3id-packages/datatable-0.11.0-cp37-cp37m-manylinux2010_x86_64.whl'}'\n !pip install --upgrade '{BASE_DIR + '/r3id-packages/tokenizers-0.9.2-cp37-cp37m-manylinux1_x86_64.whl'}'\n if not USE_GPU:\n !pip uninstall -y tensorflow\n !pip install --upgrade tensorflow==2.3.1\n !pip install --upgrade tensorflow-probability\nelse:\n !pip install '{BASE_DIR + '/r3id-packages/datatable-0.11.0-cp36-cp36m-manylinux2010_x86_64.whl'}'\n !pip install --upgrade '{BASE_DIR + '/r3id-packages/tokenizers-0.9.2-cp36-cp36m-manylinux1_x86_64.whl'}'\n\n!pip install --upgrade '{BASE_DIR + '/r3id-packages/transformers-3.4.0-py3-none-any.whl'}'\n!pip install cloud-tpu-client","execution_count":null,"outputs":[]},{"metadata":{"id":"ufr8QIu5qg-Y"},"cell_type":"markdown","source":"## Import"},{"metadata":{"trusted":true,"id":"7E5iDmMK11jf"},"cell_type":"code","source":"import os\nimport pandas as pd\nimport datatable as dt\nimport json\nfrom collections import defaultdict\nimport random\nimport math\nimport datetime\nimport collections\nimport tensorflow as tf\nos.environ['TF_DETERMINISTIC_OPS'] = '1'\nimport tensorflow_probability as tfp\nfrom copy import deepcopy\nimport gc\ngc.enable()\nfrom cloud_tpu_client import Client\n\nimport transformers\nfrom transformers import PretrainedConfig, DistilBertConfig\nfrom transformers.modeling_tf_distilbert import TFFFN, TFTransformer\nfrom transformers.modeling_tf_distilbert import TFMultiHeadSelfAttention as HFTFMultiHeadSelfAttention\nfrom transformers.modeling_tf_distilbert import TFTransformerBlock as HFTFTransformerBlock\nfrom transformers.modeling_tf_distilbert import TFFFN as HFTFFFN\n\nfrom transformers.modeling_tf_utils import (\n TFPreTrainedModel,\n TFSharedEmbeddings,\n keras_serializable,\n shape_list,\n)\nif IS_KAGGLE:\n import riiideducation\n from kaggle_datasets import KaggleDatasets","execution_count":null,"outputs":[]},{"metadata":{"id":"DdpkJbooTueu"},"cell_type":"markdown","source":"# TPU\n\nSetup TPU\n\nSince we upgrade to tensorflow 2.3 manually from the defalt version (2.2), we need to use the following code to configure TPU version.\n\n Client().configure_tpu_version(tf.__version__, restart_type='ifNeeded')\n \n"},{"metadata":{"id":"uGKoipn2Ttol","outputId":"e7d992a6-3ded-4a70-e2c7-bbd4f839d228","trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\n\ntry:\n if IS_KAGGLE:\n Client().configure_tpu_version(tf.__version__, restart_type='ifNeeded')\n # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n tpu = tf.distribute.cluster_resolver.TPUClusterResolver() \n print('Running on TPU ', tpu.master())\nexcept ValueError:\n tpu = None\n\nif tpu:\n tf.config.experimental_connect_to_cluster(tpu)\n tf.tpu.experimental.initialize_tpu_system(tpu)\n strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"id":"OuB4HtjyqyGw"},"cell_type":"markdown","source":"# Data\n\nPaths to data files used for this notebook."},{"metadata":{"id":"l8LuQDTier5n","trusted":true},"cell_type":"code","source":"n_contents_dict_path = f'{BASE_DIR}/r3id-info-public/n_contents_dict.json'\n\nquestion_tags_info_path = f'{BASE_DIR}/r3id-info-public/question_tags_info.json'\nlecture_tags_info_path = f'{BASE_DIR}/r3id-info-public/lecture_tags_info.json'\n\nquestion_part_info_path = f'{BASE_DIR}/r3id-info-public/question_part_info.json'\nlecture_part_info_path = f'{BASE_DIR}/r3id-info-public/lecture_part_info.json'\n\nvalid_info_path = f'{BASE_DIR}/r3id-info-public/valid_info_fold_1.json'\ntrain_valid_split_indices_path = f'{BASE_DIR}/r3id-info-public/train_valid_split_indices_fold_1.json'\n\ntrain_tfrec_dir_local = f'{BASE_DIR}/ednet-tfrecords-sequential'\nvalid_tfrec_dir_local = f'{BASE_DIR}/r3id-tfrecords-valid-public'\n\nif tpu is None:\n \n train_tfrec_dir = f'{BASE_DIR}/ednet-tfrecords-sequential'\n valid_tfrec_dir = f'{BASE_DIR}/r3id-tfrecords-valid-public'\n\nelif IS_KAGGLE:\n \n train_tfrec_dir = KaggleDatasets().get_gcs_path('ednet-tfrecords-sequential')\n valid_tfrec_dir = KaggleDatasets().get_gcs_path('r3id-tfrecords-valid-public')\n\n print(f'train_tfrec_dir from Kaggle = {train_tfrec_dir}\\n')\n print(f'valid_tfrec_di from Kaggle = {valid_tfrec_dir}\\n')\n\n # you can list the buckets\n !gsutil ls $train_tfrec_dir\n !gsutil ls $valid_tfrec_dir\n \nelse:\n \n train_tfrec_dir = f'{BUCKET_DIR}/ednet-tfrecords-sequential'\n valid_tfrec_dir = f'{BUCKET_DIR}/r3id-tfrecords-valid-public'","execution_count":null,"outputs":[]},{"metadata":{"id":"mf0RGxVP11kR"},"cell_type":"markdown","source":"# Configuration"},{"metadata":{"id":"NUmudgpdsKXW"},"cell_type":"markdown","source":"## Model / Training settings"},{"metadata":{"id":"NgInt0xvUF99","trusted":true},"cell_type":"code","source":"MODEL_TYPE = 'ed'\nMODEL_SIZE = 'small' # specify some model size parameters in the next cell\nMODEL_DESC = 'edformer-tpu'\n\nACTIVATION = 'gelu'\nUSE_ABS_POS = False\nSHARE_POS_EMBEDDING = True\nUSE_TAGS = True\nUSE_PART = True\nUSE_PRIOR_EXPLANATION = True\nUSE_PRIOR_QUESTION_ELAPSED_TIME_INPUT = True\nALLOW_BUNDLE_ATTEN = False\nGENERATIVE = True\n\nWINDOW_SIZE = 96\nLOSS_WEIGHT_WINDOW_SIZE = None\nBATCH_SIZE = 64 * strategy.num_replicas_in_sync\nPRED_BATCH_SIZE = 256 * strategy.num_replicas_in_sync\n\nassert BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\nassert PRED_BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\n\nN_EPOCHS = 6\nSTEPS_PER_CALL = 1000\n\nLR = 1e-3\nEND_LR = 5e-4\nWARMUP_STEPS = 4000\n\nDETERMINISTIC = False\n\nif DETERMINISTIC:\n \n SEED = 2021\n N_PARALLEL_READS = None # 1\n N_PARALLEL_CALLS = None # 1\n SHUFFLE_BUFFER_SIZE = 1\n\nelse:\n \n SEED = None\n N_PARALLEL_READS = 16\n N_PARALLEL_CALLS = tf.data.experimental.AUTOTUNE\n SHUFFLE_BUFFER_SIZE = 4096\n\nMAX_TRAIN_ITER_STEPS = None\nMAX_VALID_ITER_STEPS = None\n\nPRINTING_STEPS = 1000\n\nCKPT_DIR = f'{MODEL_TYPE}-{MODEL_SIZE}-win-{WINDOW_SIZE}-bs-{BATCH_SIZE}-{MODEL_DESC}/'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### Specify some model size parameters and assign a name to it."},{"metadata":{"trusted":true},"cell_type":"code","source":"if MODEL_SIZE == 'small':\n\n N_LAYERS = 2\n N_HEADS = 4\n DIM = 128\n HIDDEN_DIM = 4 * DIM","execution_count":null,"outputs":[]},{"metadata":{"id":"zI-dw1serob_"},"cell_type":"markdown","source":"## Running mode\n\nDetermine we want to run traning, validataion or prediction for commit / submission to this competition.\n\nBy validation, it is actually a validation pipeline I wrote for simulatesimulating the competition submission.\nThat part of code if not presented here. For this notebook, only `TRAIN = True` will be used and it contains the\nvalidation which is run in the tf.data.Dataset way - which is much faster."},{"metadata":{"trusted":true,"id":"mp1I7Ar811kZ"},"cell_type":"code","source":"DTYPE = tf.float32\ntf.keras.backend.set_floatx('float32')\n\nTRAIN = True\nVALID = False\nPRED = False\n\nRESUME_TRAINING = False\n\nN_FILES = 6\nSUBMISSION = False\n\nif IS_KAGGLE:\n N_FILES = len(os.listdir('/kaggle/input/riiid-test-answer-prediction'))\n SUBMISSION = (N_FILES != 6)\nelse:\n PRED = False\n\nif SUBMISSION:\n \n TRAIN = False\n VALID = False\n PRED = True\n\nif not TRAIN:\n RESUME_TRAINING = False\n\nDEBUG = True\nPROBE = False\n\nCKPT_TRAIN_PATH = None\nCKPT_PRED_PATH = None\nif CKPT_TRAIN_PATH is None:\n \n if IS_KAGGLE:\n CKPT_TRAIN_PATH = './'\n else:\n CKPT_TRAIN_PATH = f'{BUCKET_DIR}/r3id-ckpts/{CKPT_DIR}'\n \n if TRAIN and not RESUME_TRAINING:\n _state = !gsutil -q stat {CKPT_TRAIN_PATH}*; echo $?\n already_existed = 1 - int(_state[0])\n assert not already_existed \n\nif CKPT_PRED_PATH is None:\n \n if IS_KAGGLE:\n CKPT_PRED_PATH = f'{BASE_DIR}/{CKPT_DIR}'\n else:\n CKPT_PRED_PATH = f'{BUCKET_DIR}/r3id-ckpts/{CKPT_DIR}'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"id":"o77CsqXW11ld","outputId":"556ab83c-e6a1-4cd2-a4eb-5c0dab468151"},"cell_type":"code","source":"print(f'SUBMISSION: {SUBMISSION}')\nprint(f'TRAIN: {TRAIN}')\nprint(f'VALID: {VALID}')\nprint(f'PRED: {PRED}')\n\nprint('')\n\nprint(f'CKPT_TRAIN_PATH: {CKPT_TRAIN_PATH}')\nprint(f'CKPT_PRED_PATH: {CKPT_PRED_PATH}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Probed Info"},{"metadata":{"trusted":true},"cell_type":"code","source":"MAX_TRAIN_HISTORY_LEN = 17917\nMAX_PRED_TIME_QUESTION_BUNDLE_LEN = 10","execution_count":null,"outputs":[]},{"metadata":{"id":"BlcYdBBAr2_h"},"cell_type":"markdown","source":"## Vocabulary settings\n\nThe notions of `token` and `id` are different here. Think `id` as encoding of a `token`. It might be confusing - for example, a question with `content_id` being `100` (which is in the sense of token) will have an input id `104` to the model encoder, because we have 4 special tokens here. See [a later section](#token-id-mapping) for more details."},{"metadata":{"trusted":true,"id":"w9tdC-J811mN"},"cell_type":"code","source":"PAD_TOKEN = -2\nSTART_TOKEN = -3\nEND_TOKEN = -4\nMASK_TOKEN = -5\n\nPAD_ID = 0\nSTART_ID = 1\nEND_ID = 2\nMASK_ID = 3\n\nRESPONSE_LECTURE_TOKEN = -1\nRESPONSE_FALSE_TOKEN = 0\nRESPONSE_TRUE_TOKEN = 1\n\nRESPONSE_LECTURE_ID = 4\nRESPONSE_FALSE_ID = 5\nRESPONSE_TRUE_ID = 6\n\nNON_TARGET_ID = -100\n\nTAG_VOCAB_SIZE = 189 # including extra `PAD_TOKEN`.\nPART_VOCAB_SIZE = 9 # including extra `PAD_TOKEN`.\nPRIOR_EXPLANATION_VOCAB_SIZE = 3 # including extra `PAD_TOKEN`.\n\n# --------------------------------------------------\n\nN_TAGS_PER_CONTENT = 6\n\n# --------------------------------------------------\n\n# This is only an assumption (including `0` for padding)\nMAX_HISTORY_LEN = 20480\n# --------------------------------------------------","execution_count":null,"outputs":[]},{"metadata":{"id":"yDB5LrYw11pg"},"cell_type":"markdown","source":"# Train Manager"},{"metadata":{"id":"34iiHe4c11pk"},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"id":"njJiZx_p11pn"},"cell_type":"markdown","source":"### TFRecord files"},{"metadata":{"trusted":true,"id":"fxD0l4Vr11ps","outputId":"ff056e7a-a17c-4a43-abb4-32f8b8fca539"},"cell_type":"code","source":"train_tfrec_fns = os.listdir(train_tfrec_dir_local)\n# Sort the file - For verification purpose\ntrain_tfrec_fns = sorted(train_tfrec_fns, key=lambda x: int(x.replace('EdNet-user-history-', '').replace('.tfrecord', '')))\n \nif tpu is None:\n train_tfrec_paths = [os.path.join(train_tfrec_dir, fn) for fn in train_tfrec_fns]\nelse:\n train_tfrec_paths = [train_tfrec_dir + f'/{fn}' for fn in train_tfrec_fns]\n\ntrain_tfrec_paths","execution_count":null,"outputs":[]},{"metadata":{"id":"PT9m690JLIAk","outputId":"d652fab8-1579-4fa5-ecc9-4a1f6588e8a0","trusted":true},"cell_type":"code","source":"valid_tfrec_fns = os.listdir(valid_tfrec_dir_local)\n\nif tpu is None:\n valid_tfrec_paths = [os.path.join(valid_tfrec_dir, fn) for fn in valid_tfrec_fns]\nelse:\n valid_tfrec_paths = [valid_tfrec_dir + f'/{fn}' for fn in valid_tfrec_fns]\n\nvalid_tfrec_paths","execution_count":null,"outputs":[]},{"metadata":{"id":"old9ZSlp11qH"},"cell_type":"markdown","source":"### Load TFRecord files - tf.io.RaggedFeature\n\nWe are dealing with `tf.io.RaggedFeature` here becausse users have different lengths of interaction history. When we use `tf.data.experimental.dense_to_ragged_batch` later to batch the examples, we will get `tf.RaggedTensor`.\n\nWe also inject some extra information, like\n* the length of each sequence\n* if an interaction is in the prediction time (for which we want to predict the answer correction)\n* absolute position of an interaction in the history\n\nThe tfrecord file for the validation dataset actually contains the full history of a user (if it is selected to be in the validation dataset). But it has extra attributes, including\n* the number of blocks in a user history selected to be used for validation\n* the starting and ending (exclusive) position of each block selected\n\nThis extra information eables us to build the truncated user history for computing the validation score.\n\nThe notion of `block` is defined as:\n \n a sequence staring with a potential lecture (which must be the first element in the history, or its previous interaction is a question) and a subsequentail lectures (if any), followed by a unique question bundle (if any).\n \nFor a given user, at a time, only 1 block could in prediction time. "},{"metadata":{"trusted":true,"id":"1wcmYrGh11qK"},"cell_type":"code","source":"# --------------------------------------------------------------------------------\n# For training dataset\n\ntrain_raw_features = {\n 'user_id': tf.io.FixedLenFeature([], dtype=tf.int64),\n 'row_id': tf.io.RaggedFeature(value_key='row_id', dtype=tf.int64),\n 'timestamp': tf.io.RaggedFeature(value_key='timestamp', dtype=tf.int64),\n 'content_id': tf.io.RaggedFeature(value_key='content_id', dtype=tf.int64),\n 'content_type_id': tf.io.RaggedFeature(value_key='content_type_id', dtype=tf.int64),\n 'task_container_id': tf.io.RaggedFeature(value_key='task_container_id', dtype=tf.int64),\n 'user_answer': tf.io.RaggedFeature(value_key='user_answer', dtype=tf.int64),\n 'answered_correctly': tf.io.RaggedFeature(value_key='answered_correctly', dtype=tf.int64),\n 'prior_question_elapsed_time': tf.io.RaggedFeature(value_key='prior_question_elapsed_time', dtype=tf.float32),\n 'prior_question_had_explanation': tf.io.RaggedFeature(value_key='prior_question_had_explanation', dtype=tf.int64),\n}\n\n\ndef parse_train_example(example):\n \"\"\"Parse an example from the training tfrecord files.\n\n Add the following extra attributes:\n \n - seq_len: The length of the user interaction history for training, before the validation dataset being removed from it.\n - prev_seq_len: The length of interaction history in an example from which the current example is obtained. Here, it just\n equals to `seq_len` since it has no source.\n - start: The starting index in the interaction history of an example from which the current example is obtained. Here, it is `0`.\n - end: The ending index in the interaction history of an example from which the current example is obtained. Here, it is `seq_len - 1`.\n - pred_time_mask: A l-D `tf.Tensor` of `0` and `1`, indicating if a place is in the prediction time. Here, all of them are `0`. \n\n \"\"\"\n\n _parsed = tf.io.parse_single_example(example, train_raw_features)\n \n parsed = {}\n parsed['user_id'] = _parsed['user_id']\n parsed['seq_len'] = tf.reduce_sum(tf.ones_like(_parsed['row_id'], dtype=tf.int32))\n parsed['prev_seq_len'] = parsed['seq_len']\n parsed['start'] = tf.constant(0, dtype=tf.int32)\n parsed['end'] = parsed['seq_len'] - 1\n \n for k in train_raw_features:\n \n data = _parsed[k]\n if k not in ['row_id', 'user_id', 'timestamp', 'prior_question_elapsed_time']:\n data = tf.cast(data, dtype=tf.int32)\n elif k == 'prior_question_elapsed_time':\n data = tf.cast(data, dtype=DTYPE)\n if k != 'user_id':\n parsed[k] = data\n \n # We need to use `START_TOKEN` rather than `PAD_TOKEN` here.\n parsed['shifted_answered_correctly'] = tf.concat([[START_TOKEN], parsed['answered_correctly']], axis=0)\n\n # This should be all `0`.\n pad_mask = tf.cast(parsed['row_id'] == PAD_TOKEN, dtype=tf.int32)\n\n # The tfrecord dataset is only used for training (excluding the part used for validation), so no place should be in the prediction time.\n pred_time_mask = tf.zeros_like(parsed['timestamp'], dtype=tf.int32)\n \n # This should be all `0`.\n pred_time_mask = pred_time_mask * (1 - pad_mask) + (PAD_TOKEN) * pad_mask\n \n parsed['pred_time_mask'] = pred_time_mask\n\n parsed['abs_pos'] = tf.range(parsed['seq_len'], dtype=tf.int32)\n parsed['shifted_abs_pos'] = tf.concat([[START_TOKEN], parsed['abs_pos'][:-1]], axis=0)\n \n return parsed\n\n\n# --------------------------------------------------------------------------------\n# For validation dataset\n\ntrain_features_with_valid_info = {\n 'user_id': tf.io.FixedLenFeature([], dtype=tf.int64),\n 'row_id': tf.io.RaggedFeature(value_key='row_id', dtype=tf.int64),\n 'timestamp': tf.io.RaggedFeature(value_key='timestamp', dtype=tf.int64),\n 'content_id': tf.io.RaggedFeature(value_key='content_id', dtype=tf.int64),\n 'content_type_id': tf.io.RaggedFeature(value_key='content_type_id', dtype=tf.int64),\n 'task_container_id': tf.io.RaggedFeature(value_key='task_container_id', dtype=tf.int64),\n 'user_answer': tf.io.RaggedFeature(value_key='user_answer', dtype=tf.int64),\n 'answered_correctly': tf.io.RaggedFeature(value_key='answered_correctly', dtype=tf.int64),\n 'prior_question_elapsed_time': tf.io.RaggedFeature(value_key='prior_question_elapsed_time', dtype=tf.float32),\n 'prior_question_had_explanation': tf.io.RaggedFeature(value_key='prior_question_had_explanation', dtype=tf.int64),\n 'n_valid_blocks': tf.io.FixedLenFeature([], dtype=tf.int64),\n 'valid_blocks_start_pos': tf.io.RaggedFeature(value_key='valid_blocks_start_pos', dtype=tf.int64),\n 'valid_blocks_end_pos': tf.io.RaggedFeature(value_key='valid_blocks_end_pos', dtype=tf.int64),\n}\n\n\ndef parse_train_example_with_valid_info(example):\n\n _parsed = tf.io.parse_single_example(example, train_features_with_valid_info)\n \n parsed = {}\n parsed['user_id'] = _parsed['user_id']\n parsed['seq_len'] = tf.reduce_sum(tf.ones_like(_parsed['row_id'], dtype=tf.int32))\n parsed['prev_seq_len'] = parsed['seq_len']\n parsed['start'] = tf.constant(0, dtype=tf.int32)\n parsed['end'] = parsed['seq_len'] - 1\n \n for k in train_features_with_valid_info:\n \n data = _parsed[k]\n if k not in ['row_id', 'user_id', 'timestamp', 'prior_question_elapsed_time']:\n data = tf.cast(data, dtype=tf.int32)\n elif k == 'prior_question_elapsed_time':\n data = tf.cast(data, dtype=DTYPE)\n if k != 'user_id':\n parsed[k] = data\n\n # We need to use `START_TOKEN` rather than `PAD_TOKEN` here.\n parsed['shifted_answered_correctly'] = tf.concat([[START_TOKEN], parsed['answered_correctly']], axis=0) \n\n # This should be all `0`.\n pad_mask = tf.cast(parsed['row_id'] == PAD_TOKEN, dtype=tf.int32)\n\n # The tfrecord dataset, in the raw format, contains training examples with validation information.\n # At this step, no place is consider to be in the prediction time yet.\n # This information will be updated in a further dataset transformation.\n pred_time_mask = tf.zeros_like(parsed['timestamp'], dtype=tf.int32)\n \n # This should be all `0`.\n pred_time_mask = pred_time_mask * (1 - pad_mask) + (PAD_TOKEN) * pad_mask\n \n parsed['pred_time_mask'] = pred_time_mask\n\n parsed['abs_pos'] = tf.range(parsed['seq_len'], dtype=tf.int32)\n parsed['shifted_abs_pos'] = tf.concat([[START_TOKEN], parsed['abs_pos'][:-1]], axis=0)\n \n return parsed","execution_count":null,"outputs":[]},{"metadata":{"id":"E1MVqN03a7Bf"},"cell_type":"markdown","source":"#### check"},{"metadata":{"id":"SPsCWH3ja1kZ","outputId":"a152f6d8-69b8-4912-c596-e51735c9c56d","trusted":true},"cell_type":"code","source":"train_raw_ds = tf.data.TFRecordDataset(train_tfrec_paths, num_parallel_reads=1)\ntrain_raw_ds = train_raw_ds.map(parse_train_example, num_parallel_calls=1, deterministic=True)\nfor x in train_raw_ds.take(1):\n print(x)","execution_count":null,"outputs":[]},{"metadata":{"id":"SIG0fVP-LbRc","outputId":"2a1dffc5-2878-4a41-e53d-ee6427def267","trusted":true},"cell_type":"code","source":"valid_raw_ds = tf.data.TFRecordDataset(valid_tfrec_paths, num_parallel_reads=1)\nvalid_raw_ds = valid_raw_ds.map(parse_train_example_with_valid_info, num_parallel_calls=1, deterministic=True)\nfor x in valid_raw_ds.take(1):\n print(x)","execution_count":null,"outputs":[]},{"metadata":{"id":"PdS1fIW511qg"},"cell_type":"markdown","source":"### Split the dataset into training / validation parts\n\nThis is used to remove the interactions used for validation from a user's full interaction history."},{"metadata":{"trusted":true,"id":"MBvizSTo11ql"},"cell_type":"code","source":"def convert_split_index_dict(split_index_dict):\n\n user_ids = []\n split_indices = []\n for k, v in split_index_dict.items():\n user_ids.append(k)\n split_indices.append(v)\n\n user_id_tensor = tf.constant(user_ids, dtype=tf.int64)\n split_index_tensor = tf.constant(split_indices, dtype=tf.int32)\n\n initializer = tf.lookup.KeyValueTensorInitializer(user_id_tensor, split_index_tensor)\n\n split_index_table = tf.lookup.StaticHashTable(\n initializer, default_value=tf.int32.limits[1], name=None\n )\n\n return split_index_table\n\n\ndef split_train_example(raw_example, split_index_table):\n \"\"\"Split an original train example to actual training part and validation part, and only return the training part.\n \"\"\"\n \n user_id = raw_example['user_id']\n seq_len = raw_example['seq_len']\n \n split_index = split_index_table.lookup(user_id)\n tf.debugging.Assert(tf.reduce_all(split_index >= 0), [split_index])\n \n example = {}\n \n # `user_id` not showing in `split_index_table` - not used for validation.\n # `split_index` becomes `seq_len` - i.e. all the interactions belongs to training.\n if split_index == tf.int32.limits[1]:\n split_index = seq_len\n \n example['user_id'] = raw_example['user_id']\n example['seq_len'] = split_index\n example['prev_seq_len'] = raw_example['seq_len']\n example['start'] = tf.constant(0, dtype=tf.int32)\n example['end'] = split_index - 1 \n \n for k in raw_example:\n if k not in ['user_id', 'seq_len', 'prev_seq_len', 'start', 'end']:\n example[k] = raw_example[k][0:split_index]\n \n return example\n\n\ndef split_train_ds(train_raw_ds, split_indices, num_parallel_calls=None, deterministic=None):\n \n reduced_raw_ds = train_raw_ds.map(lambda example: split_train_example(example, split_indices), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n \n return reduced_raw_ds","execution_count":null,"outputs":[]},{"metadata":{"id":"KqpuGiNc11q5"},"cell_type":"markdown","source":"### Training dataset transformation - from tf.RaggedTensor to tf.Tensor\n\nThese are helper functions used for sampling random subsequences (of a fixed length) of user interactions (in the splitted training dataset) for training."},{"metadata":{"trusted":true,"id":"WrZZVg2G11q9"},"cell_type":"code","source":"def extract_subseqs(seqs, ending_indices, window_size, seq_len):\n\n \"\"\"Let `seqs` be a `tf.RaggedTensor` be a tensor with rank = 2, where the 1st and 2nd dimensions are\n batch dimension and temporal and the unique ragged dimesion. Let `ending_indices` be a 1-D tensor\n with the same batch dimension as `seqs`. The condition `-1 <= ending_indices[i] < len(seqs[i])`\n must holds.\n\n For each example `seqs[i]`, we extract a partial sequence\n `seqs[ending_indices[i] - window_size : ending_indices[i]]`. If `len(seqs) < window_size`, the invalid\n indices will get `PAD_TOKEN` as values.\n \n Args:\n seqs: A `tf.RaggedTensor` tensor with rank = 2. The 1st dim is the batch dimension,\n and the 2nd dim is the temporal dimension. The temporal dimension is the ragged rank,\n i.e. the unique dimension which is ragged.\n \n ending_indices: A 1-D `tf.int32` tensor with shape = [ragged_tensor.shape[0]]\n \n window_size: A scalar `tf.int32` tensor, which should be positive.\n \"\"\"\n \n orig_window_size = window_size\n\n window_size += 1\n\n tf.debugging.Assert(tf.rank(seqs) >= 2, [tf.rank(seqs)])\n \n tf.debugging.Assert(tf.reduce_all(window_size > 0), [window_size])\n tf.debugging.Assert(tf.reduce_all(ending_indices >= -1), [ending_indices])\n \n batch_size = tf.reduce_sum(tf.ones_like(seq_len, dtype=tf.int32))\n \n tf.debugging.Assert(tf.reduce_all((ending_indices < seq_len)), [ending_indices, seq_len])\n \n # shape = [batch_size, 1]\n _ending_indices = ending_indices[:, tf.newaxis]\n \n # add by 1 because we will add ...\n _ending_indices += 1\n\n # shape = [1, window_size]\n ranges = tf.range(window_size)[tf.newaxis, :]\n \n # shape = [batch_size, window_size]\n indices = _ending_indices - ranges\n \n # reverse the sequence dimension to get the correct temporal direction\n # shape = [batch_size, window_size]\n indices = tf.reverse(indices, axis=[1])\n \n # shape = [batch_size, window_size, 2]\n indices_to_seqs = tf.stack([tf.broadcast_to(tf.range(batch_size)[:, tf.newaxis], shape=[batch_size, window_size]), indices], axis=2)\n \n # Change negative indices to 0\n # shape = [batch_size, window_size, 2]\n indices_to_ext_seqs = tf.math.maximum(indices_to_seqs, 0)\n \n # Need to rework\n invalid_values = tf.cast(tf.ones(shape=seqs.shape[2:]) * PAD_TOKEN, seqs.dtype)\n \n # shape = [batch_size, 1, ...]\n invalid_values = tf.repeat([[invalid_values]], repeats=batch_size, axis=0)\n \n # Same shape as `ragged_tensor`, but each sequence has 1 more element inserted at the beginning\n extended_seqs = tf.concat([invalid_values, seqs], axis=1)\n \n # selected\n # shape = [batch_size, window_size, ...]\n sampled_subseqs = tf.gather_nd(extended_seqs, indices=indices_to_ext_seqs)\n\n # take the last `orig_window_size` part\n sampled_subseqs = sampled_subseqs[:, 1:]\n prev_last_element = sampled_subseqs[:, 0]\n \n return sampled_subseqs, prev_last_element\n\n\ndef random_ending_indices(seq_len, seed=None):\n\n batch_size = tf.reduce_sum(tf.ones_like(seq_len))\n \n max_seq_len = tf.cast(tf.math.reduce_max(seq_len), dtype=DTYPE)\n \n # shape = [batch_size]\n ending_indices = -1.0 + tf.random.uniform(minval=0.0, maxval=1.0, dtype=DTYPE, shape=[batch_size], seed=seed) * tf.cast(seq_len + 1, dtype=DTYPE)\n ending_indices = tf.cast(tf.math.floor(ending_indices), dtype=tf.int32) \n\n # sanity check\n tf.debugging.Assert(tf.math.reduce_min(seq_len - ending_indices) >= 1, [seq_len, ending_indices]) \n tf.debugging.Assert(tf.math.reduce_min(ending_indices) >= -1, [seq_len, ending_indices]) \n \n return ending_indices\n \n \ndef random_subseqs(seqs, window_size, seq_len, seed=None):\n \"\"\"Let `seqs` be a `tf.RaggedTensor` be a tensor with rank = 2, where the 1st and 2nd dimensions are\n batch dimension and temporal and the unique ragged dimesion.\n \n For each example `seqs[i]`, we extract a random partial sequence of length <= window_size\n \n Args:\n seqs: A `tf.RaggedTensor` tensor with rank = 2. The 1st dim is the batch dimension,\n and the 2nd dim is the temporal dimension. The temporal dimension is the ragged rank,\n i.e. the unique dimension which is ragged.\n \n window_size: A scalar `tf.int32` tensor, which should be positive.\n \"\"\"\n\n ending_indices = random_ending_indices(seq_len=seq_len, seed=seed)\n sampled_subseqs, _ = extract_subseqs(seqs, ending_indices, window_size, seq_len=seq_len)\n \n return sampled_subseqs, ending_indices\n\n# --------------------------------------------------------------------------------\n# For train / validation\n\ndef extract_subseqs_from_raw_batch(raw_batch, ending_indices, window_size):\n \"\"\"Extract subsequences from a training batch (containing no validation part anymore) consisting of `tf.RaggedTensor` objects.\n The subsequences will have the same length `window_size` by padding from the beginning with `PAD_TOKEN`.\n \"\"\" \n\n tf.debugging.Assert(tf.reduce_all(raw_batch['start'] == tf.zeros_like(raw_batch['start'], dtype=tf.int32)), [raw_batch['start']])\n \n batch = {}\n batch['user_id'] = raw_batch['user_id']\n\n start = tf.math.maximum(ending_indices - window_size + 1, 0)\n end = ending_indices\n seq_len = (end - start) + 1\n \n batch['seq_len'] = seq_len\n batch['prev_seq_len'] = raw_batch['seq_len'] \n batch['start'] = start\n batch['end'] = end\n \n for k in raw_batch:\n if k not in ['user_id', 'seq_len', 'prev_seq_len', 'start', 'end'] + ['n_valid_blocks', 'valid_blocks_start_pos', 'valid_blocks_end_pos', 'valid_block_pos', 'valid_block_idx', 'valid_start', 'valid_end']:\n \n subseqs, prev_last_element = extract_subseqs(seqs=raw_batch[k], ending_indices=ending_indices, window_size=window_size, seq_len=raw_batch['prev_seq_len'])\n \n batch[k] = subseqs\n\n elif k in ['n_valid_blocks', 'valid_block_idx', 'valid_start', 'valid_end']:\n batch[k] = raw_batch[k]\n\n return batch\n\n# --------------------------------------------------------------------------------\n\ndef random_subseqs_from_raw_batch(raw_batch, window_size, only_last, seed=None):\n \n if only_last:\n ending_indices = raw_batch['seq_len'] - 1\n else:\n ending_indices = random_ending_indices(raw_batch['seq_len'], seed=seed)\n ending_indices = tf.math.maximum(ending_indices, window_size - 1)\n _mask = tf.cast(ending_indices >= raw_batch['seq_len'], tf.int32)\n ending_indices = ending_indices * (1 - _mask) + (raw_batch['seq_len'] - 1) * _mask\n\n batch = extract_subseqs_from_raw_batch(raw_batch, ending_indices, window_size)\n \n return batch, ending_indices\n\n\ndef random_subseqs_from_batched_raw_ds(batched_raw_ds, window_size, only_last=False, seed=None, num_parallel_calls=None, deterministic=None):\n \n batched_ds = batched_raw_ds.map(lambda raw_batch: random_subseqs_from_raw_batch(raw_batch, window_size, only_last, seed), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n \n return batched_ds","execution_count":null,"outputs":[]},{"metadata":{"id":"MNp9cUoC11rN"},"cell_type":"markdown","source":"## Convert the dataset to model inputs"},{"metadata":{},"cell_type":"markdown","source":"#### helper functions"},{"metadata":{"id":"na7ayJgW9NuK","trusted":true},"cell_type":"code","source":"def convert_valid_info(valid_info):\n\n data = {}\n\n user_ids = list(valid_info.keys())\n\n for user_id in user_ids:\n\n user_valid_info = valid_info[user_id]\n\n user_valid_info['bundle_info']['block_starting_index_dict'] = {int(k): v for k, v in user_valid_info['bundle_info']['block_starting_index_dict'].items()}\n user_valid_info['bundle_info']['bundle_starting_index_dict'] = {int(k): v for k, v in user_valid_info['bundle_info']['bundle_starting_index_dict'].items()}\n data[int(user_id)] = user_valid_info\n\n del valid_info[user_id]\n\n return data\n\ndef load_data(path):\n\n if path.endswith('.json'):\n \n with open(path, 'r', encoding='UTF-8') as fp:\n data = json.load(fp)\n\n if path.endswith('valid_info_fold_1.json'):\n data = convert_valid_info(data)\n\n elif path.endswith('n_contents_dict.json'):\n data = {int(k): v for k, v in data.items()}\n \n elif path.endswith('train_valid_split_indices_fold_1.json'): \n data = {int(k): v for k, v in data.items()}\n\n elif path.endswith('question_tags_info.json') or path.endswith('lecture_tags_info.json'): \n data = {int(k): v for k, v in data.items()} \n\n elif path.endswith('question_part_info.json') or path.endswith('lecture_part_info.json'): \n data = {int(k): v for k, v in data.items()} \n\n return data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Build some table / information that are used for mapping ids in the dataset (consider them as tokens) to model input ids.\n\nWe use `tf.lookup.StaticHashTable` to build the mappings (for questions / lectures) instead of Python dictionary, so we can perform the transformation with `tf.data.Dataset`.\n\nFor tags and parts respectively, we just use a single tensor to store the information, and use `tf.gather` to extract the information."},{"metadata":{"trusted":true,"id":"j5VohfqM11rW"},"cell_type":"code","source":"questioin_df = pd.read_csv(f'{BASE_DIR}/riiid-test-answer-prediction/questions.csv')\nlecture_df = pd.read_csv(f'{BASE_DIR}/riiid-test-answer-prediction/lectures.csv')\n\nquestion_ids = questioin_df['question_id'].tolist()\nlecture_ids = lecture_df['lecture_id'].tolist()\n\nquestion_ids = sorted(question_ids)\nlecture_ids = sorted(lecture_ids)\n\nassert question_ids == sorted(question_ids)\nassert lecture_ids == sorted(lecture_ids)\n\nspecial_vocab = [PAD_TOKEN, START_TOKEN, END_TOKEN, MASK_TOKEN]\nquestion_vocab = question_ids\nlecture_vocab = lecture_ids\n\ncontent_vocab = special_vocab + question_vocab + lecture_vocab\n\n# Map question ids to encoder's input ids\ninput_ids = tf.range(len(special_vocab), len(special_vocab) + len(question_vocab))\ninitializer = tf.lookup.KeyValueTensorInitializer(tf.constant(question_ids, dtype=tf.int32), input_ids)\nquestion_id_to_input_id_table = tf.lookup.StaticHashTable(\n initializer, default_value=tf.int32.limits[1], name=None\n)\n\n# Map lecture ids to encoder's input ids\ninput_ids = tf.range(len(special_vocab) + len(question_vocab), len(content_vocab))\ninitializer = tf.lookup.KeyValueTensorInitializer(tf.constant(lecture_ids, dtype=tf.int32), input_ids)\nlecture_id_to_input_id_table = tf.lookup.StaticHashTable(\n initializer, default_value=tf.int32.limits[1], name=None\n)\n\nresponse_vocab = [PAD_TOKEN, START_TOKEN, END_TOKEN, MASK_TOKEN, RESPONSE_LECTURE_TOKEN, RESPONSE_TRUE_ID, RESPONSE_FALSE_ID]\n\nCONTENT_VOCAB_SIZE = len(content_vocab)\nRESPONSE_VOCAB_SIZE = len(response_vocab)\n\nquestion_tags_info = load_data(question_tags_info_path)\nlecture_tags_info = load_data(lecture_tags_info_path)\n\nquestion_part_info = load_data(question_part_info_path)\nlecture_part_info = load_data(lecture_part_info_path)\n\ntags_database = []\npart_database = []\n\nfor idx in range(len(content_vocab)):\n if idx < len(special_vocab):\n tags_database.append([-1] * N_TAGS_PER_CONTENT)\n part_database.append(-1)\n elif idx < len(special_vocab) + len(question_vocab):\n index_in_questions_ids = idx - len(special_vocab)\n question_id = question_ids[index_in_questions_ids]\n tags_database.append(question_tags_info[question_id])\n part_database.append(question_part_info[question_id])\n elif idx < len(content_vocab):\n index_in_lecture_ids = idx - (len(special_vocab) + len(question_vocab))\n lecture_id = lecture_ids[index_in_lecture_ids]\n tags_database.append(lecture_tags_info[lecture_id])\n part_database.append(lecture_part_info[lecture_id])\n\n# The tags and part here are not the ids for model inputs, but the ids defined in the questions.csv / lectures.csv\nc_inputs_ids_to_tags = tf.constant(tags_database, dtype=tf.int32)\nc_inputs_ids_to_part = tf.constant(part_database, dtype=tf.int32)","execution_count":null,"outputs":[]},{"metadata":{"id":"9Flv1sZr07CC","outputId":"576f430b-6525-468b-e1fb-373de988601d","trusted":true},"cell_type":"code","source":"print(f'CONTENT_VOCAB_SIZE: {CONTENT_VOCAB_SIZE}')\nprint(f'RESPONSE_VOCAB_SIZE: {RESPONSE_VOCAB_SIZE}')","execution_count":null,"outputs":[]},{"metadata":{"id":"aLVrl0Zy11rr"},"cell_type":"markdown","source":"### Prepare encoder / decoder input ids and attention masks\n\n### special masks\n\nWe provide the `causal attention mask` and `attention mask that can attend to the current and previous timestamp` - rembember that we can have a bunlde of questions that share the same timestamp."},{"metadata":{"trusted":true,"id":"u-JlqvWq11rv"},"cell_type":"code","source":"def get_causal_attention_mask(nd, ns, dtype, only_before):\n \"\"\"\n 1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),\n -1, ns-nd), but doesn't produce garbage on TPUs.\n \"\"\"\n \n # Remark: Think `nd` as the number of queries and `ns` as the number of keys.\n # In encoder-decoder case, the queries are the decoder features and the keys are the encoder features.\n \n i = tf.range(nd)[:, tf.newaxis] # repeat along dim 1\n j = tf.range(ns) # repeat along dim 0 \n m = i >= (j - ns + nd) + tf.cast(only_before, dtype=tf.int32)\n \n return tf.cast(m, dtype)\n\ndef get_attention_mask_from_timestamp_batch(timestamp_tensors, dtype, only_before):\n \"\"\"\n Args:\n timestamp_tensors: 2-D tf.int32 tensor, representing a batch of sequences of non-decreasing timestamps.\n \n Returns:\n attention_mask: 3-D tf.int32 tensor of shape = [batch_size, query_len, key_len], consisting of 0 and 1.\n Here `query_len` and `key_len` are actually `seq_len`. It should be reshpaed, when used to calculate \n attention scores, to [batch_size, nb_attn_head, query_len, key_len].\n \"\"\"\n \n t = timestamp_tensors\n \n batch_size = tf.math.reduce_sum(tf.ones_like(t[:, :1], dtype=tf.int32))\n seq_len = tf.math.reduce_sum(tf.ones_like(t[:1, :], dtype=tf.int32))\n \n x = tf.broadcast_to(t[:, :, tf.newaxis], shape=[batch_size, seq_len, seq_len]) # repeat along dim 2\n y = tf.broadcast_to(t[:, tf.newaxis, :], shape=[batch_size, seq_len, seq_len]) + tf.cast(only_before, dtype=tf.int64) # repeat along dim 1\n \n m = x >= y\n \n return tf.cast(m, dtype)","execution_count":null,"outputs":[]},{"metadata":{"id":"Xenht7RU11r9"},"cell_type":"markdown","source":"##### check"},{"metadata":{"trusted":true,"id":"qLweq7ry11sA","outputId":"a4f2504d-0b40-4a6a-dba9-f3f2d524e1c4"},"cell_type":"code","source":"print(get_causal_attention_mask(3, 3, tf.int32, only_before=tf.constant(False)))\nprint(get_causal_attention_mask(3, 3, tf.int32, only_before=tf.constant(True)))\n\ntimestamp_tensors = tf.constant([[0, 0, 1, 2, 2, 2, 3, 3], [0, 1, 1, 2, 2, 3, 3, 3]], dtype=tf.int64)\n\nprint(get_attention_mask_from_timestamp_batch(timestamp_tensors, tf.int32, only_before=tf.constant(False)))\nprint(get_attention_mask_from_timestamp_batch(timestamp_tensors, tf.int32, only_before=tf.constant(True)))","execution_count":null,"outputs":[]},{"metadata":{"id":"MFxIW3spKXFU"},"cell_type":"markdown","source":"### Tensors required for model (input ids, targets, attention masks)\n\nTo be continued ..."},{"metadata":{"trusted":true,"id":"NfG8YPPS11sU"},"cell_type":"code","source":"@tf.function\ndef get_attention_masks(batch, training, encoder_decoder, generative, allow_bundle_atten):\n\n pad_mask = tf.cast((batch['content_type_id'] == PAD_TOKEN), dtype=tf.int32)\n question_mask = tf.cast((batch['content_type_id'] == 0), dtype=tf.int32)\n\n # Replace `PAD_TOKEN` by `0`.\n pred_time_mask = batch['pred_time_mask'] * tf.cast(batch['pred_time_mask'] != PAD_TOKEN, dtype=tf.int32)\n\n content_ids = batch['content_id']\n\n # sanity check\n _pad_mask = tf.cast((batch['timestamp'] == PAD_TOKEN), dtype=tf.int32)\n tf.debugging.Assert(tf.reduce_all(pad_mask == _pad_mask), [pad_mask, _pad_mask]) \n\n # ----------------------------------------\n # General mask\n # `seq_len` below is actually the `window_size`.\n \n seq_len = tf.math.reduce_sum(tf.ones_like(content_ids[0, :], dtype=tf.int32))\n \n # Don't pay attention to [PAD]\n # shape = [batch_size, seq_len]\n non_pad_mask = 1 - pad_mask\n \n # Can only attent to the current and previous position\n # shape = [seq_len, seq_len]\n causal_attention_mask = get_causal_attention_mask(nd=seq_len, ns=seq_len, dtype=tf.int32, only_before=tf.constant(False))\n \n # Can only attent to the previous position \n # shape = [seq_len, seq_len]\n causal_attention_mask_only_before = get_causal_attention_mask(nd=seq_len, ns=seq_len, dtype=tf.int32, only_before=tf.constant(True))\n\n # Can only attend to the current previous timestamps\n # shape = [batch_size, seq_len, seq_len]\n timestamp_attention_mask = get_attention_mask_from_timestamp_batch(batch['timestamp'], dtype=tf.int32, only_before=tf.constant(False))\n\n # Can only attend to previous timestamps\n # shape = [batch_size, seq_len, seq_len]\n timestamp_attention_mask_only_before = get_attention_mask_from_timestamp_batch(batch['timestamp'], dtype=tf.int32, only_before=tf.constant(True))\n \n # ----------------------------------------\n # content self attention mask\n # shape = [batch_size, seq_len, seq_len]\n\n if tf.cast(allow_bundle_atten, dtype=tf.bool):\n # can see the contents in the current and previous bundle, not including [PAD]\n c_mask = non_pad_mask[:, tf.newaxis, :] * timestamp_attention_mask\n else:\n # can see the current and previous contents, not including [PAD]\n c_mask = non_pad_mask[:, tf.newaxis, :] * causal_attention_mask[tf.newaxis, :, :]\n \n # ----------------------------------------\n # response self attention mask\n # shape = [batch_size, seq_len, seq_len] \n \n # for `encoder-only` model: same as `c_mask`\n r_mask = c_mask\n\n if tf.cast(encoder_decoder, dtype=tf.bool):\n \n # for `encoder-decoder` model: should be causal masking\n d_mask = non_pad_mask[:, tf.newaxis, :] * causal_attention_mask[tf.newaxis, :, :]\n\n # if neither `training` nor `generative`\n if (1 - training) * (1 - generative) == 1:\n\n pred_time_question_mask = pred_time_mask * question_mask\n \n d_mask = d_mask * tf.math.maximum((1 - pred_time_question_mask)[:, tf.newaxis, :], tf.eye(seq_len, dtype=tf.int32)[tf.newaxis, :, :])\n r_mask = d_mask\n c_mask = c_mask * tf.math.maximum((1 - pred_time_question_mask)[:, tf.newaxis, :], tf.eye(seq_len, dtype=tf.int32)[tf.newaxis, :, :])\n\n else:\n r_mask = d_mask\n\n # ----------------------------------------\n # response to content to attention_mask\n # shape = [batch_size, seq_len, seq_len]\n \n r_c_mask = c_mask\n\n # ----------------------------------------\n # Used for encoder-only models.\n # Can see only the responses in the previous bundles, not including [PAD]\n\n # shape = [batch_size, seq_len, seq_len]\n c_r_mask = non_pad_mask[:, tf.newaxis, :] * timestamp_attention_mask_only_before\n \n return c_mask, r_mask, r_c_mask, c_r_mask\n\n\ndef add_input_ids_and_targets(batch, training, generative, use_abs_pos):\n \"\"\"Add input ids and targets for training\n \"\"\"\n \n content_ids = batch['content_id']\n content_type_ids = batch['content_type_id']\n answered_correctly = batch['answered_correctly']\n \n question_mask = tf.cast((content_type_ids == 0), dtype=tf.int32)\n lecture_mask = tf.cast((content_type_ids == 1), dtype=tf.int32)\n \n pad_mask = tf.cast((content_type_ids == PAD_TOKEN), dtype=tf.int32)\n _pad_mask = tf.cast((batch['timestamp'] == PAD_TOKEN), dtype=tf.int32)\n tf.debugging.Assert(tf.reduce_all(pad_mask == _pad_mask), [pad_mask, _pad_mask])\n \n # Replace `PAD_TOKEN` by `0`.\n pred_time_mask = batch['pred_time_mask'] * tf.cast(batch['pred_time_mask'] != PAD_TOKEN, dtype=tf.int32)\n pred_time_question_mask = pred_time_mask * question_mask\n \n # The number of questions in prediction time\n # shape = [batch_size]\n n_questions_in_pred_time = tf.math.reduce_sum(pred_time_question_mask, axis=1)\n # sanity check\n tf.debugging.Assert(tf.reduce_all(n_questions_in_pred_time >= 0), [n_questions_in_pred_time]) \n\n response_false_mask = tf.cast((answered_correctly == RESPONSE_FALSE_TOKEN), dtype=tf.int32)\n response_true_mask = tf.cast((answered_correctly == RESPONSE_TRUE_TOKEN), dtype=tf.int32)\n \n # ----------------------------------------\n\n _batch = {}\n for k in batch:\n _batch[k] = batch[k]\n\n if use_abs_pos:\n # Shifted by one, `PAD_TOKEN` --> `0`\n _batch['pos_ids'] = tf.math.maximum(0, _batch['abs_pos'] + 1)\n\n # Shifted by one, `PAD_TOKEN` --> `0`, `START_TOKEN` --> `0`.\n _batch['shifted_pos_ids'] = tf.math.maximum(0, _batch['shifted_abs_pos'] - 1)\n\n else:\n\n _batch['pos_ids'] = tf.math.cumsum(1 - pad_mask, axis=1)\n _batch['shifted_pos_ids'] = tf.math.maximum(0, _batch['pos_ids'] - 1)\n\n # ----------------------------------------\n # content_input_ids \n \n content_input_ids = question_mask * question_id_to_input_id_table.lookup(content_ids) + \\\n lecture_mask * lecture_id_to_input_id_table.lookup(content_ids) + \\\n pad_mask * PAD_ID\n \n _batch['c_input_ids'] = content_input_ids\n \n # ----------------------------------------\n # response_input_ids - only used for encoder-only models\n \n # For `RESPONSE_LECTURE_ID`, we need to multiply by `(1 - pred_time_question_mask) * lecture_mask` instead of just `lecture_mask`.\n # Reason: we might have lectures occur during the prediction time. And these should be assigned to `RESPONSE_MASK_ID`.\n # If we only multiply by `lecture_mask`, we will get `RESPONSE_MASK_ID + RESPONSE_LECTURE_ID` which gives OOV error for embedding.\n \n response_input_ids_masked = pad_mask * PAD_ID + pred_time_question_mask * MASK_ID + lecture_mask * RESPONSE_LECTURE_ID + (1 - pred_time_question_mask) * response_false_mask * RESPONSE_FALSE_ID + (1 - pred_time_question_mask) * response_true_mask * RESPONSE_TRUE_ID\n _batch['r_input_ids'] = response_input_ids_masked\n\n # ----------------------------------------\n # d_input_ids\n\n shifted_answered_correctly = batch['shifted_answered_correctly']\n\n shifted_pad_mask = tf.cast((shifted_answered_correctly == PAD_TOKEN), dtype=tf.int32)\n shifted_start_mask = tf.cast((shifted_answered_correctly == START_TOKEN), dtype=tf.int32)\n\n shifted_masking_mask = tf.cast((shifted_answered_correctly == MASK_TOKEN), dtype=tf.int32)\n shifted_response_lecture_mask = tf.cast((shifted_answered_correctly == RESPONSE_LECTURE_TOKEN), dtype=tf.int32)\n shifted_response_false_mask = tf.cast((shifted_answered_correctly == RESPONSE_FALSE_TOKEN), dtype=tf.int32)\n shifted_response_true_mask = tf.cast((shifted_answered_correctly == RESPONSE_TRUE_TOKEN), dtype=tf.int32)\n\n decoder_input_ids = shifted_pad_mask * PAD_ID + shifted_start_mask * START_ID + shifted_masking_mask * MASK_ID + shifted_response_lecture_mask * RESPONSE_LECTURE_ID + shifted_response_false_mask * RESPONSE_FALSE_ID + shifted_response_true_mask * RESPONSE_TRUE_ID\n\n if not generative:\n\n pred_time_question_start_mask = tf.cast(tf.math.cumsum(pred_time_question_mask, axis=1) == 1, dtype=tf.int32)\n # shape = [batch_size]\n pred_time_question_start_value = tf.math.reduce_sum(pred_time_question_start_mask * decoder_input_ids, axis=1)\n\n # If at the starting of questions, we get `MASK_ID`, we change it to `RESPONSE_LECTURE_ID`\n prev_lecture_mask = tf.cast(pred_time_question_start_value == MASK_ID, tf.int32) \n pred_time_question_start_value = RESPONSE_LECTURE_ID * prev_lecture_mask + pred_time_question_start_value * (1 - prev_lecture_mask) \n\n # All places in prediction time share the values at the prediction question starting places.\n decoder_input_ids = decoder_input_ids * (1 - pred_time_question_mask) + pred_time_question_mask * pred_time_question_start_value[:, tf.newaxis]\n\n # If there are remaining MASK_ID, its previous place is in prediction time, and we are sure it is a lecture at that moment.\n decoder_input_ids = decoder_input_ids * tf.cast(decoder_input_ids != MASK_ID, dtype=tf.int32) + RESPONSE_LECTURE_ID * tf.cast(decoder_input_ids == MASK_ID, dtype=tf.int32)\n\n _batch['d_input_ids'] = decoder_input_ids\n\n # ----------------------------------------\n # post processing `pos_ids` and `shifted_pos_ids`\n # Once the real decoder is implemented, we need to fix this.\n\n if not generative:\n\n pos_ids = _batch['pos_ids']\n shifted_pos_ids = _batch['shifted_pos_ids']\n\n pos_ids_question_start_value = tf.math.reduce_sum(pred_time_question_start_mask * pos_ids, axis=1)\n shifted_pos_ids_question_start_value = tf.math.reduce_sum(pred_time_question_start_mask * shifted_pos_ids, axis=1)\n\n pos_ids = pos_ids * (1 - pred_time_question_mask) + pred_time_question_mask * pos_ids_question_start_value[:, tf.newaxis]\n shifted_pos_ids = shifted_pos_ids * (1 - pred_time_question_mask) + pred_time_question_mask * shifted_pos_ids_question_start_value[:, tf.newaxis]\n\n _batch['pos_ids'] = pos_ids\n _batch['shifted_pos_ids'] = shifted_pos_ids\n\n # ----------------------------------------\n # tags\n\n _batch['tag_ids'] = tf.gather(params=c_inputs_ids_to_tags, indices=_batch['c_input_ids']) + 1 \n \n # ---------------------------------------- \n # part\n\n _batch['part_ids'] = tf.gather(params=c_inputs_ids_to_part, indices=_batch['c_input_ids']) + 1\n\n # ----------------------------------------\n # prior explanation ids - just `prior_question_had_explanation` added by 1.\n # For PAD, we get `-1` but changed to `0`.\n # only used along with `d_input_ids`.\n\n _batch['prior_explanation_ids'] = tf.math.maximum(0, _batch['prior_question_had_explanation'] + 1)\n\n # ----------------------------------------\n # `prior_question_elapsed_time` in seconds\n\n _batch['prior_question_elapsed_time_input'] = tf.cast(_batch['prior_question_elapsed_time'], dtype=DTYPE) / 1000.0\n\n # ----------------------------------------\n # lag_time --> \n\n\n # ----------------------------------------\n # target\n \n # `-2` means padding, `-1` means lecture\n answer_mask = tf.cast(batch['answered_correctly'] > -1, dtype=tf.int32)\n \n # negated values become `NON_TARGET_ID (-100)`\n _batch['target'] = batch['answered_correctly'] * answer_mask + (NON_TARGET_ID) * (1 - answer_mask)\n \n # ----------------------------------------\n # nb_pred_places\n\n targets = _batch['target']\n\n # `targets` are defined for all places (other than [PAD] and lectures).\n # However, during validation, unlike during training, we only focus on the places that are in prediction time (and being questions).\n # This should be used only in `train_step` and `valid_step`, `train` and `valid`, but not in `run_pred`.\n if training == 0:\n pred_time_mask = _batch['pred_time_mask']\n targets = targets * pred_time_mask + NON_TARGET_ID * (1 - pred_time_mask)\n\n pred_mask = targets != NON_TARGET_ID\n nb_pred_places = tf.math.reduce_sum(tf.cast(pred_mask, dtype=tf.int32))\n\n # shape = [batch_size], but it is a constant\n _batch['nb_pred_places'] = nb_pred_places * tf.ones_like(_batch['user_id'], dtype=tf.int32)\n\n return _batch\n\n\ndef prepare_training_dataset(batched_ds, generative=False, use_abs_pos=False, num_parallel_calls=None, deterministic=None):\n # Add input ids and targets for training\n \n training = tf.constant(1, dtype=tf.int32)\n\n return batched_ds.map(lambda batch, _: add_input_ids_and_targets(batch, training, generative, use_abs_pos), num_parallel_calls=num_parallel_calls, deterministic=deterministic)","execution_count":null,"outputs":[]},{"metadata":{"id":"V4YUlTldJ099"},"cell_type":"markdown","source":"### For validation dataset\n\nPreviously, each example in the validation dataset contains the full interaction history of the corresponding user. However, each example for validation should contains only a single block (see [Load TFRecord files](#load-tfrecord-files) for the definition of `block`). Therefore, we perform the following dataset transformation: for each example\n\n1. repeat it `n_valid_blocks` times\n2. for each repeated instance, add an index `n` to indicate it is the `n-th` blocks appeared in the validation time\n3. use the index calculated in 2. and the information stored in `'valid_blocks_start_pos'` and `'valid_blocks_end_pos'` to get the start / end indices in that user interaction history.\n4. use the end indices obtained in 3. to remove the interactions after the current validation block\n5. use the start / end indices to prepare the examples as preparing a training example\n * In particular, we modify the attribute `pred_time_mask` that indicates the places where interactions are in the validation.\n"},{"metadata":{"id":"n13GQqGSJuf6","trusted":true},"cell_type":"code","source":"def add_valid_block_info(valid_raw_example):\n \"\"\"\n - Add `valid_block_pos`.\n - Repeat `n_valid_blocks` times, each with a index `valid_block_idx`.\n - Transform using `trans_2`.\n \"\"\"\n \n example = {}\n for k in valid_raw_example:\n example[k] = valid_raw_example[k]\n \n valid_block_pos = tf.stack([valid_raw_example['valid_blocks_start_pos'], valid_raw_example['valid_blocks_end_pos']], axis=1)\n example['valid_block_pos'] = valid_block_pos\n \n n_valid_blocks = example['n_valid_blocks']\n \n ds_1 = tf.data.Dataset.from_tensors(example).repeat(tf.cast(n_valid_blocks, dtype=tf.int64))\n ds_2 = tf.data.Dataset.range(tf.cast(n_valid_blocks, dtype=tf.int64), output_type=tf.int32)\n ds = tf.data.Dataset.zip((ds_1, ds_2))\n \n ds = ds.map(lambda ex, valid_block_idx: add_valid_block_idx(ex, valid_block_idx))\n\n return ds\n\n\ndef add_valid_block_idx(valid_example, valid_block_idx):\n \"\"\"Add the following information:\n \n - `valid_block_idx`: The index of a block in all the blocks in a user's interaction history that are used for validation.\n - `valid_start`: The starting indices of a valid block in a user's (before being splitted) training interaction history.\n - `valid_end`: The ending indices of the a valid block in a user's (before being splitted) training interaction history.\n \n Then call `remove_future_valid_blocks()` to remove validation blocks after the current one.\n \n \"\"\"\n \n valid_example['valid_block_idx'] = valid_block_idx\n valid_example['valid_start'] = valid_example['valid_block_pos'][valid_block_idx][0]\n valid_example['valid_end'] = valid_example['valid_block_pos'][valid_block_idx][1]\n \n return valid_example\n\n\ndef remove_future_valid_blocks(valid_example):\n \"\"\"Remove the validation blocks in the full interaction history of a user after the current validation block.\n \"\"\"\n \n example = {}\n \n valid_start = valid_example['valid_start']\n valid_end = valid_example['valid_end']\n \n example['user_id'] = valid_example['user_id']\n example['seq_len'] = valid_end + 1\n example['prev_seq_len'] = valid_example['seq_len']\n example['start'] = tf.constant(0, dtype=tf.int32)\n example['end'] = valid_end\n \n for k in valid_example:\n \n if k not in ['user_id', 'seq_len', 'prev_seq_len', 'start', 'end']:\n \n if k in ['n_valid_blocks', 'valid_blocks_start_pos', 'valid_blocks_end_pos', 'valid_block_pos', 'valid_block_idx', 'valid_start', 'valid_end']:\n # attributes with single value, or the values are not required to be removed\n example[k] = valid_example[k]\n else:\n # attributes with \n example[k] = valid_example[k][0:valid_end + 1]\n \n if k == 'pred_time_mask':\n # Update `pred_time_mask` - assign `1` to the current validation places.\n \n n_valid_interactions = valid_end - valid_example['valid_start'] + 1\n example[k] = tf.concat([valid_example[k][:valid_start], tf.ones(shape=[n_valid_interactions], dtype=tf.int32)], axis=0)\n \n return example\n\n\ndef extract_ending_subseqs(raw_batch, window_size):\n \n ending_indices = raw_batch['seq_len'] - 1\n\n return extract_subseqs_from_raw_batch(raw_batch, ending_indices, window_size)\n \n\ndef prepare_validation_dataset(valid_raw_ds, batch_size=3, window_size=5, generative=False, use_abs_pos=False, seed=None, num_parallel_calls=None, deterministic=None):\n \n valid_ds = valid_raw_ds.flat_map(lambda valid_raw_example: add_valid_block_info(valid_raw_example))\n valid_ds = valid_ds.map(lambda example: remove_future_valid_blocks(example))\n \n # should be outside\n # batch examples with attributes having different lengths across examples - tf.RaggedTensor\n batched_valid_ds = valid_ds.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=batch_size, drop_remainder=(IS_KAGGLE and tpu is not None)))\n\n # batch - tf.Tensor: Extract subsequences from the ending of a fixed length \n batched_valid_ds = batched_valid_ds.map(lambda raw_batch: extract_ending_subseqs(raw_batch, window_size), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n \n training = tf.constant(0, dtype=tf.int32)\n\n valid_ds = batched_valid_ds.map(lambda batch: add_input_ids_and_targets(batch, training, generative, use_abs_pos), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n \n return valid_ds","execution_count":null,"outputs":[]},{"metadata":{"id":"9CwJ1Ki011sl"},"cell_type":"markdown","source":"### check"},{"metadata":{"id":"p4L5uloHLnWO","trusted":true},"cell_type":"code","source":"train_raw_ds = tf.data.TFRecordDataset(train_tfrec_paths, num_parallel_reads=1)\ntrain_raw_ds = train_raw_ds.map(parse_train_example, num_parallel_calls=1, deterministic=True)\n\n# Get the splitted training dataset.\ntrain_valid_split_indices = load_data(train_valid_split_indices_path)\ntrain_valid_split_table = convert_split_index_dict(train_valid_split_indices)\nreduced_raw_ds = split_train_ds(train_raw_ds, train_valid_split_table, num_parallel_calls=1, deterministic=True)\n\n# batch examples with attributes having different lengths across examples - tf.RaggedTensor\nbatched_raw_ds = reduced_raw_ds.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=5))\n\n# batch - tf.Tensor: Extract random subsequences of a fixed length\nbatched_ds = random_subseqs_from_batched_raw_ds(batched_raw_ds, window_size=10, seed=1, num_parallel_calls=1, deterministic=True)\n\n# Add input ids and targets for training\ntrain_ds = prepare_training_dataset(batched_ds, use_abs_pos=False, num_parallel_calls=1, deterministic=True)","execution_count":null,"outputs":[]},{"metadata":{"id":"PGQBbKGprMMs","outputId":"0b3e07e3-085e-493f-eb3a-bce97b419ae6","trusted":true},"cell_type":"code","source":"it = iter(train_ds.take(2))\nnext(it)","execution_count":null,"outputs":[]},{"metadata":{"id":"GCO3h1YIL59c","trusted":true},"cell_type":"code","source":"valid_raw_ds = tf.data.TFRecordDataset([valid_tfrec_paths], num_parallel_reads=1)\nvalid_raw_ds = valid_raw_ds.map(parse_train_example_with_valid_info, num_parallel_calls=1, deterministic=True)\nvalid_ds = prepare_validation_dataset(valid_raw_ds, batch_size=5, window_size=10, use_abs_pos=False)","execution_count":null,"outputs":[]},{"metadata":{"id":"zjRpC3h0wvyM","outputId":"a08a4546-08f2-4feb-d7a5-862ce42084bf","trusted":true},"cell_type":"code","source":"it = iter(valid_ds.take(2))\nnext(it)","execution_count":null,"outputs":[]},{"metadata":{"id":"6FMsKoD-11s_"},"cell_type":"markdown","source":"## Model definition"},{"metadata":{},"cell_type":"markdown","source":"### Layer initializer and some activation functions\n\nSince we use the `Glorot uniform initializer` (also called `Xavier uniform initializer`) which are used in the papers [SAINT](#https://arxiv.org/pdf/2002.07033.pdf) and [SAINT+](#https://arxiv.org/pdf/2010.12042.pdf), we need to overwrite some layers' `__init__` methods, which usually contain activation functions.\nWe copy some of them from Hugging Face `transformer` library, [see this link](#https://github.com/huggingface/transformers/blob/master/src/transformers/activations_tf.py)."},{"metadata":{"id":"7okmf9y8qYX5","trusted":true},"cell_type":"code","source":"def get_initializer(seed):\n\n return tf.keras.initializers.GlorotUniform(seed=seed)\n\n\ndef gelu(x):\n \"\"\"\n Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when\n initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):\n 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see\n https://arxiv.org/abs/1606.08415\n \"\"\"\n x = tf.convert_to_tensor(x)\n cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.math.sqrt(2.0), dtype=x.dtype)))\n\n return x * cdf\n\n\ndef gelu_new(x):\n \"\"\"\n Gaussian Error Linear Unit. This is a smoother version of the GELU. Original paper: https://arxiv.org/abs/1606.0841\n Args:\n x: float Tensor to perform activation\n Returns:\n `x` with the GELU activation applied.\n \"\"\"\n x = tf.convert_to_tensor(x)\n pi = tf.cast(math.pi, x.dtype)\n coeff = tf.cast(0.044715, x.dtype)\n cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))\n\n return x * cdf\n\n\ndef gelu_fast(x):\n x = tf.convert_to_tensor(x)\n coeff1 = tf.cast(7978845608, x.dtype)\n coeff2 = tf.cast(0.044715, x.dtype)\n\n return 0.5 * x * (1.0 + tf.tanh(x * coeff2 * (1.0 + coeff1 * x * x)))\n\n\nACT2FN = {\n \"gelu\": tf.keras.layers.Activation(gelu),\n \"relu\": tf.keras.activations.relu,\n \"swish\": tf.keras.activations.swish,\n \"silu\": tf.keras.activations.swish,\n \"gelu_new\": tf.keras.layers.Activation(gelu_new),\n \"tanh\": tf.keras.activations.tanh,\n \"gelu_fast\": tf.keras.layers.Activation(gelu_fast),\n}\n\n\ndef get_tf_activation(activation_string):\n if activation_string in ACT2FN:\n return ACT2FN[activation_string]\n else:\n raise KeyError(\"function {} not found in ACT2FN mapping {}\".format(activation_string, list(ACT2FN.keys())))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model Architecture\n\nI tried the implementations to support:\n\n- encoder only archeticture\n- encoder-decoder archeticture: without auto-regressive generation\n- encoder-decoder archeticture: with auto-regressive generation"},{"metadata":{},"cell_type":"markdown","source":"#### Building blocks\n\nThis cell contains the building blocks (layer definitions), including encoder / decoder. The final model used for training / inference is in the next cell.\n\nThis is a quite large cell and I hide it. You can click `[Code]` button to see the implementation if you are interested.\nAs mentioned above, I copied [distilbert.py](#https://github.com/huggingface/transformers/tree/master/src/transformers/models/distilbert) file from [Hugging Face's transformer library](#https://github.com/huggingface/transformers), and adding (a lot of) my personal modification.\n\nTo not make the notebook become even larger, and due to the time constraint, I removed the original doc for the methods, and added the new doc for only a few methods."},{"metadata":{"id":"lbkDSH67xuNY","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class EdFormerConfig(PretrainedConfig):\n\n def __init__(\n self,\n model_type,\n model_desc,\n model_size='none',\n content_vocab_size=CONTENT_VOCAB_SIZE,\n response_vocab_size=RESPONSE_VOCAB_SIZE,\n tag_vocab_size=TAG_VOCAB_SIZE,\n part_vocab_size=PART_VOCAB_SIZE,\n prior_explanation_vocab_size=PRIOR_EXPLANATION_VOCAB_SIZE,\n max_position_embeddings=MAX_HISTORY_LEN,\n sinusoidal_pos_embds=False,\n n_layers=4,\n n_heads=8,\n dim=512,\n hidden_dim=4 * 512,\n activation=ACTIVATION, \n dropout=0.1,\n attention_dropout=0.1,\n seq2seq_dropout=0.1,\n initializer_range=0.02,\n seed=SEED, \n pad_token_id=PAD_ID,\n use_abs_pos=USE_ABS_POS,\n share_position_embeddings=SHARE_POS_EMBEDDING,\n use_tags=USE_TAGS,\n use_part=USE_PART,\n use_prior_explanation=USE_PRIOR_EXPLANATION,\n use_prior_question_elapsed_time_input=USE_PRIOR_QUESTION_ELAPSED_TIME_INPUT,\n allow_bundle_atten=ALLOW_BUNDLE_ATTEN,\n generative=GENERATIVE,\n **kwargs\n ):\n super().__init__(**kwargs, pad_token_id=pad_token_id)\n \n self.model_type = model_type\n self.model_size = model_size\n self.model_desc = model_desc\n\n self.content_vocab_size = content_vocab_size\n self.response_vocab_size = response_vocab_size\n\n self.tag_vocab_size = tag_vocab_size\n self.part_vocab_size = part_vocab_size\n self.prior_explanation_vocab_size = prior_explanation_vocab_size\n\n self.use_abs_pos = use_abs_pos\n self.share_position_embeddings = share_position_embeddings\n self.use_tags = use_tags\n self.use_part = use_part\n self.use_prior_explanation = use_prior_explanation\n self.allow_bundle_atten = allow_bundle_atten\n self.generative = generative\n self.use_prior_question_elapsed_time_input = use_prior_question_elapsed_time_input\n\n self.max_position_embeddings = max_position_embeddings\n self.sinusoidal_pos_embds = sinusoidal_pos_embds\n \n self.n_layers = n_layers\n self.n_heads = n_heads\n self.dim = dim\n self.hidden_dim = hidden_dim\n\n self.activation = activation\n\n self.dropout = dropout\n self.attention_dropout = attention_dropout\n self.seq2seq_dropout = seq2seq_dropout\n\n self.initializer_range = initializer_range\n self.seed = seed\n\n def toJSON(self):\n\n return json.dumps(self, default=lambda o: o.__dict__, sort_keys=False, indent=4)\n\n def vocab_size(self, input_name):\n\n if input_name == 'content':\n return self.content_vocab_size\n elif input_name == 'response':\n return self.response_vocab_size\n elif input_name == 'tag':\n return self.tag_vocab_size\n elif input_name == 'part':\n return self.part_vocab_size \n elif input_name == 'prior_explanation':\n return self.prior_explanation_vocab_size\n else:\n raise ValueError('input name not used for model')\n\n @property\n def hidden_size(self):\n return self.dim\n\n @property\n def num_attention_heads(self):\n return self.n_heads\n\n @property\n def num_hidden_layers(self):\n return self.n_layers\n\n\nclass TFSharedEmbeddings(tf.keras.layers.Layer):\n\n def __init__(self, vocab_size, hidden_size, seed=None, **kwargs):\n \n super().__init__(**kwargs)\n\n self.vocab_size = vocab_size\n self.hidden_size = hidden_size\n self.seed = seed\n\n def build(self, input_shape):\n\n self.weight = self.add_weight(\n \"weight\", shape=[self.vocab_size, self.hidden_size], initializer=get_initializer(seed=self.seed)\n )\n super().build(input_shape)\n\n def call(self, inputs: tf.Tensor, mode: str = \"embedding\") -> tf.Tensor:\n\n if mode == \"embedding\":\n return self._embedding(inputs)\n elif mode == \"linear\":\n return self._linear(inputs)\n else:\n raise ValueError(\"mode {} is not valid.\".format(mode))\n\n def _embedding(self, input_ids):\n \"\"\"Applies embedding based on inputs tensor.\"\"\"\n return tf.gather(self.weight, input_ids)\n\n def _linear(self, inputs):\n\n first_dims = shape_list(inputs)[:-1]\n x = tf.reshape(inputs, [-1, self.hidden_size])\n logits = tf.matmul(x, self.weight, transpose_b=True)\n\n return tf.reshape(logits, first_dims + [self.vocab_size])\n\n\nclass TFEmbeddings(tf.keras.layers.Layer):\n\n def __init__(self, config, input_name, position_embeddings_layer=None, **kwargs):\n \n super().__init__(**kwargs)\n \n self.pad_id = config.pad_token_id\n self.input_name = input_name\n self.vocab_size = config.vocab_size(input_name)\n\n self.use_tags = config.use_tags\n self.use_part = config.use_part\n self.use_prior_explanation = config.use_prior_explanation\n self.use_prior_question_elapsed_time_input = config.use_prior_question_elapsed_time_input\n\n self.dim = config.dim\n self.seed = config.seed\n self.word_embeddings = TFSharedEmbeddings(\n self.vocab_size, config.dim, seed=config.seed, name=\"word_embeddings\"\n ) # padding_idx=0)\n \n if position_embeddings_layer is None:\n\n self.position_embeddings = tf.keras.layers.Embedding(\n config.max_position_embeddings,\n config.dim,\n embeddings_initializer=get_initializer(config.seed),\n name=\"position_embeddings\",\n )\n \n else:\n\n self.position_embeddings = position_embeddings_layer\n\n if self.input_name == 'content':\n\n if config.use_tags:\n\n self.tag_embeddings = tf.keras.layers.Embedding(\n config.tag_vocab_size,\n config.dim,\n embeddings_initializer=get_initializer(config.seed),\n name=\"tag_embeddings\",\n )\n\n if config.use_part:\n\n self.part_embeddings = tf.keras.layers.Embedding(\n config.part_vocab_size,\n config.dim,\n embeddings_initializer=get_initializer(config.seed),\n name=\"part_embeddings\",\n )\n\n elif self.input_name == 'response':\n\n if config.use_prior_explanation:\n\n self.prior_explanation_embeddings = tf.keras.layers.Embedding(\n config.prior_explanation_vocab_size,\n config.dim,\n embeddings_initializer=get_initializer(config.seed),\n name=\"prior_explanation_embeddings\",\n )\n\n if config.use_prior_question_elapsed_time_input:\n\n self.prior_question_elapsed_time_embeddings = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"prior_question_elapsed_time_embeddings\",\n )\n\n assert config.activation in [\"relu\", \"gelu\"], \"activation ({}) must be in ['relu', 'gelu']\".format(\n config.activation\n )\n self.activation = get_tf_activation(config.activation)\n\n self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name=\"LayerNorm\")\n self.dropout = tf.keras.layers.Dropout(config.dropout)\n\n def build(self, input_shape):\n \"\"\"Build shared word embedding layer \"\"\"\n \n with tf.name_scope(\"word_embeddings\"):\n # Create and initialize weights. The random normal initializer was chosen\n # arbitrarily, and works well.\n self.word_embeddings = self.add_weight(\n \"weight\", shape=[self.vocab_size, self.dim], initializer=get_initializer(self.seed)\n )\n \n super().build(input_shape)\n\n def call(self, input_ids=None, position_ids=None, tag_ids=None, part_ids=None, prior_explanation_ids=None, prior_question_elapsed_time_input=None, inputs_embeds=None, mode=\"embedding\", training=False):\n \n if mode == \"embedding\":\n return self._embedding(input_ids, position_ids, tag_ids, part_ids, prior_explanation_ids, prior_question_elapsed_time_input, inputs_embeds, training=training)\n elif mode == \"linear\":\n return self._linear(input_ids)\n else:\n raise ValueError(\"mode {} is not valid.\".format(mode))\n\n def _embedding(self, input_ids, position_ids, tag_ids, part_ids, prior_explanation_ids, prior_question_elapsed_time_input, inputs_embeds, training=False):\n \"\"\"\n \"\"\"\n \n assert not (input_ids is None and inputs_embeds is None)\n\n if input_ids is not None:\n seq_length = shape_list(input_ids)[1]\n else:\n seq_length = shape_list(inputs_embeds)[1]\n\n if inputs_embeds is None:\n inputs_embeds = tf.gather(self.word_embeddings, input_ids)\n\n if position_ids is None:\n position_ids = tf.range(1, 1 + seq_length, dtype=tf.int32)[tf.newaxis, :]\n\n position_embeddings = tf.cast(\n self.position_embeddings(position_ids), inputs_embeds.dtype\n ) # (bs, max_seq_length, dim)\n\n tag_embeddings = tf.zeros_like(inputs_embeds, dtype=inputs_embeds.dtype)\n if self.use_tags and tag_ids is not None:\n tag_embeddings = tf.cast(\n self.tag_embeddings(tag_ids), inputs_embeds.dtype\n ) # (bs, seq_len, N_TAGS_PER_CONTENT, dim)\n \n # shape = (bs, seq_len, N_TAGS_PER_CONTENT)\n tag_mask = tf.cast(tag_ids != self.pad_id, dtype=tf.int32)\n\n tag_embeddings = tag_embeddings * tf.cast(tag_mask, dtype=inputs_embeds.dtype)[:, :, :, tf.newaxis]\n \n # shape = (bs, seq_len)\n nb_tags = tf.math.reduce_sum(tag_mask, axis=2)\n nb_tags = tf.cast(nb_tags, dtype=inputs_embeds.dtype)\n nb_tags = tf.math.maximum(nb_tags, tf.cast(1.0, dtype=inputs_embeds.dtype))\n\n tag_embeddings = tf.math.reduce_sum(tag_embeddings, axis=2) / nb_tags[:, :, tf.newaxis]\n\n part_embeddings = tf.zeros_like(inputs_embeds, dtype=inputs_embeds.dtype)\n if self.use_part and part_ids is not None:\n part_embeddings = tf.cast(\n self.part_embeddings(part_ids), inputs_embeds.dtype\n ) # (bs, seq_len, dim)\n\n prior_explanation_embeddings = tf.zeros_like(inputs_embeds, dtype=inputs_embeds.dtype)\n if self.use_prior_explanation and prior_explanation_ids is not None:\n prior_explanation_embeddings = tf.cast(\n self.prior_explanation_embeddings(prior_explanation_ids), inputs_embeds.dtype\n ) # (bs, seq_len, dim)\n\n prior_question_elapsed_time_embeddings = tf.zeros_like(inputs_embeds, dtype=inputs_embeds.dtype)\n if self.use_prior_question_elapsed_time_input and prior_question_elapsed_time_input is not None:\n prior_question_elapsed_time_embeddings = self.prior_question_elapsed_time_embeddings(prior_question_elapsed_time_input[:, :, tf.newaxis])\n prior_question_elapsed_time_embeddings = self.activation(prior_question_elapsed_time_embeddings) \n\n embeddings = inputs_embeds + position_embeddings + tag_embeddings + part_embeddings + prior_explanation_embeddings + prior_question_elapsed_time_embeddings # (bs, max_seq_length, dim)\n embeddings = self.LayerNorm(embeddings) # (bs, max_seq_length, dim)\n embeddings = self.dropout(embeddings, training=training) # (bs, max_seq_length, dim)\n \n return embeddings\n\n def _linear(self, inputs):\n \"\"\"\n Computes logits by running inputs through a linear layer\n Args:\n inputs: A float32 tensor with shape [batch_size, length, hidden_size]\n Returns:\n float32 tensor with shape [batch_size, length, vocab_size].\n \"\"\"\n batch_size = shape_list(inputs)[0]\n length = shape_list(inputs)[1]\n\n x = tf.reshape(inputs, [-1, self.dim])\n logits = tf.matmul(x, self.word_embeddings, transpose_b=True)\n\n return tf.reshape(logits, [batch_size, length, self.vocab_size])\n\n\nclass TFMultiHeadSelfAttention(HFTFMultiHeadSelfAttention):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n\n self.q_lin = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"q_lin\"\n )\n self.k_lin = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"k_lin\"\n )\n self.v_lin = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"v_lin\"\n )\n self.out_lin = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"out_lin\"\n )\n\n def call(self, query, key, value, mask, head_mask, output_attentions, training=False):\n \"\"\"\n Parameters:\n query: tf.Tensor(bs, query_length, dim)\n key: tf.Tensor(bs, key_length, dim)\n value: tf.Tensor(bs, key_length, dim)\n mask: tf.Tensor(bs, query_length / 1, key_length)\n Returns:\n weights: tf.Tensor(bs, n_heads, query_length, key_length) Attention weights context: tf.Tensor(bs,\n query_length, dim) Contextualized layer. Optional: only if `output_attentions=True`\n \"\"\"\n bs, q_length, dim = shape_list(query)\n k_length = shape_list(key)[1]\n # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)\n # assert key.size() == value.size()\n dim_per_head = tf.math.divide(self.dim, self.n_heads)\n dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)\n \n def shape(x):\n \"\"\" separate heads \"\"\"\n return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))\n\n def unshape(x):\n \"\"\" group heads \"\"\"\n return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))\n\n q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head)\n k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head)\n v = shape(self.v_lin(value)) # (bs, n_heads, k_length, dim_per_head)\n q = tf.cast(q, dtype=DTYPE)\n q = tf.multiply(q, tf.math.rsqrt(tf.cast(dim_per_head, dtype=DTYPE)))\n k = tf.cast(k, dtype=q.dtype)\n scores = tf.matmul(q, k, transpose_b=True) # (bs, n_heads, q_length, k_length)\n mask = mask[:, tf.newaxis, :, :] # (bs, 1, qlen / 1, klen) --> (bs, n_heads, qlen, klen)\n # scores.masked_fill_(mask, -float('inf')) # (bs, n_heads, q_length, k_length)\n\n mask = tf.cast(mask, dtype=scores.dtype)\n scores = scores - 1e30 * (1.0 - mask)\n weights = tf.nn.softmax(scores, axis=-1) # (bs, n_heads, qlen, klen)\n weights = self.dropout(weights, training=training) # (bs, n_heads, qlen, klen)\n # This makes things more numerically stable.\n weights = weights * mask\n\n # Mask heads if we want to\n if head_mask is not None:\n weights = weights * head_mask\n\n context = tf.matmul(weights, v) # (bs, n_heads, qlen, dim_per_head)\n context = unshape(context) # (bs, q_length, dim)\n context = self.out_lin(context) # (bs, q_length, dim)\n\n if output_attentions:\n return (context, weights)\n else:\n return (context,)\n \n\nclass TFFFN(HFTFFFN):\n\n def __init__(self, config, **kwargs):\n \n super(TFFFN, self).__init__(**kwargs)\n\n self.lin1 = tf.keras.layers.Dense(\n config.hidden_dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"lin1\"\n )\n self.lin2 = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"lin2\"\n )\n\n self.activation = get_tf_activation(config.activation)\n\n\nclass TFTransformerBlock(HFTFTransformerBlock):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n\n self.attention = TFMultiHeadSelfAttention(config, name=\"attention\")\n \n \nclass TFContentBlock(TFTransformerBlock):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n\n self.attention = TFMultiHeadSelfAttention(config, name=\"c_attention\")\n \n\nclass TFResponseBlock(TFTransformerBlock):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n \n self.attention = TFMultiHeadSelfAttention(config, name=\"r_attention\")\n \n self.r_c_attentioin = TFMultiHeadSelfAttention(config, name=\"r_c_attentioin\")\n self.r_c_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name=\"r_c_attn_layer_norm\") \n \n \n def call(self, r, c_hidden, r_mask, r_c_mask, head_mask, output_attentions, training=False): # removed: src_enc=None, src_len=None\n\n r_output = self.attention(r, r, r, r_mask, head_mask, output_attentions, training=training)\n if output_attentions:\n r_output, r_weights = r_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)\n else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples\n # assert type(sa_output) == tuple\n r_output = r_output[0]\n r_output = self.sa_layer_norm(r_output + r) # (bs, seq_length, dim)\n\n r_c_output = self.r_c_attentioin(\n query=r_output,\n key=c_hidden,\n value=c_hidden,\n mask=r_c_mask,\n head_mask=head_mask,\n output_attentions=output_attentions,\n training=training\n )\n \n if output_attentions:\n r_c_output, r_c_weights = r_c_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)\n else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples\n # assert type(sa_output) == tuple\n r_c_output = r_c_output[0] \n r_c_output = self.r_c_attn_layer_norm(r_c_output + r_output)\n \n # Feed Forward Network\n ffn_output = self.ffn(r_c_output, training=training) # (bs, seq_length, dim)\n ffn_output = self.output_layer_norm(ffn_output + r_c_output) # (bs, seq_length, dim)\n\n output = (ffn_output,)\n if output_attentions:\n output = (r_weights, r_c_weights) + output\n return output\n \n def generate(self, current_pos, current_r, prev_c_hidden, prev_d_hidden, r_mask, r_c_mask, head_mask, output_attentions):\n\n r_mask = r_mask[:, current_pos:current_pos+1 ,:]\n r_c_mask = r_c_mask[:, current_pos:current_pos+1, :]\n\n d_hidden = tf.concat([prev_d_hidden[:, :current_pos, :], current_r, prev_d_hidden[:, current_pos+1:, :]], axis=1)\n r_output = self.attention(\n current_r, d_hidden, d_hidden,\n r_mask,\n head_mask,\n output_attentions,\n training=False\n )\n if output_attentions:\n r_output, r_weights = r_output # (bs, 1, dim), (bs, n_heads, 1, seq_length)\n else: \n r_output = r_output[0]\n r_output = self.sa_layer_norm(r_output + current_r) # (bs, 1, dim)\n\n r_c_output = self.r_c_attentioin(\n query=r_output,\n key=prev_c_hidden,\n value=prev_c_hidden,\n mask=r_c_mask,\n head_mask=head_mask,\n output_attentions=output_attentions,\n training=False\n )\n \n if output_attentions:\n r_c_output, r_c_weights = r_c_output # (bs, 1, dim), (bs, n_heads, 1, seq_length)\n else:\n r_c_output = r_c_output[0] \n r_c_output = self.r_c_attn_layer_norm(r_c_output + r_output) # (bs, 1, dim)\n \n # Feed Forward Network\n ffn_output = self.ffn(r_c_output, training=False) # (bs, 1, dim)\n ffn_output = self.output_layer_norm(ffn_output + r_c_output) # (bs, 1, dim)\n\n output = (ffn_output,)\n if output_attentions:\n output = (r_weights, r_c_weights) + output\n return output\n\n\nclass TFCRBlock(TFTransformerBlock):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n\n self.attention = TFMultiHeadSelfAttention(config, name=\"cr_attention\") \n\n\nclass TFContentCoder(TFTransformer):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n \n self.layer = [TFContentBlock(config, name=\"layer_._{}\".format(i)) for i in range(config.n_layers)]\n \n \nclass TFResponseCoder(TFTransformer):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs) \n \n self.layer = [TFResponseBlock(config, name=\"layer_._{}\".format(i)) for i in range(config.n_layers)]\n\n def call(self, r_embeds, c_hidden, r_mask, r_c_mask, head_mask, output_attentions, output_hidden_states, return_dict, training=False):\n # docstyle-ignore\n \"\"\"\n \"\"\"\n all_hidden_states = () if output_hidden_states else None\n all_attentions = () if output_attentions else None\n\n hidden_state = r_embeds\n for i, layer_module in enumerate(self.layer):\n if output_hidden_states:\n all_hidden_states = all_hidden_states + (hidden_state,)\n\n layer_outputs = layer_module(hidden_state, c_hidden, r_mask, r_c_mask, head_mask[i], output_attentions, training=training)\n hidden_state = layer_outputs[-1]\n\n if output_attentions:\n assert len(layer_outputs) == 3\n r_attn = layer_outputs[0]\n r_c_attn = layer_outputs[1]\n all_attentions = all_attentions + ((r_attn, r_c_attn),)\n else:\n assert len(layer_outputs) == 1, f\"Incorrect number of outputs {len(layer_outputs)} instead of 1\"\n\n # Add last layer\n if output_hidden_states:\n all_hidden_states = all_hidden_states + (hidden_state,)\n\n if not return_dict:\n return tuple(v for v in [hidden_state, all_hidden_states, all_attentions] if v is not None)\n return TFBaseModelOutput(\n last_hidden_state=hidden_state, hidden_states=all_hidden_states, attentions=all_attentions\n ) \n \n def generate(\n self,\n current_pos,\n current_d_embedding,\n prev_c_hidden,\n all_prev_d_hidden,\n r_mask,\n r_c_mask,\n head_mask,\n output_attentions\n ):\n\n all_hidden_states = ()\n all_attentions = () if output_attentions else None\n\n current_d_hidden = current_d_embedding\n for i, (layer_module, prev_d_hidden) in enumerate(zip(self.layer, all_prev_d_hidden[:-1])):\n\n all_hidden_states = all_hidden_states + (current_d_hidden,)\n\n layer_outputs = layer_module.generate(\n current_pos, current_d_hidden, prev_c_hidden, prev_d_hidden, r_mask, r_c_mask, head_mask[i], output_attentions\n )\n current_d_hidden = layer_outputs[-1]\n\n if output_attentions:\n assert len(layer_outputs) == 3\n r_attn = layer_outputs[0]\n r_c_attn = layer_outputs[1]\n all_attentions = all_attentions + ((r_attn, r_c_attn),)\n else:\n assert len(layer_outputs) == 1, f\"Incorrect number of outputs {len(layer_outputs)} instead of 1\"\n\n # Add last layer\n all_hidden_states = all_hidden_states + (current_d_hidden,)\n\n return tuple(v for v in [current_d_hidden, all_hidden_states, all_attentions] if v is not None) \n\n\nclass TFCRCoder(TFTransformer):\n \n def __init__(self, config, **kwargs):\n \n super().__init__(config, **kwargs)\n \n self.layer = [TFCRBlock(config, name=\"layer_._{}\".format(i)) for i in range(config.n_layers)]\n \n def call(self, c_embeds, r_embeds, c_mask, r_mask, c_r_mask, r_c_mask, head_mask, output_attentions, output_hidden_states, return_dict, training=False):\n\n all_hidden_states = () if output_hidden_states else None\n all_attentions = () if output_attentions else None \n \n # shape = [batch_size, 2 * seq_len, hidden_dim]\n embeds = tf.concat([c_embeds, r_embeds], axis=1)\n\n _c_mask = tf.concat([c_mask, c_r_mask], axis=2)\n _r_mask = tf.concat([r_c_mask, r_mask], axis=2)\n \n attn_mask = tf.concat([_c_mask, _r_mask], axis=1)\n \n hidden_state = embeds\n for i, layer_module in enumerate(self.layer):\n if output_hidden_states:\n all_hidden_states = all_hidden_states + (hidden_state,)\n \n layer_outputs = layer_module(hidden_state, attn_mask, head_mask[i], output_attentions, training=training)\n hidden_state = layer_outputs[-1]\n \n if output_attentions:\n assert len(layer_outputs) == 2\n attn = layer_outputs[0]\n all_attentions = all_attentions + (attn,)\n else:\n assert len(layer_outputs) == 1, f\"Incorrect number of outputs {len(layer_outputs)} instead of 1\" \n \n # Add last layer\n if output_hidden_states:\n all_hidden_states = all_hidden_states + (hidden_state,) \n \n if not return_dict:\n return tuple(v for v in [hidden_state, all_hidden_states, all_attentions] if v is not None)\n return TFBaseModelOutput(\n last_hidden_state=hidden_state, hidden_states=all_hidden_states, attentions=all_attentions\n ) \n\n \nclass TFEdFormerEncoder(tf.keras.layers.Layer):\n \n def __init__(self, config, **kwargs):\n\n super().__init__(**kwargs)\n\n self.cr_encoder = TFCRCoder(config, name='cr_encoder')\n \n def call(\n self,\n c_embeds,\n r_embeds,\n d_embeds,\n c_mask,\n r_mask, \n r_c_mask,\n c_r_mask,\n head_mask,\n output_attentions,\n output_hidden_states,\n return_dict,\n training=False,\n ):\n\n # currently, only empty tuple is returned\n c_all_hidden_states = () if output_hidden_states else None\n c_all_attentions = () if output_attentions else None \n r_all_hidden_states = () if output_hidden_states else None\n r_all_attentions = () if output_attentions else None \n\n cr_outputs = self.cr_encoder(\n c_embeds, r_embeds, c_mask, r_mask, c_r_mask, r_c_mask, head_mask,\n output_attentions, output_hidden_states, return_dict, training=training\n )\n \n if not return_dict:\n cr_hidden = cr_outputs[0]\n else:\n cr_hidden = cr_outputs.last_hidden_state\n \n c_seq_len = tf.math.reduce_sum(tf.ones_like(c_embeds[0, :, 0], dtype=tf.int32))\n r_seq_len = tf.math.reduce_sum(tf.ones_like(r_embeds[0, :, 0], dtype=tf.int32))\n \n c_hidden = cr_hidden[:, 0:c_seq_len, :]\n r_hidden = cr_hidden[:, c_seq_len:(c_seq_len + r_seq_len), :]\n\n # Adding all hidden states and attentions\n if not return_dict:\n c_outputs = tuple(v for v in [c_hidden, c_all_hidden_states, c_all_attentions] if v is not None)\n r_outputs = tuple(v for v in [r_hidden, r_all_hidden_states, r_all_attentions] if v is not None)\n else:\n c_outputs =TFBaseModelOutput(\n last_hidden_state=c_hidden, hidden_states=c_all_hidden_states, attentions=c_all_attentions\n ) \n r_outputs =TFBaseModelOutput(\n last_hidden_state=r_hidden, hidden_states=r_all_hidden_states, attentions=r_all_attentions\n )\n\n hidden_states = c_hidden\n\n return (hidden_states, c_outputs, r_outputs)\n\n\ndef process_mask(mask, input_shpae):\n \n if mask is None:\n mask = tf.ones(input_shpae)\n mask = tf.cast(mask, dtype=tf.float32) \n \n return mask\n\n\nclass TFEdFormerEncoderDecoder(tf.keras.layers.Layer):\n \n def __init__(self, config, **kwargs):\n\n super().__init__(**kwargs)\n\n self.encoder = TFContentCoder(config, name=\"encoder\") # Encoder\n self.decoder = TFResponseCoder(config, name=\"decoder\") # Decoder\n \n def call(\n self,\n c_embeds,\n r_embeds,\n d_embeds,\n c_mask,\n r_mask, \n r_c_mask,\n c_r_mask,\n head_mask,\n output_attentions,\n output_hidden_states,\n return_dict,\n training=False,\n ):\n\n c_outputs = self.encoder(\n c_embeds, c_mask, head_mask,\n output_attentions, output_hidden_states, return_dict, training\n )\n \n if not return_dict:\n c_hidden = c_outputs[0]\n else:\n c_hidden = c_outputs.last_hidden_state\n \n r_outputs = self.decoder(\n d_embeds, c_hidden, r_mask, r_c_mask, head_mask,\n output_attentions, output_hidden_states, return_dict, training\n ) \n\n if not return_dict:\n r_hidden = r_outputs[0]\n else:\n r_hidden = r_outputs.last_hidden_state\n\n hidden_states = r_hidden\n\n return (hidden_states, c_outputs, r_outputs)\n\n def generate(\n self,\n current_pos,\n current_d_embedding,\n prev_c_hidden,\n all_prev_d_hidden,\n r_mask,\n r_c_mask,\n head_mask,\n output_attentions\n ):\n\n return self.decoder.generate(current_pos, current_d_embedding, prev_c_hidden, all_prev_d_hidden, r_mask, r_c_mask, head_mask, output_attentions)\n\n\nclass TFEdFormerMainLayer(tf.keras.layers.Layer):\n \n config_class = EdFormerConfig\n\n def __init__(self, config, **kwargs):\n\n super().__init__(**kwargs)\n \n self.num_hidden_layers = config.num_hidden_layers\n self.output_attentions = config.output_attentions\n self.output_hidden_states = config.output_hidden_states\n self.return_dict = config.use_return_dict\n \n self.use_tags = config.use_tags\n self.use_part = config.use_part\n self.use_prior_explanation = config.use_prior_explanation\n self.use_prior_question_elapsed_time_input = config.use_prior_question_elapsed_time_input \n\n if config.share_position_embeddings:\n # All `TFEmbeddings` share a single `position_embeddings`.\n position_embeddings = tf.keras.layers.Embedding(\n config.max_position_embeddings,\n config.dim,\n embeddings_initializer=get_initializer(config.seed),\n name=\"position_embeddings\",\n )\n else:\n # Each `TFEmbeddings` will have its own `position_embeddings`.\n position_embeddings = None\n \n self.content_embeddings = TFEmbeddings(\n config,\n input_name = 'content',\n position_embeddings_layer=position_embeddings,\n name=\"content_embeddings\"\n )\n \n self.response_embeddings = TFEmbeddings(\n config,\n input_name = 'response',\n position_embeddings_layer=position_embeddings,\n name=\"response_embeddings\"\n )\n \n if config.model_type == 'cr':\n self.coder = TFEdFormerEncoder(config, name='coder')\n elif config.model_type == 'ed':\n self.coder = TFEdFormerEncoderDecoder(config, name='coder')\n \n def call(\n self,\n inputs,\n c_mask=None,\n r_mask=None,\n r_c_mask=None,\n c_r_mask=None,\n head_mask=None, \n output_attentions=None,\n output_hidden_states=None,\n return_dict=None,\n training=False, \n ):\n\n c_input_ids = inputs.get('c_input_ids')\n r_input_ids = inputs.get('r_input_ids')\n d_input_ids = inputs.get('d_input_ids')\n\n batch_size = tf.math.reduce_sum(tf.ones_like(c_input_ids[:, 0], dtype=tf.int32))\n c_seq_len = tf.math.reduce_sum(tf.ones_like(c_input_ids[0, :], dtype=tf.int32))\n r_seq_len = tf.math.reduce_sum(tf.ones_like(r_input_ids[0, :], dtype=tf.int32))\n \n # The simplest way to compute positions.\n pos_ids = tf.range(c_seq_len, dtype=tf.int32) + 1\n shifted_pos_ids = tf.concat([[0], pos_ids[:-1]], axis=0)\n\n # positional information provided\n pos_ids = inputs.get('pos_ids', pos_ids)\n shifted_pos_ids = inputs.get('shifted_pos_ids', shifted_pos_ids)\n\n if self.use_tags:\n tag_ids = inputs.get('tag_ids')\n else:\n tag_ids = None\n \n if self.use_part:\n part_ids = inputs.get('part_ids')\n else:\n part_ids = None\n\n if self.use_prior_explanation:\n prior_explanation_ids = inputs.get('prior_explanation_ids')\n else:\n prior_explanation_ids = None\n\n if self.use_prior_question_elapsed_time_input:\n prior_question_elapsed_time_input = inputs.get('prior_question_elapsed_time_input')\n else:\n prior_question_elapsed_time_input = None\n\n c_mask = inputs.get('c_mask', c_mask)\n r_mask = inputs.get('r_mask', r_mask)\n r_c_mask = inputs.get('r_c_mask', r_c_mask)\n c_r_mask = inputs.get('c_r_mask', c_r_mask) \n head_mask = inputs.get('head_mask', head_mask)\n \n output_attentions = inputs.get('output_attentions', output_attentions)\n output_hidden_states = inputs.get('output_hidden_states', output_hidden_states)\n return_dict = inputs.get('return_dict', return_dict)\n \n output_attentions = output_attentions if output_attentions is not None else self.output_attentions\n output_hidden_states = output_hidden_states if output_hidden_states is not None else self.output_hidden_states\n return_dict = return_dict if return_dict is not None else self.return_dict\n \n c_mask = process_mask(c_mask, [batch_size, c_seq_len, c_seq_len])\n r_mask = process_mask(r_mask, [batch_size, r_seq_len, r_seq_len])\n r_c_mask = process_mask(r_c_mask, [batch_size, r_seq_len, c_seq_len])\n c_r_mask = process_mask(c_r_mask, [batch_size, c_seq_len, r_seq_len]) \n \n # Prepare head mask if needed\n # 1.0 in head_mask indicate we keep the head\n # attention_probs has shape bsz x n_heads x N x N\n # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]\n # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]\n if head_mask is not None:\n raise NotImplementedError\n else:\n head_mask = [None] * self.num_hidden_layers\n \n c_embedding_output = self.content_embeddings(input_ids=c_input_ids, position_ids=pos_ids, tag_ids=tag_ids, part_ids=part_ids, prior_explanation_ids=None, training=training) # (bs, seq_length, dim)\n r_embedding_output = self.response_embeddings(input_ids=r_input_ids, position_ids=pos_ids, tag_ids=None, part_ids=None, prior_explanation_ids=None, training=training) # (bs, seq_length, dim)\n d_embedding_output = self.response_embeddings(input_ids=d_input_ids, position_ids=shifted_pos_ids, tag_ids=None, part_ids=None, prior_explanation_ids=prior_explanation_ids, prior_question_elapsed_time_input=prior_question_elapsed_time_input, training=training) # (bs, seq_length, dim)\n\n outputs = self.coder(\n c_embedding_output,\n r_embedding_output,\n d_embedding_output,\n c_mask,\n r_mask,\n r_c_mask,\n c_r_mask,\n head_mask,\n output_attentions,\n output_hidden_states,\n return_dict,\n training=training \n )\n \n return outputs\n\n def generate(\n self,\n inputs,\n current_pos,\n seq_len,\n updated_d_input_ids, # (bs, seq_len, dim)\n prev_c_hidden,\n all_prev_d_hidden,\n r_mask,\n r_c_mask,\n head_mask,\n output_attentions\n ):\n # The simplest way to compute positions.\n pos_ids = tf.range(seq_len, dtype=tf.int32) + 1\n shifted_pos_ids = tf.concat([[0], pos_ids[:-1]], axis=0)\n\n # positional information provided\n pos_ids = inputs.get('pos_ids', pos_ids)\n shifted_pos_ids = inputs.get('shifted_pos_ids', shifted_pos_ids)\n\n if self.use_prior_explanation:\n prior_explanation_ids = inputs.get('prior_explanation_ids')\n else:\n prior_explanation_ids = None\n\n if self.use_prior_question_elapsed_time_input:\n prior_question_elapsed_time_input = inputs.get('prior_question_elapsed_time_input')\n else:\n prior_question_elapsed_time_input = None\n\n current_d_embedding = self.response_embeddings(\n input_ids=updated_d_input_ids,\n position_ids=shifted_pos_ids,\n tag_ids=None,\n part_ids=None,\n prior_explanation_ids=prior_explanation_ids,\n prior_question_elapsed_time_input=prior_question_elapsed_time_input,\n )[:, current_pos:current_pos+1, :]\n\n outputs = self.coder.generate(\n current_pos,\n current_d_embedding,\n prev_c_hidden,\n all_prev_d_hidden,\n r_mask,\n r_c_mask,\n head_mask,\n output_attentions=output_attentions\n )\n \n return outputs \n\n\nclass TFEdFormerPreTrainedModel(TFPreTrainedModel):\n \"\"\"\n An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained\n models.\n \"\"\"\n\n config_class = EdFormerConfig\n base_model_prefix = \"edformer\"\n \n \nclass TFEdFormerModel(TFEdFormerPreTrainedModel):\n \n def __init__(self, config, *inputs, **kwargs):\n \n super().__init__(config, *inputs, **kwargs)\n \n self.edformer = TFEdFormerMainLayer(config, name=\"edformer\") # Embeddings\n\n def call(self, inputs, **kwargs):\n \n outputs = self.edformer(\n inputs, \n **kwargs\n )\n \n return outputs","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### The model\n\nThe model for answer correctness prediction.\n\nIn particular, a `generate()` method is implemented for auto-regressive prediction generation. It uses some `generate()` methods in the build block layers defined in the previous code cell.\n\n![\"Capture777\"](\"https://i.ibb.co/bW09VT1/Capture777.png\")
\n
\nauto-regressive prediction generation
\n
\n\n* credit for the picture: [Mariya Yao: Novel Methods For Text Generation Using Adversarial Learning & Autoencoders](#https://www.topbots.com/ai-research-gan-vae-text-generation/)","attachments":{}},{"metadata":{"trusted":true},"cell_type":"code","source":"class TFEdFormerAnswerPredictionModel(TFPreTrainedModel):\n \n def __init__(self, config, *inputs, **kwargs):\n \n super().__init__(config, *inputs, **kwargs)\n \n self.edformer = TFEdFormerMainLayer(config, name=\"edformer\") # Embeddings\n\n self.pre_classifier = tf.keras.layers.Dense(\n config.dim,\n kernel_initializer=get_initializer(config.seed),\n name=\"pre_classifier\",\n ) \n\n assert config.activation in [\"relu\", \"gelu\"], \"activation ({}) must be in ['relu', 'gelu']\".format(\n config.activation\n )\n self.activation = get_tf_activation(config.activation)\n\n self.classifier = tf.keras.layers.Dense(\n 1, kernel_initializer=get_initializer(config.seed), name=\"classifier\"\n )\n self.dropout = tf.keras.layers.Dropout(config.seq2seq_dropout)\n \n def call(\n self,\n inputs=None,\n c_mask=None,\n r_mask=None,\n r_c_mask=None,\n c_r_mask=None,\n head_mask=None,\n output_attentions=None,\n output_hidden_states=None,\n return_dict=None,\n labels=None,\n training=False,\n ):\n\n return_dict = return_dict if return_dict is not None else self.edformer.return_dict\n\n edformer_output = self.edformer(\n inputs,\n c_mask=c_mask,\n r_mask=r_mask,\n r_c_mask=r_c_mask,\n c_r_mask=c_r_mask,\n head_mask=head_mask, \n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n return_dict=return_dict,\n training=training,\n )\n\n hidden_states, c_outputs, r_outputs = edformer_output\n\n o = self.pre_classifier(hidden_states) # (bs, dim)\n o = self.activation(o)\n o = self.dropout(o, training=training) \n \n logits = self.classifier(o) # (bs, seq_len, dim)\n logits = logits[:, :, 0]\n\n return (logits, c_outputs, r_outputs)\n\n def generate(\n self,\n inputs,\n start_pos,\n window_size,\n dim,\n c_mask=None,\n r_mask=None,\n r_c_mask=None,\n c_r_mask=None,\n head_mask=None,\n output_attentions=None\n ):\n \"\"\"\n Generate the predictions in a auto-regressive way.\n \"\"\"\n\n question_mask = tf.cast((inputs['content_type_id'] == 0), dtype=tf.int32)\n pred_time_question_mask = inputs['pred_time_mask'] * question_mask\n pred_time_question_mask_shifted = tf.concat([tf.zeros_like(pred_time_question_mask[:, -1:], dtype=tf.int32), pred_time_question_mask[:, 0:-1]], axis=1)\n\n preds = tf.cast(inputs['target'], dtype=tf.float32)\n\n d_hidden, c_outputs, r_outputs = self.edformer(\n inputs,\n c_mask=c_mask,\n r_mask=r_mask,\n r_c_mask=r_c_mask,\n c_r_mask=c_r_mask,\n head_mask=head_mask, \n output_attentions=output_attentions,\n output_hidden_states=True,\n training=False,\n )\n\n o = self.pre_classifier(d_hidden) # (bs, seq_len, dim)\n o = self.activation(o)\n o = self.dropout(o, training=False)\n\n logits = self.classifier(o)[:, :, 0] # (bs, seq_len)\n\n prev_c_hidden = c_outputs[0]\n all_prev_d_hidden = r_outputs[1]\n\n c_input_ids = inputs.get('c_input_ids')\n d_input_ids = inputs.get('d_input_ids')\n updated_d_input_ids = d_input_ids\n\n batch_size = tf.math.reduce_sum(tf.ones_like(c_input_ids[:, 0], dtype=tf.int32))\n c_seq_len = tf.math.reduce_sum(tf.ones_like(c_input_ids[0, :], dtype=tf.int32))\n d_seq_len = tf.math.reduce_sum(tf.ones_like(updated_d_input_ids[0, :], dtype=tf.int32))\n \n r_mask = process_mask(r_mask, [batch_size, d_seq_len, d_seq_len])\n r_c_mask = process_mask(r_c_mask, [batch_size, d_seq_len, c_seq_len])\n \n if head_mask is not None:\n raise NotImplementedError\n else:\n head_mask = [None] * self.edformer.num_hidden_layers \n\n output_attentions = inputs.get('output_attentions', output_attentions)\n output_attentions = output_attentions if output_attentions is not None else self.edformer.output_attentions\n\n for current_pos in tf.range(start_pos, d_seq_len, dtype=tf.int32):\n\n _current_d_hidden, all_current_d_hidden = self.edformer.generate(\n inputs,\n current_pos,\n c_seq_len,\n updated_d_input_ids,\n prev_c_hidden,\n all_prev_d_hidden,\n r_mask,\n r_c_mask,\n head_mask,\n output_attentions=output_attentions\n )[0:2]\n\n all_prev_d_hidden = tuple([\n tf.reshape(tf.concat(\n [x[:, :current_pos, :], y, x[:, current_pos+1:d_seq_len, :]],\n axis=1\n ), shape=[batch_size, window_size, dim]) for x, y in zip(all_prev_d_hidden, all_current_d_hidden)\n ])\n\n # ------------------------------------------------------------------------------------------\n # The logit for the place at `current_pos`.\n\n o = self.pre_classifier(_current_d_hidden) # (bs, 1, dim)\n o = self.activation(o)\n o = self.dropout(o, training=False)\n\n _logits = self.classifier(o)[:, :, 0] # (bs, 1)\n _preds = tf.math.sigmoid(_logits) # (bs, 1)\n\n logits = tf.concat([logits[:, :current_pos], _logits, logits[:, current_pos+1:d_seq_len]], axis=1)\n preds = tf.concat([preds[:, :current_pos], _preds, preds[:, current_pos+1:d_seq_len]], axis=1)\n\n logits = tf.reshape(logits, shape=[batch_size, window_size])\n preds = tf.reshape(preds, shape=[batch_size, window_size])\n\n # ------------------------------------------------------------------------------------------\n # update the d_input_ids at `current_pos + 1`\n\n predicted_answered_correctly = tf.cast(_preds >= 0.5, dtype=tf.int32) # (bs, 1)\n\n if current_pos < d_seq_len - 1:\n \n mask = tf.cast(updated_d_input_ids != PAD_ID, dtype=tf.int32) * tf.cast(updated_d_input_ids != START_ID, dtype=tf.int32) * tf.cast(updated_d_input_ids != RESPONSE_LECTURE_ID, dtype=tf.int32)\n mask = mask * pred_time_question_mask_shifted\n mask = mask[:, current_pos + 1:current_pos + 2]\n\n predicted_next_response_id = (predicted_answered_correctly + RESPONSE_FALSE_ID) * mask + (1 - mask) * updated_d_input_ids[:, current_pos + 1:current_pos + 2]\n\n updated_d_input_ids = tf.concat([updated_d_input_ids[:, :current_pos+1], predicted_next_response_id[:, :1], updated_d_input_ids[:, current_pos+2:]], axis=1)\n\n updated_d_input_ids = tf.reshape(updated_d_input_ids, shape=[batch_size, window_size])\n # ------------------------------------------------------------------------------------------\n\n return preds, logits","execution_count":null,"outputs":[]},{"metadata":{"id":"kgXCR1dp11te"},"cell_type":"markdown","source":"### check"},{"metadata":{"trusted":true,"id":"pbwmHbYc11ti","outputId":"bb64c680-82b5-485e-d5c2-fa8a714e806b"},"cell_type":"code","source":"content_input_ids = tf.constant(1, shape=[3, 5])\nresponse_input_ids = tf.constant(1, shape=[3, 5])\ndecoder_input_ids = tf.constant(1, shape=[3, 5])\ntag_ids = tf.constant(0, shape=[3, 5, N_TAGS_PER_CONTENT])\npart_ids = tf.constant(1, shape=[3, 5])\nprior_explanation_ids = tf.constant(1, shape=[3, 5])\nprior_question_elapsed_time_inputs = tf.constant(1.0, shape=[3, 5])\n\ninputs = {\n 'c_input_ids': content_input_ids,\n 'r_input_ids': response_input_ids,\n 'd_input_ids': decoder_input_ids,\n 'tag_ids': tag_ids,\n 'part_ids': part_ids,\n 'prior_explanation_ids': prior_explanation_ids,\n 'prior_question_elapsed_time_inputs': prior_question_elapsed_time_inputs\n}\n\nfor model_type in ['cr', 'ed']:\n config = EdFormerConfig(model_type=model_type, model_desc='dummy', share_position_embeddings=True, use_tags=True, user_part=True, use_prior_explanation=True, use_prior_question_elapsed_time_input=True)\n predictor = TFEdFormerAnswerPredictionModel(config)\n logits, c_outputs, r_outputs = predictor(inputs=inputs, output_attentions=True, output_hidden_states=True, return_dict=False)\n print(logits)","execution_count":null,"outputs":[]},{"metadata":{"id":"srxuqrpW11ur"},"cell_type":"markdown","source":"## Training Manager"},{"metadata":{},"cell_type":"markdown","source":"### Learing rate with warmup\n\nTwo learning rate schedules are provided: `Linear` and `Noam`, both with linear warmup. The `Linear` schedule has linear decay."},{"metadata":{"trusted":true},"cell_type":"code","source":"class LinearLR(tf.keras.optimizers.schedules.LearningRateSchedule):\n \n def __init__(self, total_steps, lr=1e-4, end_lr=1e-6, warmup_steps=WARMUP_STEPS):\n \n self.total_steps = tf.cast(total_steps, dtype=DTYPE)\n self.lr = lr\n self.end_lr = end_lr\n self.warmup_steps = tf.cast(warmup_steps, dtype=DTYPE)\n \n def __call__(self, step):\n \n is_warmup = tf.cast(step < self.warmup_steps, dtype=DTYPE)\n \n warmup_lr = is_warmup * self.lr * (step + 1) / self.warmup_steps \n decay_lr = (1 - is_warmup) * (self.lr - (self.lr - self.end_lr) / tf.math.maximum(self.total_steps - self.warmup_steps, 1) * (step - self.warmup_steps + 1))\n decay_lr = (1 - is_warmup) * tf.math.maximum(self.end_lr, decay_lr)\n\n lr = warmup_lr + decay_lr\n\n return lr\n\nclass NoamLR(tf.keras.optimizers.schedules.LearningRateSchedule):\n\n def __init__(self, hidden_size, lr, end_lr, warmup_steps):\n\n self.lr = tf.cast(lr, DTYPE)\n self.end_lr = tf.cast(end_lr, DTYPE)\n self.warmup_steps = tf.cast(warmup_steps, DTYPE)\n self.hidden_size = tf.cast(hidden_size, DTYPE)\n\n def __call__(self, step):\n\n scaling = self.lr / (self.hidden_size**-0.5 * self.warmup_steps**-0.5)\n\n lr = scaling * self.hidden_size**-0.5 * tf.math.minimum((step + 1) * self.warmup_steps**-1.5, (step + 1)**-0.5)\n\n is_warmup = tf.cast(step < self.warmup_steps, dtype=DTYPE)\n\n lr = lr * is_warmup + (1.0 - is_warmup) * tf.math.maximum(self.end_lr, lr) \n\n return lr","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model input signatures"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_input_signatures(batch_size, seq_len, valid=False):\n\n input_signatures = {\n 'user_id': tf.TensorSpec(shape=[batch_size], dtype=tf.int64, name='user_id'),\n 'seq_len': tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='seq_len'),\n 'prev_seq_len': tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='prev_seq_len'),\n 'start': tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='start'),\n 'end': tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='end'),\n 'row_id': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int64, name='row_id'),\n 'timestamp': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int64, name='timestamp'),\n 'content_id': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='content_id'),\n 'content_type_id': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='content_type_id'),\n 'task_container_id': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='task_container_id'),\n 'user_answer': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='user_answer'),\n 'answered_correctly': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='answered_correctly'),\n 'shifted_answered_correctly': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='shifted_answered_correctly'),\n 'prior_question_elapsed_time': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='prior_question_elapsed_time'),\n 'prior_question_elapsed_time_input': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='prior_question_elapsed_time_input'),\n 'prior_question_had_explanation': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='prior_question_had_explanation'),\n 'pred_time_mask': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='pred_time_mask'),\n 'abs_pos': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='abs_pos'),\n 'shifted_abs_pos': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='shifted_abs_pos'),\n 'pos_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='pos_ids'),\n 'shifted_pos_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='shifted_pos_ids'), \n 'c_input_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='c_input_ids'),\n 'r_input_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='r_input_ids'),\n 'd_input_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='d_input_ids'),\n 'tag_ids': tf.TensorSpec(shape=[batch_size, seq_len, N_TAGS_PER_CONTENT], dtype=tf.int32, name='tag_ids'),\n 'part_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_ids'), \n 'prior_explanation_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='prior_explanation_ids'), \n 'target': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='target'),\n 'nb_pred_places': tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='nb_pred_places')\n }\n\n if valid:\n\n input_signatures['n_valid_blocks'] = tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='n_valid_blocks')\n input_signatures['valid_block_idx'] = tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='valid_block_idx')\n input_signatures['valid_start'] = tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='valid_start')\n input_signatures['valid_end'] = tf.TensorSpec(shape=[batch_size], dtype=tf.int32, name='valid_end')\n\n return [input_signatures]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training Manager"},{"metadata":{"trusted":true,"id":"s0MnJ1m111uu"},"cell_type":"code","source":"class TrainConfig:\n\n def __init__(\n self,\n ckpt_path=CKPT_TRAIN_PATH,\n window_size=WINDOW_SIZE,\n loss_weight_window_size=LOSS_WEIGHT_WINDOW_SIZE,\n n_epochs=N_EPOCHS,\n batch_size=BATCH_SIZE,\n pred_batch_size=PRED_BATCH_SIZE,\n shuffle_buf_size=SHUFFLE_BUFFER_SIZE,\n seed=SEED,\n deterministic=DETERMINISTIC,\n num_parallel_reads=N_PARALLEL_READS,\n num_parallel_calls=N_PARALLEL_CALLS,\n steps_per_call=STEPS_PER_CALL\n ):\n\n self.ckpt_path = ckpt_path\n self.window_size = window_size\n self.loss_weight_window_size = loss_weight_window_size\n self.n_epochs = n_epochs\n self.batch_size = batch_size\n self.pred_batch_size = pred_batch_size\n self.shuffle_buf_size = shuffle_buf_size\n self.seed=seed\n self.deterministic = deterministic\n self.num_parallel_reads = num_parallel_reads\n self.num_parallel_calls = num_parallel_calls\n self.steps_per_call = steps_per_call\n\n n_contents_dict = load_data(n_contents_dict_path)\n n_training_examples = 0\n for k, v in n_contents_dict.items():\n n_training_examples += math.ceil(v * 1.0 / self.window_size)\n del n_contents_dict\n\n valid_info = load_data(valid_info_path)\n n_valid_examples = 0\n for k, v in valid_info.items():\n n_valid_examples += v['bundle_info']['n_blocks']\n del valid_info\n\n # This is slightly higher than the actual number of examples used for training.\n self.n_training_examples = n_training_examples\n n_training_steps_per_epoch = self.n_training_examples // self.batch_size\n self.n_training_calls_per_epoch = n_training_steps_per_epoch // self.steps_per_call\n self.n_training_steps_per_epoch = self.n_training_calls_per_epoch * self.steps_per_call\n self.n_training_steps = self.n_epochs * self.n_training_steps_per_epoch\n\n self.n_valid_examples = n_valid_examples\n if IS_KAGGLE and tpu is not None:\n self.n_valid_examples -= self.n_valid_examples % self.pred_batch_size\n self.n_valid_steps = self.n_valid_examples // self.pred_batch_size + int(self.n_valid_examples % self.pred_batch_size > 0)\n self.n_steps_in_last_valid_call = self.n_valid_steps % self.steps_per_call\n self.n_valid_calls = self.n_valid_steps // self.steps_per_call + int(self.n_steps_in_last_valid_call > 0)\n\n self.lr = None\n self.end_lr = None\n self.warmup_steps = None\n\n def toJSON(self):\n\n return json.dumps(self, default=lambda o: o.__dict__, sort_keys=False, indent=4)\n\nclass Train_Manager:\n \n def __init__(self, config, train_config):\n \n self.config = config\n self.train_config = train_config\n\n self.train_ds = strategy.experimental_distribute_dataset(self.get_train_ds())\n self.valid_ds = strategy.experimental_distribute_dataset(self.get_valid_ds())\n \n def toJSON(self):\n\n config = json.loads(self.config.toJSON())\n train_config = json.loads(self.train_config.toJSON())\n\n for k, v in train_config.items():\n config[k] = v\n\n return json.dumps(config, sort_keys=False, ensure_ascii=False, indent=4)\n\n def get_train_ds(self):\n\n train_raw_ds = tf.data.TFRecordDataset(train_tfrec_paths, num_parallel_reads=self.train_config.num_parallel_reads)\n train_raw_ds = train_raw_ds.map(parse_train_example, num_parallel_calls=self.train_config.num_parallel_calls, deterministic=self.train_config.deterministic)\n\n\n # Get the splitted training dataset.\n train_valid_split_indices = load_data(train_valid_split_indices_path) \n train_valid_split_table = convert_split_index_dict(train_valid_split_indices)\n reduced_raw_ds = split_train_ds(train_raw_ds, train_valid_split_table, num_parallel_calls=self.train_config.num_parallel_calls, deterministic=self.train_config.deterministic)\n\n # repeat different times for each user, according to their history lengths.\n # This is used to avoid overfitting on users with much shorter history sequences.\n reduced_raw_ds_augmented = reduced_raw_ds.flat_map(\n lambda x: tf.data.Dataset.from_tensors(x).repeat(\n tf.cast(\n tf.math.ceil(tf.cast(x['seq_len'], tf.float32) / self.train_config.window_size), dtype=tf.int64\n )\n )\n )\n\n # batch examples with attributes having different lengths across examples - tf.RaggedTensor\n batched_raw_ds = reduced_raw_ds_augmented.repeat().shuffle(self.train_config.shuffle_buf_size, seed=self.train_config.seed).apply(tf.data.experimental.dense_to_ragged_batch(self.train_config.batch_size))\n \n # batch - tf.Tensor: Extract random subsequences of a fixed length\n batched_ds = random_subseqs_from_batched_raw_ds(batched_raw_ds, window_size=self.train_config.window_size, only_last=False, seed=self.train_config.seed, num_parallel_calls=self.train_config.num_parallel_calls, deterministic=self.train_config.deterministic)\n \n # Add input ids and targets for training \n train_ds = prepare_training_dataset(batched_ds, generative=self.config.generative, use_abs_pos=self.config.use_abs_pos, num_parallel_calls=self.train_config.num_parallel_calls, deterministic=self.train_config.deterministic)\n \n train_ds = train_ds.prefetch(8)\n\n return train_ds\n\n def get_valid_ds(self):\n\n p = valid_tfrec_paths\n num_parallel_reads = self.train_config.num_parallel_reads\n num_parallel_calls = self.train_config.num_parallel_calls\n deterministic = self.train_config.deterministic\n\n valid_raw_ds = tf.data.TFRecordDataset(p, num_parallel_reads=num_parallel_reads)\n valid_raw_ds = valid_raw_ds.map(parse_train_example_with_valid_info, num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n valid_ds = prepare_validation_dataset(valid_raw_ds, batch_size=self.train_config.pred_batch_size, window_size=self.train_config.window_size, generative=self.config.generative, use_abs_pos=self.config.use_abs_pos, num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n \n valid_ds = valid_ds.prefetch(8)\n\n return valid_ds\n\n def get_train_objs(self, lr=LR, end_lr=END_LR, warmup_steps=WARMUP_STEPS):\n \n self.train_config.lr = lr\n self.train_config.end_lr = end_lr,\n self.train_config.warmup_steps = warmup_steps\n\n with strategy.scope():\n\n predictor = TFEdFormerAnswerPredictionModel(self.config)\n\n _lr = NoamLR(hidden_size=self.config.dim, lr=lr, end_lr=end_lr, warmup_steps=warmup_steps)\n\n # optimizer = tf.keras.optimizers.Adam(learning_rate=_lr, beta_1=0.9, beta_2=0.999, epsilon=1e-08)\n optimizer = tf.keras.optimizers.Adam(\n _lr,\n beta_1=tf.Variable(0.9),\n beta_2=tf.Variable(0.999),\n epsilon=tf.Variable(1e-8)\n )\n optimizer.iterations # this access will invoke optimizer._iterations method and create optimizer.iter attribute\n optimizer.decay = tf.Variable(0.0) # Adam.__init__ assumes ``decay`` is a float object, so this needs to be converted to tf.Variable **after** __init__ method.\n\n loss_obj = tf.keras.losses.BinaryCrossentropy(\n from_logits=True, label_smoothing=0, reduction=tf.keras.losses.Reduction.NONE,\n name='binary_crossentropy'\n )\n\n loss_metric = tf.keras.metrics.Sum()\n acc_metric = tf.keras.metrics.BinaryAccuracy()\n auc_metric = tf.keras.metrics.AUC(num_thresholds=2000)\n\n # --------------------------------------------------\n\n metrics = (loss_metric, acc_metric, auc_metric)\n \n return (predictor, optimizer, loss_obj, metrics)\n\n def train_valid(self, predictor, optimizer, ckpt_manager, loss_obj, metrics, last_epoch=0, valid_epochs=None, only_valid=False):\n\n batch_size = self.train_config.batch_size\n seq_len = self.train_config.window_size\n\n encoder_decoder = tf.constant(self.config.model_type=='ed', dtype=tf.int32)\n generative = tf.constant(self.config.generative, dtype=tf.int32)\n allow_bundle_atten = tf.constant(self.config.allow_bundle_atten, dtype=tf.int32)\n\n loss_metric, acc_metric, auc_metric = metrics\n\n @tf.function(\n input_signature=get_input_signatures(int(batch_size / strategy.num_replicas_in_sync), seq_len, valid=False)\n )\n def train_step(train_batch):\n\n training = tf.constant(1, dtype=tf.int32)\n \n c_mask, r_mask, r_c_mask, c_r_mask = get_attention_masks(train_batch, training, encoder_decoder, generative, allow_bundle_atten)\n\n targets = train_batch['target']\n pred_mask = targets != NON_TARGET_ID\n pred_indices = tf.where(pred_mask)\n\n selected_targets = tf.gather_nd(targets, pred_indices)\n \n # Compute the gradients for a list of variables.\n with tf.GradientTape() as tape:\n\n (logits, c_outputs, r_outputs) = predictor(\n train_batch,\n c_mask,\n r_mask,\n r_c_mask,\n c_r_mask,\n output_attentions=False, output_hidden_states=False, training=True\n )\n selected_logits = tf.gather_nd(logits, pred_indices)\n\n # Need to have the 1st dimension to make the losses not averaged\n losses = loss_obj(selected_targets[:, tf.newaxis], selected_logits[:, tf.newaxis])\n total_loss = tf.math.reduce_sum(losses)\n\n # `train_batch['nb_pred_places'][0]` is the total number of places used for calculating loss across replicas\n loss = total_loss / tf.cast(train_batch['nb_pred_places'][0], dtype=DTYPE)\n \n grads = tape.gradient(loss, predictor.trainable_variables)\n\n # Ask the optimizer to apply the processed gradients.\n optimizer.apply_gradients(zip(grads, predictor.trainable_variables))\n \n preds = tf.math.sigmoid(logits)\n selected_preds = tf.math.sigmoid(selected_logits)\n \n loss_metric.update_state(total_loss)\n acc_metric.update_state(selected_targets[:, tf.newaxis], selected_preds[:, tf.newaxis])\n auc_metric.update_state(selected_targets, selected_preds)\n\n batch_size = None\n @tf.function(\n input_signature=get_input_signatures(batch_size, seq_len, valid=True)\n )\n def valid_step(valid_batch):\n\n training = tf.constant(0, dtype=tf.int32)\n\n c_mask, r_mask, r_c_mask, c_r_mask = get_attention_masks(valid_batch, training, encoder_decoder, generative, allow_bundle_atten)\n \n # This only works with CPU / GPU. On TPU, XLA can't compile the graph for this part.\n# if generative == 1:\n \n# start_pos = self.train_config.window_size - MAX_PRED_TIME_QUESTION_BUNDLE_LEN\n# _, logits = predictor.generate(\n# valid_batch, start_pos=start_pos, window_size=self.train_config.window_size, dim=self.config.dim,\n# c_mask=c_mask, r_mask=r_mask, r_c_mask=r_c_mask, c_r_mask=c_r_mask\n# )\n\n# else:\n\n logits, _, _ = predictor(\n valid_batch,\n c_mask,\n r_mask,\n r_c_mask,\n c_r_mask,\n output_attentions=False, output_hidden_states=False, training=False\n )\n \n # `targets` are defined for all places (other than [PAD] and lectures).\n # However, during validation, unlike during training, we only focus on the places that are in prediction time (and being questions).\n pred_time_mask = valid_batch['pred_time_mask']\n targets = valid_batch['target']\n\n # only select the last `MAX_PRED_TIME_QUESTION_BUNDLE_LEN`\n pred_time_mask = pred_time_mask[:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:]\n targets = targets[:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:]\n logits = logits[:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:]\n row_ids = valid_batch['row_id'][:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:]\n user_ids = valid_batch['user_id'][:, tf.newaxis] * tf.ones_like(row_ids, dtype=tf.int64)\n\n targets = targets * pred_time_mask + NON_TARGET_ID * (1 - pred_time_mask)\n\n preds = tf.math.sigmoid(logits)\n\n pred_mask = targets != NON_TARGET_ID\n pred_indices = tf.where(pred_mask)\n\n selected_targets = tf.gather_nd(targets, pred_indices)\n selected_logits = tf.gather_nd(logits, pred_indices)\n selected_preds = tf.math.sigmoid(selected_logits)\n\n # Need to have the 1st dimension to make the losses not averaged\n losses = loss_obj(selected_targets[:, tf.newaxis], selected_logits[:, tf.newaxis])\n total_loss = tf.math.reduce_sum(losses)\n\n # `train_batch['nb_pred_places'][0]` is the total number of places used for calculating loss\n loss = total_loss / tf.cast(tf.math.reduce_mean(valid_batch['nb_pred_places']), dtype=DTYPE)\n\n loss_metric.update_state(total_loss)\n acc_metric.update_state(selected_targets[:, tf.newaxis], selected_preds[:, tf.newaxis])\n auc_metric.update_state(selected_targets, selected_preds)\n\n return targets, preds, row_ids, user_ids, pred_mask\n\n @tf.function\n def dist_train_step(dist_train_batch):\n \n strategy.run(train_step, args=(dist_train_batch,))\n \n @tf.function\n def dist_valid_step(dist_valid_batch):\n \n targets, preds, row_ids, user_ids, pred_mask = strategy.run(valid_step, args=(dist_valid_batch,))\n \n return targets, preds, row_ids, user_ids, pred_mask\n\n @tf.function\n def dist_train_multi_steps(dist_train_iter):\n\n total_nb_pred_places = tf.constant(0.0, dtype=DTYPE)\n \n for _ in tf.range(self.train_config.steps_per_call):\n\n dist_train_batch = next(dist_train_iter)\n dist_train_step(dist_train_batch)\n \n if type(dist_train_batch['nb_pred_places']) == tf.Tensor:\n nb_pred_places = tf.math.reduce_mean(dist_train_batch['nb_pred_places'])\n else: # type(dist_train_batch['nb_pred_places']) == tf.python.distribute.values.PerReplica\n # `dist_train_batch['nb_pred_places']` is PerPlica --> Use `.values` attribute to access tensors.\n nb_pred_places = tf.math.reduce_mean(tf.concat(dist_train_batch['nb_pred_places'].values, axis=0))\n\n total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\n\n return total_nb_pred_places\n\n @tf.function\n def dist_valid_multi_steps(dist_valid_iter):\n\n total_nb_pred_places = tf.constant(0.0, dtype=DTYPE)\n\n t_arr_1 = tf.TensorArray(tf.int32, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_2 = tf.TensorArray(DTYPE, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_3 = tf.TensorArray(tf.bool, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_4 = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_5 = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=False)\n\n for step_idx in tf.range(self.train_config.steps_per_call):\n\n dist_valid_batch = next(dist_valid_iter)\n _targets, _preds, _row_ids, _user_ids, _pred_mask = dist_valid_step(dist_valid_batch)\n \n if type(dist_valid_batch['nb_pred_places']) == tf.Tensor:\n nb_pred_places = tf.math.reduce_mean(dist_valid_batch['nb_pred_places'])\n else:\n\n _targets = tf.concat(_targets.values, axis=0)\n _preds = tf.concat(_preds.values, axis=0)\n _pred_mask = tf.concat(_pred_mask.values, axis=0)\n\n _row_ids = tf.concat(_row_ids.values, axis=0)\n _user_ids = tf.concat(_user_ids.values, axis=0)\n\n # `dist_valid_batch['nb_pred_places']` is PerPlica --> Use `.values` attribute to access tensors.\n nb_pred_places = tf.math.reduce_mean(tf.concat(dist_valid_batch['nb_pred_places'].values, axis=0))\n\n t_arr_1 = t_arr_1.write(step_idx, _targets)\n t_arr_2 = t_arr_2.write(step_idx, _preds)\n t_arr_3 = t_arr_3.write(step_idx, _pred_mask)\n t_arr_4 = t_arr_4.write(step_idx, _row_ids)\n t_arr_5 = t_arr_5.write(step_idx, _user_ids) \n \n total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\n \n targets = t_arr_1.concat()\n preds = t_arr_2.concat()\n pred_mask = t_arr_3.concat()\n row_ids = t_arr_4.concat()\n user_ids = t_arr_5.concat()\n\n return targets, preds, row_ids, user_ids, pred_mask, total_nb_pred_places\n\n @tf.function\n def dist_valid_multi_steps_last_call(dist_valid_iter):\n\n total_nb_pred_places = tf.constant(0.0, dtype=DTYPE)\n\n t_arr_1 = tf.TensorArray(tf.int32, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_2 = tf.TensorArray(DTYPE, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_3 = tf.TensorArray(tf.bool, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_4 = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=False)\n t_arr_5 = tf.TensorArray(tf.int64, size=0, dynamic_size=True, clear_after_read=False)\n\n for step_idx in tf.range(self.train_config.n_steps_in_last_valid_call):\n\n dist_valid_batch = next(dist_valid_iter)\n _targets, _preds, _row_ids, _user_ids, _pred_mask = dist_valid_step(dist_valid_batch)\n \n if type(dist_valid_batch['nb_pred_places']) == tf.Tensor:\n nb_pred_places = tf.math.reduce_mean(dist_valid_batch['nb_pred_places'])\n else:\n\n _targets = tf.concat(_targets.values, axis=0)\n _preds = tf.concat(_preds.values, axis=0)\n _pred_mask = tf.concat(_pred_mask.values, axis=0)\n\n _row_ids = tf.concat(_row_ids.values, axis=0)\n _user_ids = tf.concat(_user_ids.values, axis=0)\n\n nb_pred_places = tf.math.reduce_mean(tf.concat(dist_valid_batch['nb_pred_places'].values, axis=0))\n\n t_arr_1 = t_arr_1.write(step_idx, _targets)\n t_arr_2 = t_arr_2.write(step_idx, _preds)\n t_arr_3 = t_arr_3.write(step_idx, _pred_mask)\n t_arr_4 = t_arr_4.write(step_idx, _row_ids)\n t_arr_5 = t_arr_5.write(step_idx, _user_ids) \n \n total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\n\n targets = t_arr_1.concat()\n preds = t_arr_2.concat()\n pred_mask = t_arr_3.concat()\n row_ids = t_arr_4.concat()\n user_ids = t_arr_5.concat()\n \n return targets, preds, row_ids, user_ids, pred_mask, total_nb_pred_places\n\n n_training_steps_per_epoch = self.train_config.n_training_steps_per_epoch\n n_training_steps = self.train_config.n_epochs * n_training_steps_per_epoch\n\n if MAX_TRAIN_ITER_STEPS is not None:\n n_training_steps_per_epoch = MAX_TRAIN_ITER_STEPS\n n_training_steps = self.train_config.n_epochs * n_training_steps_per_epoch\n \n self.train_config.n_training_steps_per_epoch = n_training_steps_per_epoch\n self.train_config.n_training_steps = n_training_steps \n\n with open('train_config.json', 'w', encoding='UTF-8') as fp:\n json.dump(json.loads(self.toJSON()), fp, ensure_ascii=False, indent=4)\n if not IS_KAGGLE:\n !gsutil cp -r './train_config.json' '{self.train_config.ckpt_path}'\n \n n_epochs = self.train_config.n_epochs\n n_training_examples = self.train_config.n_training_examples\n\n print(f'n_epochs: {n_epochs}')\n print(f'n_training_examples: {n_training_examples}')\n print(f'n_training_steps_per_epoch: {n_training_steps_per_epoch}')\n print(f'n_training_steps: {n_training_steps}')\n print(f'n_valid_examples: {self.train_config.n_valid_examples}')\n print(f'n_valid_steps: {self.train_config.n_valid_steps}') \n \n training_history = dict()\n\n def train(last_epoch=0, valid_epochs=None):\n\n start_epoch = last_epoch + 1\n\n end_epoch = self.train_config.n_epochs + 1\n train_ds = self.train_ds\n\n dist_train_iter = iter(train_ds)\n\n for epoch in range(start_epoch, end_epoch):\n\n training_history[epoch] = {\n 'loss': [],\n 'acc': [],\n 'auc': [],\n 'timing': [] \n }\n\n n_steps = 0\n n_preds = 0\n n_last_preds = 0\n\n total_nb_pred_places = 0.0\n\n start_epoch_t = datetime.datetime.now()\n\n for call_idx in range(self.train_config.n_training_calls_per_epoch):\n\n start_t = datetime.datetime.now()\n\n nb_pred_places = dist_train_multi_steps(dist_train_iter) \n n_steps += self.train_config.steps_per_call\n\n total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\n \n avg_loss = (loss_metric.result() / total_nb_pred_places).numpy()\n acc = acc_metric.result().numpy()\n auc = auc_metric.result().numpy()\n\n # -------------------------------------------------- \n\n if (call_idx + 1) % max(1, self.train_config.n_training_calls_per_epoch // 100) == 0:\n\n training_history[epoch]['loss'].append(float(avg_loss))\n training_history[epoch]['acc'].append(float(acc))\n training_history[epoch]['auc'].append(float(auc))\n\n # -------------------------------------------------- \n\n elapsed = (datetime.datetime.now() - start_t).total_seconds()\n\n if (call_idx + 1) % max(1, self.train_config.n_training_calls_per_epoch // 10) == 0:\n\n print(f'epoch: {epoch} - step: {n_steps}')\n print(f'timing per step: {elapsed / self.train_config.steps_per_call}')\n print(f'loss: {avg_loss}')\n print(f'acc: {acc}')\n print(f'auc: {auc}')\n print(f'lr: {optimizer._decayed_lr(DTYPE)}') \n print('-' * 80)\n\n end_epoch_t = datetime.datetime.now()\n elapsed_epoch_t = (end_epoch_t - start_epoch_t).total_seconds()\n training_history[epoch]['timing'] = elapsed_epoch_t\n\n training_history[epoch]['train_loss'] = float(avg_loss)\n training_history[epoch]['train_acc'] = float(acc)\n training_history[epoch]['train_auc'] = float(auc)\n\n # In order to save the checkpoints.\n if IS_KAGGLE and tpu:\n ckpt_manager.checkpoint.save(\n file_prefix=ckpt_manager.directory + 'ckpt', options=tf.train.CheckpointOptions(experimental_io_device=\"/job:localhost\")\n )\n else:\n ckpt_manager.save()\n \n if not IS_KAGGLE and self.train_config.ckpt_path.startswith('gs://') and tpu is None:\n !gsutil cp -r './checkpoint' '{self.train_config.ckpt_path}'\n !gsutil cp -r './ckpt-{epoch}.index' '{self.train_config.ckpt_path}'\n !gsutil cp -r './ckpt-{epoch}.data-00000-of-00001' '{self.train_config.ckpt_path}'\n !rm -rf './ckpt-{epoch}.index'\n !rm -rf './ckpt-{epoch}.data-00000-of-00001' \n\n print(f'epoch: {epoch} - train')\n print(f'train loss: {avg_loss}')\n print(f'train acc: {acc}')\n print(f'train auc: {auc}')\n print('-' * 80)\n\n loss_metric.reset_states()\n acc_metric.reset_states()\n auc_metric.reset_states()\n\n # --------------------------------------------------\n # saving\n\n fn = f'training_history-{epoch}.json'\n with open(fn, 'w', encoding='UTF-8') as fp:\n json.dump(training_history, fp, ensure_ascii=False, indent=4)\n\n if not IS_KAGGLE:\n !gsutil cp -r './{fn}' '{self.train_config.ckpt_path}'\n !rm -rf './{fn}'\n\n # --------------------------------------------------\n # validation\n\n if valid_epochs is None:\n valid_epochs = [self.train_config.n_epochs]\n \n if epoch in valid_epochs:\n valid(epoch)\n\n def valid(epoch):\n\n if epoch not in training_history:\n\n training_history[epoch] = {\n 'loss': [],\n 'acc': [],\n 'auc': [],\n 'timing': [] \n } \n\n valid_user_ids = []\n valid_row_ids = []\n valid_targets = []\n valid_preds = []\n\n n_steps = 0\n total_nb_pred_places = 0.0\n\n dist_valid_iter = iter(self.valid_ds)\n\n start_valid = datetime.datetime.now()\n\n for call_idx in range(self.train_config.n_valid_calls):\n\n start_t = datetime.datetime.now()\n\n last_call = (self.train_config.n_steps_in_last_valid_call > 0) and (call_idx == self.train_config.n_valid_calls - 1)\n \n if last_call:\n targets, preds, row_ids, user_ids, pred_mask, nb_pred_places = dist_valid_multi_steps_last_call(dist_valid_iter)\n n_steps += self.train_config.n_steps_in_last_valid_call\n else:\n targets, preds, row_ids, user_ids, pred_mask, nb_pred_places = dist_valid_multi_steps(dist_valid_iter) \n n_steps += self.train_config.steps_per_call\n \n total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\n \n pred_indices = tf.where(pred_mask)\n\n selected_targets = tf.gather_nd(targets, pred_indices)\n selected_preds = tf.gather_nd(preds, pred_indices)\n\n selected_row_ids = tf.gather_nd(row_ids, pred_indices)\n selected_user_ids = tf.gather_nd(user_ids, pred_indices)\n\n avg_loss = (loss_metric.result() / total_nb_pred_places).numpy()\n acc = acc_metric.result().numpy()\n auc = auc_metric.result().numpy()\n\n elapsed = (datetime.datetime.now() - start_t).total_seconds()\n\n if call_idx % max(1, self.train_config.n_valid_calls // 10) == 0:\n\n print(f'epoch: {epoch} - valid step: {n_steps}')\n print(f'valid timing per steps: {elapsed / self.train_config.steps_per_call}')\n print(f'valid loss: {avg_loss}')\n print(f'valid acc: {acc}')\n print(f'vald auc: {auc}') \n print('-' * 80)\n \n valid_row_ids.extend(selected_row_ids.numpy().tolist())\n valid_user_ids.extend(selected_user_ids.numpy().tolist())\n valid_targets.extend(selected_targets.numpy().tolist())\n valid_preds.extend(selected_preds.numpy().tolist())\n\n end_valid = datetime.datetime.now()\n elapsed_valid = (end_valid - start_valid).total_seconds()\n training_history[epoch]['valid_timing'] = elapsed_valid\n\n training_history[epoch]['valid_loss'] = float(avg_loss)\n training_history[epoch]['valid_acc'] = float(acc)\n training_history[epoch]['valid_auc'] = float(auc)\n\n print(f'epoch: {epoch} - valid')\n print(f'valid loss: {avg_loss}')\n print(f'valid acc: {acc}')\n print(f'vald auc: {auc}') \n print('=' * 80)\n\n loss_metric.reset_states()\n acc_metric.reset_states()\n auc_metric.reset_states()\n \n # --------------------------------------------------\n\n valid_submission = pd.DataFrame.from_dict(\n {\n 'row_ids': valid_row_ids,\n 'user_ids': valid_user_ids,\n 'targets': valid_targets,\n 'preds': valid_preds\n }\n )\n valid_submission.to_csv(f'valid_submission_epoch_{epoch}.csv', index=False)\n\n if not IS_KAGGLE:\n !gsutil cp -r './valid_submission_epoch_{epoch}.csv' '{self.train_config.ckpt_path}' \n !rm -rf './valid_submission_epoch_{epoch}.csv'\n\n # --------------------------------------------------\n\n fn = f'training_history-{epoch}.json'\n if only_valid:\n fn = f'training_history-{epoch}-only-valid.json'\n with open(fn, 'w', encoding='UTF-8') as fp:\n json.dump(training_history, fp, ensure_ascii=False, indent=4)\n\n if not IS_KAGGLE:\n !gsutil cp -r './{fn}' '{self.train_config.ckpt_path}'\n !rm -rf './{fn}'\n\n # --------------------------------------------------\n\n if only_valid:\n valid(last_epoch)\n else:\n train(last_epoch, valid_epochs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train / Valid"},{"metadata":{},"cell_type":"markdown","source":"## A method to make the model initialize its variables"},{"metadata":{"id":"tH-cnl3Le3HT","trusted":true},"cell_type":"code","source":" def run_dummy_inputs(predictor):\n\n @tf.function\n def foo(inputs):\n\n logits, _, _ = predictor(inputs=inputs)\n return logits\n\n with strategy.scope():\n\n content_input_ids = tf.constant(1, shape=[3, 5])\n response_input_ids = tf.constant(1, shape=[3, 5])\n d_input_ids = tf.constant(1, shape=[3, 5])\n tag_ids = tf.constant(0, shape=[3, 5, N_TAGS_PER_CONTENT])\n part_ids = tf.constant(1, shape=[3, 5])\n prior_explanation_ids = tf.constant(1, shape=[3, 5])\n prior_question_elapsed_time_input = tf.constant(1.0, shape=[3, 5])\n \n inputs = {\n 'c_input_ids': content_input_ids,\n 'r_input_ids': response_input_ids,\n 'd_input_ids': decoder_input_ids,\n 'tag_ids': tag_ids,\n 'part_ids': part_ids,\n 'prior_explanation_ids': prior_explanation_ids,\n 'prior_question_elapsed_time_input': prior_question_elapsed_time_input \n }\n\n logits = foo(inputs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## A method to load the checkpoints"},{"metadata":{"trusted":true,"id":"Ltf4Pc_z11w7"},"cell_type":"code","source":"def load_ckpt(predictor, optimizer, ckpt_path, ckpt_no=-1):\n\n if ckpt_path.startswith('gs://') and tpu is None:\n\n assert ckpt_no >= 1\n\n f1 = f'{ckpt_path}checkpoint'\n f2 = f'{ckpt_path}ckpt-{ckpt_no}.index'\n f3 = f'{ckpt_path}ckpt-{ckpt_no}.data-00000-of-00001'\n\n !gsutil cp -r {f1} './'\n !gsutil cp -r {f2} './'\n !gsutil cp -r {f3} './'\n\n ckpt_path = './'\n\n with strategy.scope():\n\n # ----------------------------------------------------------------------\n # Init model's weight if necessary\n\n run_dummy_inputs(predictor) \n\n # ----------------------------------------------------------------------\n\n if optimizer is None:\n optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5, beta_1=0.9, beta_2=0.999, epsilon=1e-08)\n checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=predictor)\n ckpt_manager = tf.train.CheckpointManager(checkpoint, directory=ckpt_path, max_to_keep=None)\n\n last_ckpt_no = 0\n if ckpt_manager.latest_checkpoint:\n if len(ckpt_manager.latest_checkpoint.split('/ckpt-')) > 0:\n last_ckpt_no = int(ckpt_manager.latest_checkpoint.split('/ckpt-')[-1])\n \n if last_ckpt_no > 0:\n print(f'Latest checkpoint found. Model trained for {last_ckpt_no} epochs.')\n else:\n print(\"No checkpoint found.\")\n\n if ckpt_no == -1:\n ckpt_no_to_load = last_ckpt_no\n else:\n ckpt_no_to_load = ckpt_no\n\n assert ckpt_no_to_load <= last_ckpt_no\n\n ckpt_file = os.path.join(ckpt_path, f'ckpt-{ckpt_no_to_load}')\n\n # ----------------------------------------------------------------------\n # Load ckpt\n\n print(f'try to load {ckpt_file}')\n if ckpt_no_to_load > 0:\n status = checkpoint.restore(ckpt_file)\n print(f'ckpt-{ckpt_no_to_load} is restored.')\n loaded_ckpt_no = ckpt_no_to_load\n else:\n print(f'no ckpt is found and restored.')\n loaded_ckpt_no = 0\n\n return ckpt_manager, loaded_ckpt_no","execution_count":null,"outputs":[]},{"metadata":{"id":"6tsgdgCu11xI"},"cell_type":"markdown","source":"### Train / Valid / Pred configurations"},{"metadata":{"id":"fWVxPT8E2Vpb","trusted":true},"cell_type":"code","source":"max_position_embeddings = MAX_HISTORY_LEN\nif not USE_ABS_POS:\n # `0` is used for padding, the real position is from `1` to `WINDOW_SIZE`.\n max_position_embeddings = WINDOW_SIZE + 1\n\nconfig = EdFormerConfig(\n model_type=MODEL_TYPE,\n model_desc=MODEL_DESC,\n model_size=MODEL_SIZE,\n content_vocab_size=CONTENT_VOCAB_SIZE,\n response_vocab_size=RESPONSE_VOCAB_SIZE,\n tag_vocab_size=TAG_VOCAB_SIZE,\n part_vocab_size=PART_VOCAB_SIZE,\n prior_explanation_vocab_size=PRIOR_EXPLANATION_VOCAB_SIZE,\n use_prior_question_elapsed_time_input=USE_PRIOR_QUESTION_ELAPSED_TIME_INPUT,\n max_position_embeddings=max_position_embeddings,\n sinusoidal_pos_embds=False,\n n_layers=N_LAYERS,\n n_heads=N_HEADS,\n dim=DIM,\n hidden_dim=HIDDEN_DIM,\n activation=ACTIVATION, \n dropout=0.1,\n attention_dropout=0.1,\n seq2seq_dropout=0.1,\n initializer_range=0.02,\n seed=SEED,\n pad_token_id=PAD_ID,\n use_abs_pos=USE_ABS_POS,\n share_position_embeddings=SHARE_POS_EMBEDDING,\n use_tags=USE_TAGS,\n use_part=USE_PART,\n use_prior_explanation=USE_PRIOR_EXPLANATION,\n allow_bundle_atten=ALLOW_BUNDLE_ATTEN,\n generative=GENERATIVE\n)\n\ntrain_config = TrainConfig(\n ckpt_path=CKPT_TRAIN_PATH,\n window_size=WINDOW_SIZE,\n loss_weight_window_size=LOSS_WEIGHT_WINDOW_SIZE,\n n_epochs=N_EPOCHS,\n shuffle_buf_size=SHUFFLE_BUFFER_SIZE,\n batch_size=BATCH_SIZE,\n pred_batch_size=PRED_BATCH_SIZE,\n seed=SEED,\n deterministic=DETERMINISTIC,\n num_parallel_reads=N_PARALLEL_READS,\n num_parallel_calls=N_PARALLEL_CALLS,\n steps_per_call=STEPS_PER_CALL\n)","execution_count":null,"outputs":[]},{"metadata":{"id":"fQ2AaCkjWjq9"},"cell_type":"markdown","source":"### Now comes the training!"},{"metadata":{"trusted":true,"id":"AxvDlySG11yq","outputId":"f447f165-7644-4e97-cd57-54bc10db2b75"},"cell_type":"code","source":"if TRAIN:\n \n only_valid = False\n ckpts = [0]\n valid_epochs = [3, 6]\n\n train_manager = Train_Manager(config, train_config)\n predictor, optimizer, loss_obj, metrics = train_manager.get_train_objs(lr=LR, end_lr=END_LR, warmup_steps=WARMUP_STEPS)\n \n # ----------------------------------------------------------------------------------------------------\n\n run_dummy_inputs(predictor)\n\n print('-' * 40)\n print('trainable variables:')\n for v in predictor.trainable_variables:\n print(v.name)\n print('-' * 40)\n\n for ckpt_no in ckpts:\n\n # ----------------------------------------------------------------------------------------------------\n\n if RESUME_TRAINING or only_valid:\n\n # reload ckpt\n ckpt_manager, loaded_ckpt_no = load_ckpt(predictor, optimizer, CKPT_TRAIN_PATH, ckpt_no=ckpt_no)\n\n else:\n\n with strategy.scope():\n\n loaded_ckpt_no = 0\n checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=predictor)\n ckpt_path = CKPT_TRAIN_PATH\n if ckpt_path.startswith('gs://') and tpu is None:\n ckpt_path = './'\n ckpt_manager = tf.train.CheckpointManager(checkpoint, directory=ckpt_path, max_to_keep=None)\n\n train_manager.train_valid(predictor, optimizer, ckpt_manager, loss_obj, metrics, last_epoch=loaded_ckpt_no, valid_epochs=valid_epochs, only_valid=only_valid)","execution_count":null,"outputs":[]},{"metadata":{"id":"2COvA_dof18z","trusted":true},"cell_type":"code","source":"CKPT_TRAIN_PATH","execution_count":null,"outputs":[]},{"metadata":{"id":"ZN8mrPjvJOuR","trusted":true},"cell_type":"code","source":"!ls -l","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Conclusion\n\n1. We compared the usage of single/multiple steprs per TPU call. From [Single step per TPU call](#https://www.kaggle.com/yihdarshieh/r3id-tf-tpu/output?scriptVersionId=49071414) and [1000 steps per TPU call](#https://www.kaggle.com/yihdarshieh/r3id-tf-tpu/output?scriptVersionId=49075674), we see that:\n\n timing per step: (with batch sizes 128/256)\n * single step per TPU call: \n training: 178 ms\n validation: 93 ms\n * 1000 steps per TPU call: 32 m\n training: 32 ms\n validation: 65 ms\n\nIt is evident that multiple steprs per TPU call speed up the training (5.5x faster) - and a bit faster for validation. For validation, since `PRED_BATCH_SIZE = 256` is twice large as \n`BATCH_SIZE = 128`, the timing for train/valid with `1000 steps per TPU call` is almost the same if they would have the same batch size. It is not clear for me why the difference for validatioin timining is less large than the difference for training - potentially due to the overhead coming from accumulating the predictions. For this version, we use larger batch size (512/2048) - this should reduce further the timing (once normalized to the same batch size during comparison).\n\n2. We see how to save checkpoints locally on Kaggle via `tf.train.CheckpointManager`.\n\n3. We learn how to accumulate the validation predictions while using TPU with multi-stpes by using `tf.TensorArray`.\n\n4. An input pipeline with `tf.RaggedTensor` and a series of dataset transformations to obtain the training/validation dataset formats specific to this competition is demonstrated.\n\n5. Several model options are provided, in particular:\n\n * encoder only\n * encoder decoder \n - treat each question in a bundle during inference as a single question in the bundle\n - auto-regressive prediction generation\n \n6. A notebook works both on Kaggle and Colab (with a minimal necessary change).\n\nHope you enjoy it. Good luck for the competition!"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}