{"cells":[{"metadata":{"id":"UnkZKyp_11iC"},"cell_type":"markdown","source":"# About this notebook"},{"metadata":{"id":"CKfIrA-M3JNc"},"cell_type":"markdown","source":"# Environment"},{"metadata":{"trusted":true},"cell_type":"code","source":"import datetime\n\nkernel_start = datetime.datetime.now()\n\nprint(f'kernel_start: {kernel_start}')","execution_count":null,"outputs":[]},{"metadata":{"id":"wMXsMdHE34ZE","outputId":"9f74ebdd-cad1-4058-f0b8-817d3c599523","trusted":true},"cell_type":"code","source":"import os\n\nIS_KAGGLE = os.path.isdir('/kaggle/input')\n\n\n# 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.\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'IS_KAGGLE: {IS_KAGGLE}')        \nprint(f'USE_GPU: {USE_GPU}')","execution_count":null,"outputs":[]},{"metadata":{"id":"Vu1E8hiGHZTZ","outputId":"c335aeea-5bde-4cc9-895f-0b1be35e68c8","trusted":true},"cell_type":"code","source":"!pip uninstall -y cloud-tpu-client\n\nif not IS_KAGGLE:\n\n    !pip uninstall -y tensorflow\n    !pip install --upgrade tensorflow==2.3.0\n    !pip uninstall -y tensorflow-gcs-config\n    !pip install --upgrade tensorflow-gcs-config==2.3.0\n    \n\nelif IS_KAGGLE and not USE_GPU:\n    # Kaggle TPU\n    \n    !pip uninstall -y tensorflow\n    !pip install --upgrade tensorflow==2.3.1\n\nimport tensorflow as tf\nfrom tensorflow.python.tpu.client import client\n\nprint(f'TensorFlow: {tf.__version__}')\nprint(client.Client)","execution_count":null,"outputs":[]},{"metadata":{"id":"42P7vuxWpQYT","outputId":"88ece806-fbc0-468c-e912-138c98f39989","trusted":true},"cell_type":"code","source":"if IS_KAGGLE:\n    BASE_DIR = '/kaggle/input'\nelse:\n    BUCKET_DIR = 'gs://shieh-tpu/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 = 'shieh-tpu'\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://shieh-tpu/r3id' \"/content/drive/My Drive\"\n\n!ls -l '{BASE_DIR}'","execution_count":null,"outputs":[]},{"metadata":{"id":"Dz-MF4fX66sD","outputId":"613521f4-7a3b-400e-9a97-3c8e0e8de92c","trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\n\ntry:\n    client.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    if IS_KAGGLE:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver() \n    else:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR']) \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":"qquMF3i1pNTE","outputId":"22313d52-2986-4485-cec7-b63208b49bd6","trusted":true},"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\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'}'","execution_count":null,"outputs":[]},{"metadata":{"id":"5VXzmicseiDu"},"cell_type":"markdown","source":"# Packages"},{"metadata":{"id":"EbTukhm4qeyt"},"cell_type":"markdown","source":"## Install"},{"metadata":{"id":"ufr8QIu5qg-Y"},"cell_type":"markdown","source":"## Import"},{"metadata":{"id":"7E5iDmMK11jf","trusted":true},"cell_type":"code","source":"import os\nimport psutil\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport datatable as dt\nimport json\nimport pickle\nfrom collections import defaultdict\nimport random\nimport math\nimport datetime\nimport collections\nos.environ['TF_DETERMINISTIC_OPS'] = '1'\nif IS_KAGGLE and USE_GPU:\n    import tensorflow_probability as tfp\nfrom copy import deepcopy\nimport gc\ngc.enable()\n\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    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"},{"metadata":{"id":"OuB4HtjyqyGw"},"cell_type":"markdown","source":"# Data"},{"metadata":{"id":"l8LuQDTier5n","trusted":true},"cell_type":"code","source":"train_dt_path_or_obj = f'{BASE_DIR}/r3id-trainjay/train.jay'\n\nunique_user_id_train_path = f'{BASE_DIR}/r3id-info/unique_user_id_train.json'\nunique_question_id_train_path = f'{BASE_DIR}/r3id-info/unique_question_id_train.json'\nunique_lecture_id_train_path = f'{BASE_DIR}/r3id-info/unique_lecture_id_train.json'\nuser_id_to_row_id_train_path = f'{BASE_DIR}/r3id-info/user_id_to_row_id_train.json'\n\nunique_user_id_splitted_train_path = f'{BASE_DIR}/r3id-info-shieh/unique_user_id_splitted_train.json'\nunique_question_id_splitted_train_path = f'{BASE_DIR}/r3id-info-shieh/unique_question_id_splitted_train.json'\nunique_lecture_id_splitted_train_path = f'{BASE_DIR}/r3id-info-shieh/unique_lecture_id_splitted_train.json'\n\nn_contents_dict_path = f'{BASE_DIR}/r3id-info-shieh/n_contents_dict.json'\n\nquestion_tags_info_path = f'{BASE_DIR}/r3id-info-shieh/question_tags_info.json'\nlecture_tags_info_path = f'{BASE_DIR}/r3id-info-shieh/lecture_tags_info.json'\n\nquestion_part_info_path = f'{BASE_DIR}/r3id-info-shieh/question_part_info.json'\nlecture_part_info_path = f'{BASE_DIR}/r3id-info-shieh/lecture_part_info.json'\n\ncorrect_answer_info_path = f'{BASE_DIR}/r3id-info-shieh/correct_answer_info.json'\n\nquestion_history_path = f'{BASE_DIR}/r3id-info-shieh/question_history.json'\nquestion_history_at_training_end_path = f'{BASE_DIR}/r3id-info-shieh/question_history_at_training_end.json'\nsingle_question_history_at_training_end_optimized_path = f'{BASE_DIR}/r3id-info-shieh/single_question_history_at_training_end.json'\n\nuser_performance_hdf5_path = f'{BASE_DIR}/user-performance/user_performance.hdf5'\n\n# ------------------------------------------------------------------------------------------\n# this should be replaced with the corresponding files for validation\nquestion_history_for_valid_path = f'{BASE_DIR}/r3id-info-shieh/question_history_fold_1.json'\nquestion_history_at_training_end_for_valid_path = f'{BASE_DIR}/r3id-info-shieh/question_history_at_training_end_fold_1.json'\nsingle_question_history_at_training_end_optimized_for_valid_path = f'{BASE_DIR}/r3id-info-shieh/single_question_history_at_training_end_fold_1.json'\n\n# ------------------------------------------------------------------------------------------\n\nvalid_info_paths = [\n    f'{BASE_DIR}/r3id-info-shieh/valid_info.json',\n    f'{BASE_DIR}/r3id-info-shieh/valid_info_fold_1.json',\n    f'{BASE_DIR}/r3id-info-shieh/valid_info_fold_2.json',\n    f'{BASE_DIR}/r3id-info-shieh/valid_info_fold_3.json',\n]\n\ntrain_valid_split_indices_paths = [\n        f'{BASE_DIR}/r3id-info-shieh/train_valid_split_indices.json',\n        f'{BASE_DIR}/r3id-info-shieh/train_valid_split_indices_fold_1.json',\n        f'{BASE_DIR}/r3id-info-shieh/train_valid_split_indices_fold_2.json',\n        f'{BASE_DIR}/r3id-info-shieh/train_valid_split_indices_fold_3.json',\n]\n\n# train_tfrec_dir_local = f'{BASE_DIR}/ednet-tfrecords-sequential-more'\n# valid_tfrec_dir_local = f'{BASE_DIR}/r3id-tfrecords-valid-more'\n\nif tpu is None:\n    \n    TFREC_DIR = BASE_DIR\n    train_tfrec_dir = f'{BASE_DIR}/ednet-tfrecords-sequential-more'\n    valid_tfrec_dir = f'{BASE_DIR}/r3id-tfrecords-valid-more'\n\nelif IS_KAGGLE:\n    \n    train_tfrec_dir = KaggleDatasets().get_gcs_path('ednet-tfrecords-sequential-more')\n    valid_tfrec_dir = KaggleDatasets().get_gcs_path('r3id-tfrecords-valid-more')\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    TFREC_DIR = BUCKET_DIR\n    train_tfrec_dir = f'{BUCKET_DIR}/ednet-tfrecords-sequential-more'\n    valid_tfrec_dir = f'{BUCKET_DIR}/r3id-tfrecords-valid-more'\n\n# Load the big file - train.dt - once\ntrain_dt = dt.fread(train_dt_path_or_obj)","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 = 'b-2'\nMODEL_DESC = 'raw'\n\nACTIVATION = 'relu'\nUSE_PRE_CLASSIFIER = True\nUSE_SOFTMAX = False\nUSE_USER_ANSWER = False\nUSE_USER_ANSWER_LOSS = False\nUSE_CORRECT_ANSWER_FOR_ENCODER = False\nUSE_CORRECT_ANSWER_FOR_DECODER = False\nUSE_ABS_POS = False\nUSE_TASK_CONTAINER_POS = False\nSHARE_POS_EMBEDDING = True\nUSE_TAGS = False\nUSE_PART = False\nUSE_PRIOR_EXPLANATION = False\nUSE_PRIOR_QUESTION_ELAPSED_TIME_INPUT = False\nUSE_LAG_TIME = False\nUSE_LAG_TIME_FOR_ENCODER = False\nUSE_USER_LEVEL_AGGREGATED_HISTORICAL_INFO = False\nUSE_PART_AGGREGATED_HISTORICAL_INFO = False\nUSE_CORRECT_ANSWER_AGGREGATED_HISTORICAL_INFO = False\nUSE_QUESTION_LEVEL_AGGREGATED_HISTORICAL_INFO = False\nALLOW_BUNDLE_ATTEN = False\nGENERATIVE = False\n\nVALID_FOLD = 0\n\nWINDOW_SIZE = 128\nLOSS_WEIGHT_WINDOW_SIZE = None\nBATCH_SIZE = 16 * strategy.num_replicas_in_sync\nif not tpu:\n    BATCH_SIZE = 8 * BATCH_SIZE\nPRED_BATCH_SIZE = 256 * strategy.num_replicas_in_sync\nif tpu:\n    assert BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\nassert PRED_BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\nN_EPOCHS = 2\nSTEPS_PER_CALL = 1000\nMAX_N_CONTENTS_PER_USER_FOR_SAMPLING_PROB = 512\n\nLR = 1e-4\nEND_LR = 1e-4\nWARMUP_STEPS = 5000\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}-{MODEL_DESC}/'","execution_count":null,"outputs":[]},{"metadata":{"id":"zI-dw1serob_"},"cell_type":"markdown","source":"## Running mode"},{"metadata":{"id":"mp1I7Ar811kZ","trusted":true},"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":{"id":"o77CsqXW11ld","outputId":"624356cd-271b-47bb-9bf9-d123ed10511f","trusted":true},"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":{"id":"BlcYdBBAr2_h"},"cell_type":"markdown","source":"## Vocabulary settings"},{"metadata":{"id":"w9tdC-J811mN","trusted":true},"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\nANSWER_LECTURE_TOKEN = -1\nANSWER_0_TOKEN = 0\nANSWER_1_TOKEN = 1\nANSWER_2_TOKEN = 2\nANSWER_3_TOKEN = 3\n\n\nRESPONSE_LECTURE_ID = 4\nRESPONSE_FALSE_ID = 5\nRESPONSE_TRUE_ID = 6\n\nANSWER_LECTURE_ID = 4\nANSWER_0_ID = 5\nANSWER_1_ID = 6\nANSWER_2_ID = 7\nANSWER_3_ID = 8\n\nNON_TARGET_ID = -100\n\nTAG_VOCAB_SIZE = 189  # Original: [0, 2, ..., 187]. Size is `189` because we will consider `-1` for special vocab.\nPART_VOCAB_SIZE = 9  # Original: [1, 2, ..., 7]. We include `0` although this doesn't exist in the dataset. Size is `9` because we will consider `-1` for special vocab.\nPRIOR_EXPLANATION_VOCAB_SIZE = 3  # Original: [0, 1]. Size is `3` because we will consider `-1` for special vocab.\n\n# --------------------------------------------------\n\nN_TAGS_PER_CONTENT = 6\n\n# --------------------------------------------------\n# To be removed once decided not to use abs_pos\nMAX_TRAIN_HISTORY_LEN = 17917\n# This is only an assumption (including `0` for padding)\nMAX_HISTORY_LEN = 20480\n# --------------------------------------------------\n\n# Maybe to probe this info.\n# Actually, it is `9999`\nMAX_TASK_CONTAINER_ID = 10000\n\n# --------------------------------------------------\n\nN_AGGREGATED_QUESTION_SCALING_FACTOR = 1000","execution_count":null,"outputs":[]},{"metadata":{"id":"cl_fxZCOLdkd"},"cell_type":"markdown","source":"# Probe Info"},{"metadata":{"id":"qcarGDRkLdG-","trusted":true},"cell_type":"code","source":"MAX_PRED_TIME_QUESTION_BUNDLE_LEN = 10","execution_count":null,"outputs":[]},{"metadata":{"id":"yDB5LrYw11pg"},"cell_type":"markdown","source":"# Train Manager"},{"metadata":{"id":"34iiHe4c11pk"},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"id":"CEwdpMQkfrIG"},"cell_type":"markdown","source":"### Tables mapping question / lectures / tags ids (etc.) to encoder / decoder input ids"},{"metadata":{"id":"LDpq9eIDfeJf","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.json') or 'valid_info_fold_' in path:\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('user_id_to_row_id_train.json') or 'train_valid_split_indices' in path:\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        elif path.endswith('correct_answer_info.json'):           \n            data = {int(k): v for k, v in data.items()}\n              \n        elif path.endswith('question_history.json'):           \n            pass      \n                \n        elif path.endswith('question_history_at_training_end.json'):           \n            pass      \n            \n        elif path.endswith('single_question_history_at_training_end.json'):           \n            pass             \n              \n        return data\n        \n    elif path.endswith('train.jay'):\n        \n        train_dt = dt.fread(train_dt_path_or_obj)\n        \n        return train_dt\n\n    else:\n\n        raise ValueError('The path provided is not the one needed to create an instance of `Valid_Manager`!')","execution_count":null,"outputs":[]},{"metadata":{"id":"SAKL8S2Ffh-j","trusted":true},"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_TOKEN, RESPONSE_FALSE_TOKEN]\n\nanswer_vocab = [\n    PAD_TOKEN, START_TOKEN, END_TOKEN, MASK_TOKEN,\n    ANSWER_LECTURE_TOKEN,\n    ANSWER_0_TOKEN, ANSWER_1_TOKEN, ANSWER_2_TOKEN, ANSWER_3_TOKEN\n]\n\nCONTENT_VOCAB_SIZE = len(content_vocab)\nRESPONSE_VOCAB_SIZE = len(response_vocab)\nANSWER_VOCAB_SIZE = len(answer_vocab)\n\nquestion_tags_info = load_data(question_tags_info_path)\nlecture_tags_info = load_data(lecture_tags_info_path)\ncorrect_answer_info = load_data(correct_answer_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 = []\ncorrect_answer_table = []\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        correct_answer_table.append(idx)\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        correct_answer_table.append(correct_answer_info[question_id] + ANSWER_0_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        correct_answer_table.append(ANSWER_LECTURE_ID)\n\nc_inputs_ids_to_tags = tf.constant(tags_database, dtype=tf.int32)\nc_inputs_ids_to_part = tf.constant(part_database, dtype=tf.int32)\nc_inputs_ids_to_correct_answer_id = tf.constant(correct_answer_table, dtype=tf.int32)\n\n# --------------------------------------------------------------------------------------------------------------\n# This is much faster!\nquestion_id_to_input_id_dict = {\n    k: v for k, v in zip(question_ids, range(len(special_vocab), len(special_vocab) + len(question_vocab)))\n}\nc_inputs_ids_to_part_dict = {k: v for k, v in enumerate(part_database)}\nc_inputs_ids_to_correct_answer_id_dict = {k: v for k, v in enumerate(correct_answer_table)}","execution_count":null,"outputs":[]},{"metadata":{"id":"Dc7tQ-PIfzdz","outputId":"8d559922-2e74-428c-fdc6-1708dc9bcb87","trusted":true},"cell_type":"code","source":"print(f'CONTENT_VOCAB_SIZE: {CONTENT_VOCAB_SIZE}')\nprint(f'RESPONSE_VOCAB_SIZE: {RESPONSE_VOCAB_SIZE}')\nprint(f'ANSWER_VOCAB_SIZE: {ANSWER_VOCAB_SIZE}')","execution_count":null,"outputs":[]},{"metadata":{"id":"njJiZx_p11pn"},"cell_type":"markdown","source":"### TFRecord files"},{"metadata":{"id":"fxD0l4Vr11ps","outputId":"007f3c0e-7af0-493c-bd91-4917491eb3b9","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    train_tfrec_fns = [f'EdNet-user-history-{idx}.tfrecord' for idx in range(41)]\n    # Sort the file - For verification purpose\n    train_tfrec_fns = sorted(train_tfrec_fns, key=lambda x: int(x.replace('EdNet-user-history-', '').replace('.tfrecord', '')))\n\n    if tpu is None or IS_KAGGLE:\n        train_tfrec_paths = [os.path.join(train_tfrec_dir, fn) for fn in train_tfrec_fns]\n    else:\n        train_tfrec_paths = [train_tfrec_dir + f'/{fn}' for fn in train_tfrec_fns]\n\n    train_tfrec_paths","execution_count":null,"outputs":[]},{"metadata":{"id":"PT9m690JLIAk","outputId":"84d37bfe-e1b4-4473-cbe7-792611a81379","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    valid_tfrec_fns = [\n        'EdNet-user-history-with-valid-block-pos.tfrecord',\n        # 'EdNet-user-history-with-valid-block-pos-fold-1.tfrecord',\n        # 'EdNet-user-history-with-valid-block-pos-fold-2.tfrecord',\n        # 'EdNet-user-history-with-valid-block-pos-fold-3.tfrecord',\n    ]\n\n    if tpu is None:\n        valid_tfrec_paths = [os.path.join(valid_tfrec_dir, fn) for fn in valid_tfrec_fns]\n    else:\n        valid_tfrec_paths = [valid_tfrec_dir + f'/{fn}' for fn in valid_tfrec_fns]\n\n    valid_tfrec_paths","execution_count":null,"outputs":[]},{"metadata":{"id":"old9ZSlp11qH"},"cell_type":"markdown","source":"### Load TFRecord files - tf.io.RaggedFeature"},{"metadata":{"id":"1wcmYrGh11qK","trusted":true},"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    # extra information\n    'n_prev_seen': tf.io.RaggedFeature(value_key='n_prev_seen', dtype=tf.int64),\n    'n_prev_correctness': tf.io.RaggedFeature(value_key='n_prev_correctness', dtype=tf.int64),\n    # -------------------------------------------------------------------\n\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    # This makes `parsed['shifted_answered_correctly']` has 1 more extra element at this moment.\n    parsed['shifted_answered_correctly'] = tf.concat([[START_TOKEN], parsed['answered_correctly']], axis=0)\n\n    # We need to use `START_TOKEN` rather than `PAD_TOKEN` here.\n    # This makes `parsed['shifted_user_answer']` has 1 more extra element at this moment.\n    parsed['shifted_user_answer'] = tf.concat([[START_TOKEN], parsed['user_answer']], 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    # ----------------------------------------------------------------------------------------------------\n    # To get the lag time\n\n    diff_t = parsed['timestamp'] - tf.concat([[1], parsed['timestamp'][:-1]], axis=0)\n    bundle_start_mask = tf.cast(diff_t != 0, dtype=tf.int32)\n    # The jump values at bundle starting places\n    jump_t = tf.boolean_mask(diff_t, bundle_start_mask > 0)\n    # The task_container_id at bundle starting places\n    task_container_id_at_jumps = tf.boolean_mask(parsed['task_container_id'], bundle_start_mask > 0)\n    # Maps task_container_id to jump values\n    task_containder_id_to_jump_t = tf.scatter_nd(\n        indices=task_container_id_at_jumps[:, tf.newaxis], updates=jump_t, shape=tf.math.reduce_max(task_container_id_at_jumps + 1)[tf.newaxis], name=None\n    )\n    # `Our own version` of lag time.\n    lag_time = tf.gather(params=task_containder_id_to_jump_t , indices=parsed['task_container_id'])\n    # Make sure >= `0`\n    lag_time = tf.math.maximum(tf.cast(0, dtype=tf.int64), lag_time)\n\n    # set to 1 second if `lag_time == 0` but they are not in the same bunlde \n    _mask = tf.cast(tf.math.logical_and(lag_time == 0, parsed['timestamp'] > 0), dtype=tf.int64) \n    lag_time = lag_time * (1 - _mask) + 1000 * _mask \n\n    parsed['lag_time'] = lag_time\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    # extra information\n    'n_prev_seen': tf.io.RaggedFeature(value_key='n_prev_seen', dtype=tf.int64),\n    'n_prev_correctness': tf.io.RaggedFeature(value_key='n_prev_correctness', dtype=tf.int64),\n    # -------------------------------------------------------------------\n    # validation information\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\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    # This makes `parsed['shifted_answered_correctly']` has 1 more extra element at this moment.\n    parsed['shifted_answered_correctly'] = tf.concat([[START_TOKEN], parsed['answered_correctly']], axis=0)           \n\n    # We need to use `START_TOKEN` rather than `PAD_TOKEN` here.\n    # This makes `parsed['shifted_user_answer']` has 1 more extra element at this moment.\n    parsed['shifted_user_answer'] = tf.concat([[START_TOKEN], parsed['user_answer']], 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    # ----------------------------------------------------------------------------------------------------\n    # To get the lag time\n    # Block copied from `parse_train_example`.\n\n    diff_t = parsed['timestamp'] - tf.concat([[1], parsed['timestamp'][:-1]], axis=0)\n    bundle_start_mask = tf.cast(diff_t != 0, dtype=tf.int32)\n    # The jump values at bundle starting places\n    jump_t = tf.boolean_mask(diff_t, bundle_start_mask > 0)\n    # The task_container_id at bundle starting places\n    task_container_id_at_jumps = tf.boolean_mask(parsed['task_container_id'], bundle_start_mask > 0)\n    # Maps task_container_id to jump values\n    task_containder_id_to_jump_t = tf.scatter_nd(\n        indices=task_container_id_at_jumps[:, tf.newaxis], updates=jump_t, shape=tf.math.reduce_max(task_container_id_at_jumps + 1)[tf.newaxis], name=None\n    )\n    # `Our own version` of lag time.\n    lag_time = tf.gather(params=task_containder_id_to_jump_t , indices=parsed['task_container_id'])\n    # Make sure >= `0`\n    lag_time = tf.math.maximum(tf.cast(0, dtype=tf.int64), lag_time)\n\n    # set to 1 second if `lag_time == 0` but they are not in the same bunlde \n    _mask = tf.cast(tf.math.logical_and(lag_time == 0, parsed['timestamp'] > 0), dtype=tf.int64) \n    lag_time = lag_time * (1 - _mask) + 1000 * _mask \n\n    parsed['lag_time'] = lag_time\n    \n    return parsed","execution_count":null,"outputs":[]},{"metadata":{"id":"I2X60psHfLk8","outputId":"7d94592f-322a-49c4-f515-5b399f63c8a0","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    train_raw_ds = tf.data.TFRecordDataset(train_tfrec_paths, num_parallel_reads=1)\n    train_raw_ds = train_raw_ds.map(parse_train_example, num_parallel_calls=1, deterministic=True)\n    for x in train_raw_ds.take(1):\n        print(x)","execution_count":null,"outputs":[]},{"metadata":{"id":"E1MVqN03a7Bf"},"cell_type":"markdown","source":"#### check"},{"metadata":{"id":"SIG0fVP-LbRc","outputId":"fec026dc-facf-46f3-c91e-c9b5c87ab791","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    valid_raw_ds = tf.data.TFRecordDataset(valid_tfrec_paths, num_parallel_reads=1)\n    valid_raw_ds = valid_raw_ds.map(parse_train_example_with_valid_info, num_parallel_calls=1, deterministic=True)\n    for x in valid_raw_ds.take(1):\n        # pass\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 split the training and validation datasets."},{"metadata":{"id":"MBvizSTo11ql","trusted":true},"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\nWe sample random subsequences of user interactions (in the splitted training dataset) for training."},{"metadata":{"id":"WrZZVg2G11q9","trusted":true},"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 \n        - either a training batch (containing no validation part anymore) consisting of `tf.RaggedTensor` objects.\n        - or a validatioin batch (should be the ending part only)\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            'n_valid_blocks', 'valid_blocks_start_pos', 'valid_blocks_end_pos', 'valid_block_pos', 'valid_block_idx', 'valid_start', 'valid_end'\n        ]:\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            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":{"id":"o86M4bJy11rR"},"cell_type":"markdown","source":"### Tables mapping question / lectures / tags ids (etc.) to encoder / decoder input ids"},{"metadata":{"id":"L--Hbi4P11rm"},"cell_type":"markdown","source":"### Prepare encoder / decoder input ids and attention masks"},{"metadata":{"id":"aLVrl0Zy11rr"},"cell_type":"markdown","source":"#### special masks"},{"metadata":{"id":"u-JlqvWq11rv","trusted":true},"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":{"id":"qLweq7ry11sA","outputId":"b08be98c-efbe-4e0d-94d2-bc24f9cc9ac6","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    print(get_causal_attention_mask(3, 3, tf.int32, only_before=tf.constant(False)))\n    print(get_causal_attention_mask(3, 3, tf.int32, only_before=tf.constant(True)))\n\n    timestamp_tensors = tf.constant([[0, 0, 1, 2, 2, 2, 3, 3], [0, 1, 1, 2, 2, 3, 3, 3]], dtype=tf.int64)\n\n    print(get_attention_mask_from_timestamp_batch(timestamp_tensors, tf.int32, only_before=tf.constant(False)))\n    print(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)"},{"metadata":{"id":"NfG8YPPS11sU","trusted":true},"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 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    user_answer = batch['user_answer']\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    # Here, we use global variables, which is not good. But this is for fast experiments!!!\n    # Shifted by one, `PAD_TOKEN` --> `0`\n    _batch['abs_pos_ids'] = tf.cast(tf.math.minimum(tf.math.maximum(0, _batch['abs_pos'] + 1), MAX_HISTORY_LEN), dtype=DTYPE) / tf.cast(MAX_HISTORY_LEN, dtype=DTYPE) * WINDOW_SIZE\n    # Shifted by one, `PAD_TOKEN` --> `0`, `START_TOKEN` --> `0`.\n    _batch['shifted_abs_pos_ids'] = tf.cast(tf.math.minimum(tf.math.maximum(0, _batch['shifted_abs_pos'] + 1), MAX_HISTORY_LEN - 1), dtype=DTYPE) / tf.cast(MAX_HISTORY_LEN, dtype=DTYPE) * WINDOW_SIZE\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    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    ### Fixed\n    #decoder_input_ids = shifted_pad_mask * PAD_ID + shifted_start_mask * START_ID + shifted_response_lecture_mask * RESPONSE_LECTURE_ID + shifted_response_false_mask * RESPONSE_FALSE_ID + shifted_response_true_mask * RESPONSE_TRUE_ID\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    # ----------------------------------------\n    # d_ans_input_ids\n\n    shifted_user_answer = batch['shifted_user_answer']\n    shifted_answer_0_mask = tf.cast((shifted_user_answer == ANSWER_0_TOKEN), dtype=tf.int32)\n    shifted_answer_1_mask = tf.cast((shifted_user_answer == ANSWER_1_TOKEN), dtype=tf.int32)\n    shifted_answer_2_mask = tf.cast((shifted_user_answer == ANSWER_2_TOKEN), dtype=tf.int32)\n    shifted_answer_3_mask = tf.cast((shifted_user_answer == ANSWER_3_TOKEN), dtype=tf.int32)\n\n    decoder_answer_input_ids = shifted_pad_mask * PAD_ID + shifted_start_mask * START_ID + shifted_masking_mask * MASK_ID + \\\n        shifted_response_lecture_mask * ANSWER_LECTURE_ID + \\\n        shifted_answer_0_mask * ANSWER_0_ID + shifted_answer_1_mask * ANSWER_1_ID + \\\n        shifted_answer_2_mask * ANSWER_2_ID + shifted_answer_3_mask * ANSWER_3_ID\n\n    # ----------------------------------------\n    # post processing of decoder inputs\n\n    if not generative:\n\n        # This either contains exactly one place with `1`, or all places are `0`.\n        pred_time_question_start_mask = tf.cast(tf.math.cumsum(pred_time_question_mask, axis=1) == 1, dtype=tf.int32)\n        \n        # shape = [batch_size]\n        pred_time_question_start_value_response = tf.math.reduce_sum(pred_time_question_start_mask * decoder_input_ids, axis=1)\n        pred_time_question_start_value_answer = tf.math.reduce_sum(pred_time_question_start_mask * decoder_answer_input_ids, axis=1)\n\n        # If at the starting of questions, we get `MASK_ID`, we change it to `RESPONSE_LECTURE_ID` / `ANSWER_LECTURE_ID`\n        prev_lecture_mask = tf.cast(pred_time_question_start_value_response == MASK_ID, tf.int32)            \n        \n        pred_time_question_start_value_response = RESPONSE_LECTURE_ID * prev_lecture_mask + pred_time_question_start_value_response * (1 - prev_lecture_mask)\n        pred_time_question_start_value_answer = ANSWER_LECTURE_ID * prev_lecture_mask + pred_time_question_start_value_answer * (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_response[:, tf.newaxis]\n        decoder_answer_input_ids = decoder_answer_input_ids * (1 - pred_time_question_mask) + pred_time_question_mask * pred_time_question_start_value_answer[:, 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        decoder_answer_input_ids = decoder_answer_input_ids * tf.cast(decoder_answer_input_ids != MASK_ID, dtype=tf.int32) + ANSWER_LECTURE_ID * tf.cast(decoder_answer_input_ids == MASK_ID, dtype=tf.int32)\n\n   # ----------------------------------------\n   # add to `_batch`\n\n    _batch['d_input_ids'] = decoder_input_ids\n    _batch['d_ans_input_ids'] = decoder_answer_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        abs_pos_ids = _batch['abs_pos_ids']\n        shifted_abs_pos_ids = _batch['shifted_abs_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        abs_pos_ids_question_start_value = tf.math.reduce_sum(tf.cast(pred_time_question_start_mask, dtype=DTYPE) * abs_pos_ids, axis=1)\n        shifted_abs_pos_ids_question_start_value = tf.math.reduce_sum(tf.cast(pred_time_question_start_mask, dtype=DTYPE) * shifted_abs_pos_ids, axis=1)\n\n        abs_pos_ids = abs_pos_ids * tf.cast(1 - pred_time_question_mask, dtype=DTYPE) + tf.cast(pred_time_question_mask, dtype=DTYPE) * abs_pos_ids_question_start_value[:, tf.newaxis]\n        shifted_abs_pos_ids = shifted_abs_pos_ids * tf.cast(1 - pred_time_question_mask, dtype=DTYPE) + tf.cast(pred_time_question_mask, dtype=DTYPE) * shifted_abs_pos_ids_question_start_value[:, tf.newaxis]\n\n        _batch['pos_ids'] = pos_ids\n        _batch['shifted_pos_ids'] = shifted_pos_ids\n        _batch['abs_pos_ids'] = abs_pos_ids\n        _batch['shifted_abs_pos_ids'] = shifted_abs_pos_ids        \n\n    # ----------------------------------------\n    # tags\n\n    # The original tags are added by `1`\n    _batch['tag_ids'] = tf.gather(params=c_inputs_ids_to_tags, indices=_batch['c_input_ids']) + 1 \n        \n    # ----------------------------------------        \n    # part\n\n    # The original parts are added by `1`\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    # scaled to [0, 1]\n\n    _batch['prior_question_elapsed_time_input'] = tf.cast(_batch['prior_question_elapsed_time'], dtype=DTYPE) / 1000.0 / 300.0\n\n    # ----------------------------------------\n    # normalized `lag_time`\n\n    lag_time = tf.cast(_batch['lag_time'], dtype=DTYPE)\n    # nb. of hours\n    lag_time = lag_time / 1000.0 / 3600.0\n    # If `lag_time` > `72 hours` --> set it to `72 hours`.\n    lag_time = tf.math.minimum(lag_time, 72.0)\n\n    _batch['lag_time'] = lag_time\n\n    # ----------------------------------------\n    # use `task_container_ids` as positional information\n\n    task_container_pos_ids = tf.cast(_batch['task_container_id'], dtype=DTYPE) / MAX_TASK_CONTAINER_ID * 10.0\n    _batch['task_container_pos_ids'] = task_container_pos_ids\n\n    # ----------------------------------------\n    # answer correctness 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    # answer target\n        \n    # negated values become `NON_TARGET_ID (-100)`\n    _batch['answer_target'] = batch['user_answer'] * answer_mask + (NON_TARGET_ID) * (1 - answer_mask)\n    \n    # ----------------------------------------\n    # correct_answer_id  \n        \n    # Unlike `tag_ids` or `part_ids`, we don't need to have `+ 1` because `c_inputs_ids_to_correct_answer_id` is built in a slightly different way.\n    _batch['correct_answer_id'] = tf.gather(params=c_inputs_ids_to_correct_answer_id, indices=_batch['c_input_ids'])\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        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    # ----------------------------------------------------------------------------------------------------\n    # To process aggregated historical information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    scaling_factor = tf.constant(N_AGGREGATED_QUESTION_SCALING_FACTOR, dtype=tf.float32)\n    \n    _batch['n_questions_answered_scaled'] = tf.cast(_batch['n_questions_answered'], dtype=tf.float32) / scaling_factor\n    _batch['n_lectures_watched_scaled'] = tf.cast(_batch['n_lectures_watched'], dtype=tf.float32) / scaling_factor    \n    \n    # set minimum to `1` to avoid division by `0` error.\n    n_questions_answered = tf.cast(tf.math.maximum(_batch['n_questions_answered'], 1), dtype=tf.float32)\n            \n    _batch['answered_correctly_ratio'] = tf.cast(_batch['n_questions_answered_correctly'], dtype=tf.float32) / n_questions_answered\n    \n    # ----------------------------------------------------------------------------------------------------    \n    # To process aggregated historical part information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    part_count_scaled = []\n    part_correctness_ratio = []\n    \n    for part_idx in range(2, 9):\n        \n        key = f'part_{part_idx}_count'\n        _part_count = tf.cast(_batch[key], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _part_count_scaled = _part_count / scaling_factor\n        \n        part_count_scaled.append(_part_count_scaled)\n        \n        key_2 = f'part_{part_idx}_correctness_count'\n        _part_correctness_count = tf.cast(_batch[key_2], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _part_correctness_ratio = _part_correctness_count / tf.math.maximum(_part_count, 1.0)\n        \n        part_correctness_ratio.append(_part_correctness_ratio)\n        \n        # We don't want to keep these keys\n        del _batch[key]\n        del _batch[key_2]\n        \n    part_count_scaled = tf.stack(part_count_scaled, axis=-1)\n    part_correctness_ratio = tf.stack(part_correctness_ratio, axis=-1)\n    \n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE - 2]\n    _batch['part_count_scaled'] = part_count_scaled\n    _batch['part_correctness_ratio'] = part_correctness_ratio\n\n    # ----------------------------------------------------------------------------------------------------  \n    # To process `current` aggregated historical part information (i.e. not per part level)\n\n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE]\n    _part_count_scaled = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            part_count_scaled\n        ],\n        axis=-1\n    )\n\n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE]\n    _part_correctness_ratio = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            part_correctness_ratio\n        ],\n        axis=-1\n    )\n\n    current_part_count_scaled = tf.gather(params=_part_count_scaled, indices=_batch['part_ids'], batch_dims=2)\n    current_part_correctness_ratio = tf.gather(params=_part_correctness_ratio, indices=_batch['part_ids'], batch_dims=2)\n\n    # shape = [batch_size, seq_len]\n    _batch['current_part_count_scaled'] = current_part_count_scaled\n    _batch['current_part_correctness_ratio'] = current_part_correctness_ratio\n    \n    # ----------------------------------------------------------------------------------------------------      \n    # To process aggregated historical correct answer information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    correct_answer_count_scaled = []\n    correct_answer_correctness_ratio = []\n    \n    for correct_answer_idx in range(ANSWER_0_ID, ANSWER_3_ID + 1):\n        \n        key = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_count'\n        _correct_answer_count = tf.cast(_batch[key], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _correct_answer_count_scaled = _correct_answer_count / scaling_factor\n        \n        correct_answer_count_scaled.append(_correct_answer_count_scaled)\n        \n        key_2 = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_answered_correctly_count'\n        _correct_answer_answered_correctly_count = tf.cast(_batch[key_2], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _correct_answer_answered_correctly_ratio = _correct_answer_answered_correctly_count / tf.math.maximum(_correct_answer_count, 1.0)\n        \n        correct_answer_correctness_ratio.append(_correct_answer_answered_correctly_ratio)\n        \n        # We don't want to keep these keys\n        del _batch[key]\n        del _batch[key_2]\n        \n    correct_answer_count_scaled = tf.stack(correct_answer_count_scaled, axis=-1)\n    correct_answer_correctness_ratio = tf.stack(correct_answer_correctness_ratio, axis=-1)\n    \n    # shape = [batch_size, seq_len, 7]\n    _batch['correct_answer_count_scaled'] = correct_answer_count_scaled\n    _batch['correct_answer_correctness_ratio'] = correct_answer_correctness_ratio\n    \n    # ----------------------------------------------------------------------------------------------------     \n    # To process `current` aggregated historical correct answer information (i.e. not per correct answer level)\n\n    # shape = [batch_size, seq_len, ANSWER_VOCAB_SIZE - 3]\n    _correct_answer_count_scaled = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,  # For lectures\n            correct_answer_count_scaled\n        ],\n        axis=-1\n    )\n\n    # shape = [batch_size, seq_len, ANSWER_VOCAB_SIZE - 3]\n    _correct_answer_correctness_ratio = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],  # For lectures\n            correct_answer_correctness_ratio\n        ],\n        axis=-1\n    )\n\n    # `PAD_ID` gives index `0` and `ANSWER_LECTURE_ID` gives index `1`.\n    # `ANSWER_START_ID`, `ANSWER_END_ID` and `ANSWER_MASK_ID` doesn't exist in `correct_answer_id`.\n    _indices = tf.math.maximum(_batch['correct_answer_id'] - ANSWER_LECTURE_ID + 1, 0)\n\n    current_correct_answer_count_scaled = tf.gather(params=_correct_answer_count_scaled, indices=_indices, batch_dims=2)\n    current_correct_answer_correctness_ratio = tf.gather(params=_correct_answer_correctness_ratio, indices=_indices, batch_dims=2)\n\n    # shape = [batch_size, seq_len]\n    _batch['current_correct_answer_count_scaled'] = current_correct_answer_count_scaled\n    _batch['current_correct_answer_correctness_ratio'] = current_correct_answer_correctness_ratio    \n\n    # ----------------------------------------------------------------------------------------------------\n        \n    _batch['current_question_count_scaled'] = tf.cast(_batch['n_prev_seen'], dtype=tf.float32) / scaling_factor\n\n    # set minimum to `1` to avoid division by `0` error.\n    current_question_count = tf.cast(tf.math.maximum(_batch['n_prev_seen'], 1), dtype=tf.float32)\n            \n    _batch['current_question_correctness_ratio'] = tf.cast(_batch['n_prev_correctness'], dtype=tf.float32) / current_question_count\n\n    # ----------------------------------------------------------------------------------------------------\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":{"trusted":true},"cell_type":"code","source":"def get_input_signatures_2():\n    \n    batch_size = None\n    seq_len = 128\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        'shifted_user_answer': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='shifted_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_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        'lag_time': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int64, name='lag_time'),\n        'n_questions_answered': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_questions_answered'),  \n        'n_lectures_watched': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_lectures_watched'),  \n        'n_questions_answered_correctly': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_questions_answered_correctly'),  \n        'n_prev_seen': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_prev_seen'),\n        'n_prev_correctness': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_prev_correctness'),\n        'part_2_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_2_count'),\n        'part_3_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_3_count'),\n        'part_4_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_4_count'),\n        'part_5_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_5_count'),\n        'part_6_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_6_count'),\n        'part_7_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_7_count'),\n        'part_8_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_8_count'),\n        'part_2_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_2_correctness_count'),\n        'part_3_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_3_correctness_count'),\n        'part_4_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_4_correctness_count'),\n        'part_5_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_5_correctness_count'),\n        'part_6_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_6_correctness_count'),\n        'part_7_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_7_correctness_count'),\n        'part_8_correctness_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='part_8_correctness_count'),\n        'correct_answer_0_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_0_count'),\n        'correct_answer_1_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_1_count'),\n        'correct_answer_2_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_2_count'),\n        'correct_answer_3_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_3_count'),\n        'correct_answer_0_answered_correctly_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_0_answered_correctly_count'),\n        'correct_answer_1_answered_correctly_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_1_answered_correctly_count'),\n        'correct_answer_2_answered_correctly_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_2_answered_correctly_count'),\n        'correct_answer_3_answered_correctly_count': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_3_answered_correctly_count'),\n    }\n \n    return [input_signatures]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function(\n    input_signature=get_input_signatures_2()\n)\ndef add_input_ids_and_targets_2(batch):\n    \"\"\"Add input ids and targets for training\n    \"\"\"\n        \n    training = tf.constant(0, dtype=tf.int32)\n    generative = tf.constant(False, dtype=tf.bool)\n   \n    content_ids = batch['content_id']\n    content_type_ids = batch['content_type_id']\n    answered_correctly = batch['answered_correctly']\n    user_answer = batch['user_answer']\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    # Here, we use global variables, which is not good. But this is for fast experiments!!!\n    # Shifted by one, `PAD_TOKEN` --> `0`\n    _batch['abs_pos_ids'] = tf.cast(tf.math.minimum(tf.math.maximum(0, _batch['abs_pos'] + 1), MAX_HISTORY_LEN), dtype=DTYPE) / tf.cast(MAX_HISTORY_LEN, dtype=DTYPE) * WINDOW_SIZE\n    # Shifted by one, `PAD_TOKEN` --> `0`, `START_TOKEN` --> `0`.\n    _batch['shifted_abs_pos_ids'] = tf.cast(tf.math.minimum(tf.math.maximum(0, _batch['shifted_abs_pos'] + 1), MAX_HISTORY_LEN - 1), dtype=DTYPE) / tf.cast(MAX_HISTORY_LEN, dtype=DTYPE) * WINDOW_SIZE\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    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    ### Fixed\n    #decoder_input_ids = shifted_pad_mask * PAD_ID + shifted_start_mask * START_ID + shifted_response_lecture_mask * RESPONSE_LECTURE_ID + shifted_response_false_mask * RESPONSE_FALSE_ID + shifted_response_true_mask * RESPONSE_TRUE_ID\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    # ----------------------------------------\n    # d_ans_input_ids\n\n    shifted_user_answer = batch['shifted_user_answer']\n    shifted_answer_0_mask = tf.cast((shifted_user_answer == ANSWER_0_TOKEN), dtype=tf.int32)\n    shifted_answer_1_mask = tf.cast((shifted_user_answer == ANSWER_1_TOKEN), dtype=tf.int32)\n    shifted_answer_2_mask = tf.cast((shifted_user_answer == ANSWER_2_TOKEN), dtype=tf.int32)\n    shifted_answer_3_mask = tf.cast((shifted_user_answer == ANSWER_3_TOKEN), dtype=tf.int32)\n\n    decoder_answer_input_ids = shifted_pad_mask * PAD_ID + shifted_start_mask * START_ID + shifted_masking_mask * MASK_ID + \\\n        shifted_response_lecture_mask * ANSWER_LECTURE_ID + \\\n        shifted_answer_0_mask * ANSWER_0_ID + shifted_answer_1_mask * ANSWER_1_ID + \\\n        shifted_answer_2_mask * ANSWER_2_ID + shifted_answer_3_mask * ANSWER_3_ID\n\n    # ----------------------------------------\n    # post processing of decoder inputs\n\n    ### if not generative:\n\n    # This either contains exactly one place with `1`, or all places are `0`.\n    pred_time_question_start_mask = tf.cast(tf.math.cumsum(pred_time_question_mask, axis=1) == 1, dtype=tf.int32)\n\n    # shape = [batch_size]\n    pred_time_question_start_value_response = tf.math.reduce_sum(pred_time_question_start_mask * decoder_input_ids, axis=1)\n    pred_time_question_start_value_answer = tf.math.reduce_sum(pred_time_question_start_mask * decoder_answer_input_ids, axis=1)\n\n    # If at the starting of questions, we get `MASK_ID`, we change it to `RESPONSE_LECTURE_ID` / `ANSWER_LECTURE_ID`\n    prev_lecture_mask = tf.cast(pred_time_question_start_value_response == MASK_ID, tf.int32)            \n\n    pred_time_question_start_value_response = RESPONSE_LECTURE_ID * prev_lecture_mask + pred_time_question_start_value_response * (1 - prev_lecture_mask)\n    pred_time_question_start_value_answer = ANSWER_LECTURE_ID * prev_lecture_mask + pred_time_question_start_value_answer * (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_response[:, tf.newaxis]\n    decoder_answer_input_ids = decoder_answer_input_ids * (1 - pred_time_question_mask) + pred_time_question_mask * pred_time_question_start_value_answer[:, 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    decoder_answer_input_ids = decoder_answer_input_ids * tf.cast(decoder_answer_input_ids != MASK_ID, dtype=tf.int32) + ANSWER_LECTURE_ID * tf.cast(decoder_answer_input_ids == MASK_ID, dtype=tf.int32)\n\n   # ----------------------------------------\n   # add to `_batch`\n\n    _batch['d_input_ids'] = decoder_input_ids\n    _batch['d_ans_input_ids'] = decoder_answer_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    abs_pos_ids = _batch['abs_pos_ids']\n    shifted_abs_pos_ids = _batch['shifted_abs_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    abs_pos_ids_question_start_value = tf.math.reduce_sum(tf.cast(pred_time_question_start_mask, dtype=DTYPE) * abs_pos_ids, axis=1)\n    shifted_abs_pos_ids_question_start_value = tf.math.reduce_sum(tf.cast(pred_time_question_start_mask, dtype=DTYPE) * shifted_abs_pos_ids, axis=1)\n\n    abs_pos_ids = abs_pos_ids * tf.cast(1 - pred_time_question_mask, dtype=DTYPE) + tf.cast(pred_time_question_mask, dtype=DTYPE) * abs_pos_ids_question_start_value[:, tf.newaxis]\n    shifted_abs_pos_ids = shifted_abs_pos_ids * tf.cast(1 - pred_time_question_mask, dtype=DTYPE) + tf.cast(pred_time_question_mask, dtype=DTYPE) * shifted_abs_pos_ids_question_start_value[:, tf.newaxis]\n\n    _batch['pos_ids'] = pos_ids\n    _batch['shifted_pos_ids'] = shifted_pos_ids\n    _batch['abs_pos_ids'] = abs_pos_ids\n    _batch['shifted_abs_pos_ids'] = shifted_abs_pos_ids        \n\n    # ----------------------------------------\n    # tags\n\n    # The original tags are added by `1`\n    _batch['tag_ids'] = tf.gather(params=c_inputs_ids_to_tags, indices=_batch['c_input_ids']) + 1 \n        \n    # ----------------------------------------        \n    # part\n\n    # The original parts are added by `1`\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    # scaled to [0, 1]\n\n    _batch['prior_question_elapsed_time_input'] = tf.cast(_batch['prior_question_elapsed_time'], dtype=DTYPE) / 1000.0 / 300.0\n\n    # ----------------------------------------\n    # normalized `lag_time`\n\n    lag_time = tf.cast(_batch['lag_time'], dtype=DTYPE)\n    # nb. of hours\n    lag_time = lag_time / 1000.0 / 3600.0\n    # If `lag_time` > `72 hours` --> set it to `72 hours`.\n    lag_time = tf.math.minimum(lag_time, 72.0)\n\n    _batch['lag_time'] = lag_time\n\n    # ----------------------------------------\n    # use `task_container_ids` as positional information\n\n    task_container_pos_ids = tf.cast(_batch['task_container_id'], dtype=DTYPE) / MAX_TASK_CONTAINER_ID * 10.0\n    _batch['task_container_pos_ids'] = task_container_pos_ids\n\n    # ----------------------------------------\n    # answer correctness 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    # answer target\n        \n    # negated values become `NON_TARGET_ID (-100)`\n    _batch['answer_target'] = batch['user_answer'] * answer_mask + (NON_TARGET_ID) * (1 - answer_mask)\n    \n    # ----------------------------------------\n    # correct_answer_id  \n        \n    # Unlike `tag_ids` or `part_ids`, we don't need to have `+ 1` because `c_inputs_ids_to_correct_answer_id` is built in a slightly different way.\n    _batch['correct_answer_id'] = tf.gather(params=c_inputs_ids_to_correct_answer_id, indices=_batch['c_input_ids'])\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        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    # ----------------------------------------------------------------------------------------------------\n    # To process aggregated historical information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    scaling_factor = tf.constant(N_AGGREGATED_QUESTION_SCALING_FACTOR, dtype=tf.float32)\n    \n    _batch['n_questions_answered_scaled'] = tf.cast(_batch['n_questions_answered'], dtype=tf.float32) / scaling_factor\n    _batch['n_lectures_watched_scaled'] = tf.cast(_batch['n_lectures_watched'], dtype=tf.float32) / scaling_factor    \n    \n    # set minimum to `1` to avoid division by `0` error.\n    n_questions_answered = tf.cast(tf.math.maximum(_batch['n_questions_answered'], 1), dtype=tf.float32)\n            \n    _batch['answered_correctly_ratio'] = tf.cast(_batch['n_questions_answered_correctly'], dtype=tf.float32) / n_questions_answered\n    \n    # ----------------------------------------------------------------------------------------------------    \n    # To process aggregated historical part information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    part_count_scaled = []\n    part_correctness_ratio = []\n    \n    for part_idx in range(2, 9):\n        \n        key = f'part_{part_idx}_count'\n        _part_count = tf.cast(_batch[key], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _part_count_scaled = _part_count / scaling_factor\n        \n        part_count_scaled.append(_part_count_scaled)\n        \n        key_2 = f'part_{part_idx}_correctness_count'\n        _part_correctness_count = tf.cast(_batch[key_2], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _part_correctness_ratio = _part_correctness_count / tf.math.maximum(_part_count, 1.0)\n        \n        part_correctness_ratio.append(_part_correctness_ratio)\n        \n        # We don't want to keep these keys\n        del _batch[key]\n        del _batch[key_2]\n        \n    part_count_scaled = tf.stack(part_count_scaled, axis=-1)\n    part_correctness_ratio = tf.stack(part_correctness_ratio, axis=-1)\n    \n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE - 2]\n    _batch['part_count_scaled'] = part_count_scaled\n    _batch['part_correctness_ratio'] = part_correctness_ratio\n\n    # ----------------------------------------------------------------------------------------------------  \n    # To process `current` aggregated historical part information (i.e. not per part level)\n\n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE]\n    _part_count_scaled = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            part_count_scaled\n        ],\n        axis=-1\n    )\n\n    # shape = [batch_size, seq_len, PART_VOCAB_SIZE]\n    _part_correctness_ratio = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            part_correctness_ratio\n        ],\n        axis=-1\n    )\n\n    current_part_count_scaled = tf.gather(params=_part_count_scaled, indices=_batch['part_ids'], batch_dims=2)\n    current_part_correctness_ratio = tf.gather(params=_part_correctness_ratio, indices=_batch['part_ids'], batch_dims=2)\n\n    # shape = [batch_size, seq_len]\n    _batch['current_part_count_scaled'] = current_part_count_scaled\n    _batch['current_part_correctness_ratio'] = current_part_correctness_ratio\n    \n    # ----------------------------------------------------------------------------------------------------      \n    # To process aggregated historical correct answer information\n    # Be careful, the aggreated information doesn't contain the current place.\n    # shape = [batch_size, seq_len]\n\n    correct_answer_count_scaled = []\n    correct_answer_correctness_ratio = []\n    \n    for correct_answer_idx in range(ANSWER_0_ID, ANSWER_3_ID + 1):\n        \n        key = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_count'\n        _correct_answer_count = tf.cast(_batch[key], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _correct_answer_count_scaled = _correct_answer_count / scaling_factor\n        \n        correct_answer_count_scaled.append(_correct_answer_count_scaled)\n        \n        key_2 = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_answered_correctly_count'\n        _correct_answer_answered_correctly_count = tf.cast(_batch[key_2], dtype=tf.float32)\n        # shape = [batch_size, seq_len]\n        _correct_answer_answered_correctly_ratio = _correct_answer_answered_correctly_count / tf.math.maximum(_correct_answer_count, 1.0)\n        \n        correct_answer_correctness_ratio.append(_correct_answer_answered_correctly_ratio)\n        \n        # We don't want to keep these keys\n        del _batch[key]\n        del _batch[key_2]\n        \n    correct_answer_count_scaled = tf.stack(correct_answer_count_scaled, axis=-1)\n    correct_answer_correctness_ratio = tf.stack(correct_answer_correctness_ratio, axis=-1)\n    \n    # shape = [batch_size, seq_len, 7]\n    _batch['correct_answer_count_scaled'] = correct_answer_count_scaled\n    _batch['correct_answer_correctness_ratio'] = correct_answer_correctness_ratio\n    \n    # ----------------------------------------------------------------------------------------------------     \n    # To process `current` aggregated historical correct answer information (i.e. not per correct answer level)\n\n    # shape = [batch_size, seq_len, ANSWER_VOCAB_SIZE - 3]\n    _correct_answer_count_scaled = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis] / scaling_factor,  # For lectures\n            correct_answer_count_scaled\n        ],\n        axis=-1\n    )\n\n    # shape = [batch_size, seq_len, ANSWER_VOCAB_SIZE - 3]\n    _correct_answer_correctness_ratio = tf.concat(\n        [\n            tf.constant(PAD_TOKEN, dtype=tf.float32) * tf.ones_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],\n            tf.zeros_like(_batch['content_id'], dtype=tf.float32)[:, :, tf.newaxis],  # For lectures\n            correct_answer_correctness_ratio\n        ],\n        axis=-1\n    )\n\n    # `PAD_ID` gives index `0` and `ANSWER_LECTURE_ID` gives index `1`.\n    # `ANSWER_START_ID`, `ANSWER_END_ID` and `ANSWER_MASK_ID` doesn't exist in `correct_answer_id`.\n    _indices = tf.math.maximum(_batch['correct_answer_id'] - ANSWER_LECTURE_ID + 1, 0)\n\n    current_correct_answer_count_scaled = tf.gather(params=_correct_answer_count_scaled, indices=_indices, batch_dims=2)\n    current_correct_answer_correctness_ratio = tf.gather(params=_correct_answer_correctness_ratio, indices=_indices, batch_dims=2)\n\n    # shape = [batch_size, seq_len]\n    _batch['current_correct_answer_count_scaled'] = current_correct_answer_count_scaled\n    _batch['current_correct_answer_correctness_ratio'] = current_correct_answer_correctness_ratio    \n\n    # ----------------------------------------------------------------------------------------------------\n        \n    _batch['current_question_count_scaled'] = tf.cast(_batch['n_prev_seen'], dtype=tf.float32) / scaling_factor\n\n    # set minimum to `1` to avoid division by `0` error.\n    current_question_count = tf.cast(tf.math.maximum(_batch['n_prev_seen'], 1), dtype=tf.float32)\n            \n    _batch['current_question_correctness_ratio'] = tf.cast(_batch['n_prev_correctness'], dtype=tf.float32) / current_question_count\n\n    # ----------------------------------------------------------------------------------------------------\n    \n    return _batch","execution_count":null,"outputs":[]},{"metadata":{"id":"0IM9EODsUpZA"},"cell_type":"markdown","source":""},{"metadata":{"id":"12dqrCj-UqNJ","trusted":true},"cell_type":"code","source":" def add_aggregated_historical_information(parsed):\n    \"\"\"\n    For validation dataset, this must be performed after `remove_future_valid_blocks` is applied (so we have the actual `pred_time_mask`).\n    For tranining dataset, the `pred_time_mask` is (and should be) all `0`, so every places are considered.\n    \"\"\"\n \n    # ----------------------------------------------------------------------------------------------------\n    # To get aggregated historical information\n    # Be careful, the aggreated information shouldn't contain the current place.\n    \n    question_mask = tf.cast(parsed['content_type_id'] == 0, dtype=tf.int32)\n    lecture_mask = tf.cast(parsed['content_type_id'] == 1, dtype=tf.int32)\n    correction_mask = tf.cast(parsed['answered_correctly'] == 1, dtype=tf.int32)\n    content_ids = parsed['content_id']\n    non_pred_time_mask = 1 - parsed['pred_time_mask']\n    non_pred_time_question_mask = question_mask * non_pred_time_mask\n    non_pred_time_question_correction_mask = non_pred_time_question_mask * correction_mask\n    \n    # shape = [seq_len]\n    parsed['n_questions_answered'] = tf.concat(\n        [\n            tf.constant([0], dtype=tf.int32),\n            tf.math.cumsum(non_pred_time_question_mask[:-1], axis=-1)\n        ],\n        axis=-1\n    )\n    \n    parsed['n_questions_answered_correctly'] = tf.concat(\n        [\n            tf.constant([0], dtype=tf.int32),\n            tf.math.cumsum(non_pred_time_question_correction_mask[:-1], axis=-1)\n        ],\n        axis=-1\n    )\n    \n    parsed['n_lectures_watched'] = tf.concat(\n        [\n            tf.constant([0], dtype=tf.int32),\n            tf.math.cumsum(lecture_mask[:-1], axis=-1)\n        ],\n        axis=-1\n    )\n    \n    # ----------------------------------------------------------------------------------------------------\n    # Aggregated historical information for part\n    \n    content_input_ids = question_mask * question_id_to_input_id_table.lookup(content_ids) + lecture_mask * lecture_id_to_input_id_table.lookup(content_ids)\n        \n    # The original parts are added by `1`\n    part_ids = tf.gather(params=c_inputs_ids_to_part, indices=content_input_ids) + 1\n    \n    for part_idx in range(2, 9):\n        \n        key = f'part_{part_idx}_count'\n        \n        part_mask = tf.cast(part_ids == part_idx, dtype=tf.int32)\n        # Only count for questions that are not in `pred_time`.\n        part_mask = part_mask * non_pred_time_question_mask\n        part_correct_mask = part_mask * non_pred_time_question_correction_mask\n        \n        parsed[key] = tf.concat(\n            [\n                tf.constant([0], dtype=tf.int32),\n                tf.math.cumsum(part_mask[:-1], axis=-1)\n            ],\n            axis=-1\n        )\n        \n        key = f'part_{part_idx}_correctness_count'\n        \n        parsed[key] = tf.concat(\n            [\n                tf.constant([0], dtype=tf.int32),\n                tf.math.cumsum(part_correct_mask[:-1], axis=-1)\n            ],\n            axis=-1\n        )  \n       \n    # ----------------------------------------------------------------------------------------------------\n    # Aggregated historical information for correct answer\n        \n    # There is an offset of `ANSWER_0_ID`.\n    correct_answer_id = tf.gather(params=c_inputs_ids_to_correct_answer_id, indices=content_input_ids)\n    \n    for correct_answer_idx in range(ANSWER_0_ID, ANSWER_3_ID + 1):\n        \n        key = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_count'\n        \n        correct_answer_mask = tf.cast(correct_answer_id == correct_answer_idx, dtype=tf.int32)\n        # Only count for questions that are not in `pred_time`.\n        correct_answer_mask = correct_answer_mask * non_pred_time_question_mask        \n\n        correct_answer_answered_correctly_mask = correct_answer_mask * non_pred_time_question_correction_mask\n        \n        parsed[key] = tf.concat(\n            [\n                tf.constant([0], dtype=tf.int32),\n                tf.math.cumsum(correct_answer_mask[:-1], axis=-1)\n            ],\n            axis=-1\n        )\n        \n        key = f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_answered_correctly_count'\n        \n        parsed[key] = tf.concat(\n            [\n                tf.constant([0], dtype=tf.int32),\n                tf.math.cumsum(correct_answer_answered_correctly_mask[:-1], axis=-1)\n            ],\n            axis=-1\n        )      \n\n    return parsed","execution_count":null,"outputs":[]},{"metadata":{"id":"V4YUlTldJ099"},"cell_type":"markdown","source":"#### For validation dataset"},{"metadata":{"id":"n13GQqGSJuf6","trusted":true},"cell_type":"code","source":"def add_valid_block_info(valid_raw_example, num_parallel_calls=None, deterministic=None):\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), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\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, num_parallel_calls=num_parallel_calls, deterministic=deterministic))\n    valid_ds = valid_ds.map(lambda example: remove_future_valid_blocks(example), num_parallel_calls=num_parallel_calls, deterministic=deterministic)\n    \n    # This must be performed after `remove_future_valid_blocks`\n    valid_ds = valid_ds.map(lambda example: add_aggregated_historical_information(example), num_parallel_calls=num_parallel_calls, deterministic=deterministic)   \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    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":"if not IS_KAGGLE:\n\n    train_raw_ds = tf.data.TFRecordDataset(train_tfrec_paths, num_parallel_reads=1)\n    train_raw_ds = train_raw_ds.map(parse_train_example, num_parallel_calls=1, deterministic=True)\n\n    # add aggregated historical information\n    train_raw_ds = train_raw_ds.map(lambda example: add_aggregated_historical_information(example))\n\n    # Get the splitted training dataset.\n    train_valid_split_indices = load_data(train_valid_split_indices_paths[0])\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=1, deterministic=True)\n\n    # batch examples with attributes having different lengths across examples - tf.RaggedTensor\n    batched_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\n    batched_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\n    train_ds = prepare_training_dataset(batched_ds, use_abs_pos=False, num_parallel_calls=1, deterministic=True)","execution_count":null,"outputs":[]},{"metadata":{"id":"PGQBbKGprMMs","outputId":"f9165d60-9f64-45fe-f369-80668e31c020","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    it = iter(train_ds.take(2))\n    next(it)","execution_count":null,"outputs":[]},{"metadata":{"id":"odrCZd7J132N","outputId":"a5290332-e5e8-4752-f46f-0d198b305968","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    next(it)['pred_time_mask']","execution_count":null,"outputs":[]},{"metadata":{"id":"GCO3h1YIL59c","outputId":"b1c0c7f8-7c1b-4e46-b970-5c3b8c35b0e9","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n\n    valid_raw_ds = tf.data.TFRecordDataset([valid_tfrec_paths], num_parallel_reads=1)\n    valid_raw_ds = valid_raw_ds.map(parse_train_example_with_valid_info, num_parallel_calls=1, deterministic=True)\n\n    for x in valid_raw_ds.take(1):\n        # print(x)\n        for k in x:\n            print(f'{k} : shape = {x[k].shape}')\n        print('--------')\n\n    print('============================================')\n\n    valid_ds = prepare_validation_dataset(valid_raw_ds, batch_size=5, window_size=10, use_abs_pos=False)\n\n    for x in valid_ds.take(1):\n        # print(x)\n        for k in x:\n            print(f'{k} : shape = {x[k].shape}')\n\n        print('--------')","execution_count":null,"outputs":[]},{"metadata":{"id":"zjRpC3h0wvyM","outputId":"a3e5688f-927a-4868-897b-ff632a92034c","trusted":true},"cell_type":"code","source":"# it = iter(valid_ds.take(2))\n# next(it)","execution_count":null,"outputs":[]},{"metadata":{"id":"6FMsKoD-11s_"},"cell_type":"markdown","source":"## Model definition"},{"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":{"id":"lbkDSH67xuNY","trusted":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=WINDOW_SIZE + 1,\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_user_answer=USE_USER_ANSWER,\n        use_user_answer_loss=USE_USER_ANSWER_LOSS,\n        use_correct_answer_for_encoder=USE_CORRECT_ANSWER_FOR_ENCODER,\n        use_correct_answer_for_decoder=USE_CORRECT_ANSWER_FOR_DECODER,\n        use_abs_pos=USE_ABS_POS,\n        use_task_container_pos=USE_TASK_CONTAINER_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        use_lag_time=USE_LAG_TIME,\n        use_lag_time_for_encoder=USE_LAG_TIME_FOR_ENCODER,\n        use_user_level_aggregated_historical_info=USE_USER_LEVEL_AGGREGATED_HISTORICAL_INFO,\n        use_part_aggregated_historical_info=USE_PART_AGGREGATED_HISTORICAL_INFO,\n        use_correct_answer_aggregated_historical_info=USE_CORRECT_ANSWER_AGGREGATED_HISTORICAL_INFO,\n        use_question_level_aggregated_historical_info=USE_QUESTION_LEVEL_AGGREGATED_HISTORICAL_INFO,\n        allow_bundle_atten=ALLOW_BUNDLE_ATTEN,\n        generative=GENERATIVE,\n        use_pre_classifier=USE_PRE_CLASSIFIER,\n        use_softmax=USE_SOFTMAX,\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_user_answer = use_user_answer\n        self.use_correct_answer_for_encoder = use_correct_answer_for_encoder\n        self.use_correct_answer_for_decoder = use_correct_answer_for_decoder\n        self.use_user_answer_loss = use_user_answer_loss\n        self.use_abs_pos = use_abs_pos\n        self.use_task_container_pos = use_task_container_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        self.use_part_aggregated_historical_info = use_part_aggregated_historical_info\n        self.use_correct_answer_aggregated_historical_info=use_correct_answer_aggregated_historical_info\n        self.use_question_level_aggregated_historical_info=use_question_level_aggregated_historical_info\n        self.use_lag_time = use_lag_time\n        self.use_lag_time_for_encoder = use_lag_time_for_encoder\n        self.use_user_level_aggregated_historical_info = use_user_level_aggregated_historical_info\n        self.use_pre_classifier = use_pre_classifier\n        self.use_softmax = use_softmax\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__(\n        self, config, input_name,\n        position_embeddings_layer=None,       \n        **kwargs\n    ):\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_user_answer = config.use_user_answer\n        self.use_correct_answer_for_encoder = config.use_correct_answer_for_encoder\n        self.use_correct_answer_for_decoder = config.use_correct_answer_for_decoder\n        self.use_abs_pos = config.use_abs_pos\n        self.use_task_container_pos = config.use_task_container_pos\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        self.use_lag_time = config.use_lag_time\n        self.use_lag_time_for_encoder = config.use_lag_time_for_encoder\n        self.use_user_level_aggregated_historical_info = config.use_user_level_aggregated_historical_info\n        self.use_part_aggregated_historical_info = config.use_part_aggregated_historical_info\n        self.use_correct_answer_aggregated_historical_info = config.use_correct_answer_aggregated_historical_info\n        self.use_question_level_aggregated_historical_info = config.use_question_level_aggregated_historical_info\n\n        self.dim = config.dim\n        self.seed = config.seed\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.word_embeddings = TFSharedEmbeddings(\n            self.vocab_size, config.dim, seed=config.seed, name=\"word_embeddings\"\n        )  # padding_idx=0)\n        \n        self.word_embeddings_2 = TFSharedEmbeddings(\n            ANSWER_VOCAB_SIZE, config.dim, seed=config.seed, name=\"word_embeddings_2\"\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        self.correct_answer_embeddings = tf.keras.layers.Embedding(\n            ANSWER_VOCAB_SIZE,\n            config.dim,\n            embeddings_initializer=get_initializer(config.seed),\n            name=\"correct_answer_embeddings\",\n        )\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        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        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        self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name=\"LayerNorm\")\n        self.dropout = tf.keras.layers.Dropout(0.05)\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        if self.input_name == 'response':\n            \n            with tf.name_scope(\"word_embeddings_2\"):\n                # Create and initialize weights. The random normal initializer was chosen\n                # arbitrarily, and works well.\n                self.word_embeddings_2 = self.add_weight(\n                    \"weight\", shape=[ANSWER_VOCAB_SIZE, self.dim], initializer=get_initializer(self.seed)\n                )\n\n        super().build(input_shape)\n\n    def call(\n        self,\n        input_ids=None, input_ids_2=None, position_ids=None,\n        tag_ids=None, part_ids=None,\n        prior_explanation_ids=None,\n        correct_answer_ids=None,\n        dense_embeddings_encoder=None,\n        dense_embeddings_decoder=None,\n        mode=\"embedding\", training=False\n    ):\n       \n        if mode == \"embedding\":\n            return self._embedding(\n                input_ids, input_ids_2, position_ids,\n                tag_ids, part_ids,\n                prior_explanation_ids,\n                correct_answer_ids,\n                dense_embeddings_encoder,\n                dense_embeddings_decoder,\n                training=training\n            )\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(\n        self,\n        input_ids, input_ids_2, position_ids,\n        tag_ids, part_ids,\n        prior_explanation_ids,\n        correct_answer_ids,\n        dense_embeddings_encoder,\n        dense_embeddings_decoder,\n        training=False\n    ):\n        \n        seq_length = shape_list(input_ids)[1]\n\n        inputs_embeds = tf.gather(self.word_embeddings, input_ids)\n\n        position_embeddings = tf.cast(\n            self.position_embeddings(position_ids), inputs_embeds.dtype\n        )  # (bs, max_seq_length, dim)\n        \n        if self.input_name == 'content':\n            \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.cast(\n                self.part_embeddings(part_ids), inputs_embeds.dtype\n            )  # (bs, seq_len, dim)            \n            \n            correct_answer_embeddings = tf.cast(self.correct_answer_embeddings(correct_answer_ids), inputs_embeds.dtype)\n                        \n            # shape = [n_embeddings, batch_size, seq, dim]\n            concated_embeddings = tf.concat(\n                [\n                    inputs_embeds[tf.newaxis, :, :, :],\n                    position_embeddings[tf.newaxis, :, :, :],\n                    tag_embeddings[tf.newaxis, :, :, :],\n                    part_embeddings[tf.newaxis, :, :, :],\n                    correct_answer_embeddings[tf.newaxis, :, :, :],\n                    dense_embeddings_encoder\n                ],\n                axis=0\n            )            \n            \n        if self.input_name == 'response':\n            \n            # use_user_answer\n            inputs_embeds_2 = tf.gather(self.word_embeddings_2, input_ids_2)            \n            \n            prior_explanation_embeddings = tf.cast(\n                self.prior_explanation_embeddings(prior_explanation_ids), inputs_embeds.dtype\n            )  # (bs, seq_len, dim)\n            \n            # ----------------------------------------------------------------------------------------------------\n\n            # shape = [n_embeddings, batch_size, seq, dim]\n            concated_embeddings = tf.concat(\n                [                    \n                    inputs_embeds[tf.newaxis, :, :, :],\n                    inputs_embeds_2[tf.newaxis, :, :, :],\n                    position_embeddings[tf.newaxis, :, :, :],\n                    prior_explanation_embeddings[tf.newaxis, :, :, :], \n                    dense_embeddings_decoder,\n                ],\n                axis=0\n            )\n\n        # shape = [batch_size, seq_len, dim]\n        embeddings = tf.math.reduce_sum(concated_embeddings, axis=0)\n        \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    \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\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        \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\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_user_answer = config.use_user_answer\n        self.use_correct_answer_for_encoder = config.use_correct_answer_for_encoder\n        self.use_correct_answer_for_decoder = config.use_correct_answer_for_decoder\n        self.use_abs_pos = config.use_abs_pos\n        self.use_task_container_pos = config.use_task_container_pos\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        self.use_lag_time = config.use_lag_time\n        self.use_lag_time_for_encoder = config.use_lag_time_for_encoder\n        \n        self.use_user_level_aggregated_historical_info = config.use_user_level_aggregated_historical_info\n        self.use_part_aggregated_historical_info = config.use_part_aggregated_historical_info\n        self.use_correct_answer_aggregated_historical_info = config.use_correct_answer_aggregated_historical_info\n        self.use_question_level_aggregated_historical_info = config.use_question_level_aggregated_historical_info\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            position_embeddings = None\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        self.lag_time_embeddings = tf.keras.layers.Dense(\n            config.dim,\n            kernel_initializer=get_initializer(config.seed),\n            name=\"lag_time_embeddings\",\n        )            \n        \n        self.key_names_for_encoder = {\n            'n_questions_answered_scaled',\n            'n_lectures_watched_scaled',\n            'part_count_scaled',\n            'current_part_count_scaled',\n            'correct_answer_count_scaled',\n            'current_correct_answer_count_scaled',\n            'current_question_count_scaled',\n        }\n        \n        self.key_name_need_newaxis = {\n            'n_questions_answered_scaled',\n            'n_lectures_watched_scaled',\n            'current_part_count_scaled',\n            'current_correct_answer_count_scaled',\n            'current_question_count_scaled',\n            'prior_question_elapsed_time_input',\n            'lag_time',\n            'answered_correctly_ratio',\n            'current_part_correctness_ratio',\n            'current_correct_answer_correctness_ratio',\n            'current_question_correctness_ratio',            \n        }    \n        \n        self.dense_layer_name_mapping = {\n            'n_questions_answered_scaled': ['n_questions_answered_embeddings', 'n_questions_answered_layer_norm'],\n            'n_lectures_watched_scaled': ['n_lectures_watched_embeddings', 'n_lectures_watched_layer_norm'],\n            'part_count_scaled': ['part_count_embeddings', 'part_count_layer_norm'],\n            'current_part_count_scaled': ['current_part_count_embeddings', 'current_part_count_layer_norm'],\n            'correct_answer_count_scaled': ['correct_answer_count_embeddings', 'correct_answer_count_layer_norm'],\n            'current_correct_answer_count_scaled': ['current_correct_answer_count_embeddings', 'current_correct_answer_count_layer_norm'],\n            'current_question_count_scaled': ['current_question_count_embeddings', 'current_question_count_layer_norm'],\n            'prior_question_elapsed_time_input': ['prior_question_elapsed_time_input_embeddings', 'prior_question_elapsed_time_input_layer_norm'],\n            'lag_time': ['lag_time_embeddings', 'lag_time_layer_norm'],\n            'answered_correctly_ratio': ['answered_correctly_ratio_embeddings', 'answered_correctly_ratio_layer_norm'],\n            'part_correctness_ratio': ['part_correctness_ratio_embeddings', 'part_correctness_ratio_layer_norm'],\n            'current_part_correctness_ratio': ['current_part_correctness_ratio_embeddings', 'current_part_correctness_ratio_layer_norm'],   \n            'correct_answer_correctness_ratio': ['correct_answer_correctness_ratio_embeddings', 'correct_answer_correctness_ratio_layer_norm'], \n            'current_correct_answer_correctness_ratio': ['current_correct_answer_correctness_ratio_embeddings', 'current_correct_answer_correctness_ratio_layer_norm'],   \n            'current_question_correctness_ratio': ['current_question_correctness_ratio_embeddings', 'current_question_correctness_ratio_layer_norm'],     \n        }\n\n        self.dense_layer_keys = [\n            'n_questions_answered_scaled',\n            'n_lectures_watched_scaled',\n            'part_count_scaled',\n            'current_part_count_scaled',\n            'correct_answer_count_scaled',\n            'current_correct_answer_count_scaled',\n            'current_question_count_scaled',\n            'prior_question_elapsed_time_input',\n            'lag_time',\n            'answered_correctly_ratio',\n            'part_correctness_ratio',\n            'current_part_correctness_ratio',\n            'correct_answer_correctness_ratio',\n            'current_correct_answer_correctness_ratio',\n            'current_question_correctness_ratio'  \n        ]      \n        \n        self.dense_layer_mapping = dict()        \n        \n        for key_name, layer_names in self.dense_layer_name_mapping.items():\n        \n            setattr(\n                self,\n                layer_names[0],\n                tf.keras.layers.Dense(\n                    config.dim,\n                    kernel_initializer=get_initializer(config.seed),\n                    name=layer_names[0],\n                )    \n            )\n  \n            setattr(\n                self,\n                layer_names[1],\n                tf.keras.layers.LayerNormalization(epsilon=1e-12, name=layer_names[1])  \n            )\n    \n            self.dense_layer_mapping[key_name] = [getattr(self, layer_names[0]), getattr(self, layer_names[1])]\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        self.dropout = tf.keras.layers.Dropout(0.05)\n        # --------------------------------------------------------------------------------\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        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        d_ans_input_ids = inputs.get('d_ans_input_ids')\n\n        tag_ids = inputs.get('tag_ids')\n        part_ids = inputs.get('part_ids')\n\n        # --------------------------------------------------------------------------------\n\n        batch_size = tf.math.reduce_sum(tf.ones_like(c_input_ids[:, 0], dtype=tf.int32))\n        seq_len = tf.math.reduce_sum(tf.ones_like(c_input_ids[0, :], dtype=tf.int32))\n    \n        # positional information provided\n        pos_ids = inputs['pos_ids']\n        shifted_pos_ids = inputs['shifted_pos_ids']\n\n        prior_explanation_ids = inputs.get('prior_explanation_ids')\n        \n        correct_answer_ids = inputs.get('correct_answer_id')\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        \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, seq_len, seq_len])\n        r_mask = process_mask(r_mask, [batch_size, seq_len, seq_len])\n        r_c_mask = process_mask(r_c_mask, [batch_size, seq_len, seq_len])\n        c_r_mask = process_mask(c_r_mask, [batch_size, seq_len, seq_len])        \n        \n        head_mask = [None] * self.num_hidden_layers\n            \n        dense_embeddings_encoder = []\n        dense_embeddings_decoder = []\n        for key_name, layers in self.dense_layer_mapping.items():\n            \n            dense_layer = layers[0]\n            layer_norm = layers[1]\n            \n            layer_input = inputs[key_name]\n            \n            if key_name in self.key_name_need_newaxis:\n                layer_input = layer_input[:, :, tf.newaxis]\n\n            layer_output = dense_layer(layer_input)\n            layer_output = layer_norm(\n                self.dropout(\n                    self.activation(\n                        layer_output\n                    ),\n                    training=training\n                ) + layer_output\n            )\n            \n            if key_name in self.key_names_for_encoder:\n                \n                dense_embeddings_encoder.append(layer_output)\n\n                # ------------------------------------------------------------\n                # run again\n                layer_output = dense_layer(layer_input)\n                layer_output = layer_norm(\n                    self.dropout(\n                        self.activation(\n                            layer_output\n                        ),\n                        training=training\n                    ) + layer_output\n                )\n                dense_embeddings_decoder.append(layer_output)\n                # ------------------------------------------------------------\n            else:\n                dense_embeddings_decoder.append(layer_output)\n            \n        dense_embeddings_encoder = tf.stack(dense_embeddings_encoder, axis=0)\n        dense_embeddings_decoder= tf.stack(dense_embeddings_decoder, axis=0)\n                \n        # ----------------------------------------------------------------------------------------------------\n        \n        c_embedding_output = self.content_embeddings(\n            input_ids=c_input_ids, input_ids_2=d_ans_input_ids, position_ids=pos_ids,\n            tag_ids=tag_ids, part_ids=part_ids,\n            prior_explanation_ids=prior_explanation_ids,\n            correct_answer_ids=correct_answer_ids,\n            dense_embeddings_encoder=dense_embeddings_encoder,\n            dense_embeddings_decoder=dense_embeddings_decoder,\n            training=training\n        )  # (bs, seq_length, dim)\n                \n        d_embedding_output = self.response_embeddings(\n            input_ids=d_input_ids, input_ids_2=d_ans_input_ids, position_ids=shifted_pos_ids,\n            tag_ids=tag_ids, part_ids=part_ids,\n            prior_explanation_ids=prior_explanation_ids,\n            correct_answer_ids=correct_answer_ids,\n            dense_embeddings_encoder=dense_embeddings_encoder,\n            dense_embeddings_decoder=dense_embeddings_decoder,\n            training=training\n        )  # (bs, seq_length, dim)\n\n        r_embedding_output = d_embedding_output\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    \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\n    \n    \nclass TFEdFormerAnswerPredictionModel(TFPreTrainedModel):\n    \n    def __init__(self, config, *inputs, **kwargs):\n        \n        super().__init__(config, *inputs, **kwargs)\n     \n        self.use_pre_classifier = config.use_pre_classifier\n\n        self.edformer = TFEdFormerMainLayer(config, name=\"edformer\")  # Embeddings\n\n        if self.use_pre_classifier:\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            self.pre_classifier_2 = tf.keras.layers.Dense(\n                config.dim,\n                kernel_initializer=get_initializer(config.seed),\n                name=\"pre_classifier_2\",\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.dropout = tf.keras.layers.Dropout(config.seq2seq_dropout)\n\n        n_targets = 1\n        if config.use_softmax:\n            n_targets = 2\n\n        self.classifier = tf.keras.layers.Dense(\n            n_targets, kernel_initializer=get_initializer(config.seed), name=\"classifier\"\n        )\n\n        self.classifier_2 = tf.keras.layers.Dense(\n            4, kernel_initializer=get_initializer(config.seed), name=\"classifier_2\"\n        )\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        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 = hidden_states\n        o_2 = o\n\n        if self.use_pre_classifier:\n\n            o = self.pre_classifier(o)  # (bs, seq_len, dim)\n            o = self.activation(o)\n            o = self.dropout(o, training=training)        \n        \n            o_2 = self.pre_classifier_2(o_2)  # (bs, seq_len, dim)\n            o_2 = self.activation(o_2)\n            o_2 = self.dropout(o_2, training=training)       \n\n        logits = self.classifier(o)  # (bs, seq_len, n_targets), `n_targets = 2` if using softmax. \n        answer_logits = self.classifier_2(o_2)  # (bs, seq_len, 4)\n\n        return (logits, answer_logits, c_outputs, r_outputs)","execution_count":null,"outputs":[]},{"metadata":{"id":"kgXCR1dp11te"},"cell_type":"markdown","source":"### check"},{"metadata":{"id":"pbwmHbYc11ti","outputId":"dc836f59-e10d-4247-af63-21a87c1d125f","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE or True:\n\n    content_input_ids = tf.constant(1, shape=[3, 5])\n    response_input_ids = tf.constant(1, shape=[3, 5])\n    pos_ids = tf.constant(1, shape=[3, 5])\n    shifted_pos_ids = tf.constant(1, shape=[3, 5])\n    d_input_ids = tf.constant(1, shape=[3, 5])\n    d_ans_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    lag_time = tf.constant(1.0, shape=[3, 5])\n    abs_pos_ids = tf.constant(1.0, shape=[3, 5])\n    shifted_abs_pos_ids = tf.constant(1.0, shape=[3, 5])\n    task_container_pos_ids = tf.constant(1.0, shape=[3, 5])\n    correct_answer_id = tf.constant(1, shape=[3, 5])\n    n_questions_answered_scaled = tf.constant(1.0, shape=[3, 5])\n    n_lectures_watched_scaled = tf.constant(1.0, shape=[3, 5])\n    answered_correctly_ratio = tf.constant(1.0, shape=[3, 5])\n    part_correctness_ratio = tf.constant(1.0, shape=[3, 5, PART_VOCAB_SIZE - 2])\n    part_count_scaled = tf.constant(1.0, shape=[3, 5, PART_VOCAB_SIZE - 2])\n    correct_answer_count_scaled = tf.constant(1.0, shape=[3, 5, ANSWER_3_ID - ANSWER_0_ID + 1])\n    correct_answer_correctness_ratio = tf.constant(1.0, shape=[3, 5, ANSWER_3_ID - ANSWER_0_ID + 1])\n    current_part_count_scaled = tf.constant(1.0, shape=[3, 5])\n    current_part_correctness_ratio = tf.constant(1.0, shape=[3, 5])\n    current_correct_answer_count_scaled = tf.constant(1.0, shape=[3, 5])\n    current_correct_answer_correctness_ratio = tf.constant(1.0, shape=[3, 5])\n    current_question_count_scaled = tf.constant(1.0, shape=[3, 5])\n    current_question_correctness_ratio = 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': d_input_ids,\n        'd_ans_input_ids': d_ans_input_ids,\n        'pos_ids': pos_ids,\n        'shifted_pos_ids': shifted_pos_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        'lag_time': lag_time,\n        'abs_pos_ids': abs_pos_ids,\n        'shifted_abs_pos_ids': shifted_abs_pos_ids,\n        'task_container_pos_ids': task_container_pos_ids,\n        'correct_answer_id': correct_answer_id,\n        'n_questions_answered_scaled': n_questions_answered_scaled,\n        'n_lectures_watched_scaled': n_lectures_watched_scaled,\n        'answered_correctly_ratio': answered_correctly_ratio,\n        'part_correctness_ratio': part_correctness_ratio,\n        'part_count_scaled': part_count_scaled,\n        'correct_answer_count_scaled': correct_answer_count_scaled,\n        'correct_answer_correctness_ratio': correct_answer_correctness_ratio,\n        'current_part_count_scaled': current_part_count_scaled,\n        'current_part_correctness_ratio': current_part_correctness_ratio,\n        'current_correct_answer_count_scaled': current_correct_answer_count_scaled,\n        'current_correct_answer_correctness_ratio': current_correct_answer_correctness_ratio,\n        'current_question_count_scaled': current_question_count_scaled,\n        'current_question_correctness_ratio': current_question_correctness_ratio,\n    }\n\n    model_type = 'ed'\n    config = EdFormerConfig(\n        model_type=model_type, model_desc='dummy',\n        share_position_embeddings=True, use_tags=True, user_part=True,\n        use_prior_explanation=True, use_prior_question_elapsed_time_input=True,\n        use_lag_time=True,\n        use_lag_time_for_encoder=True,\n        use_user_answer=True\n    )\n    predictor = TFEdFormerAnswerPredictionModel(config)\n\n    @tf.function\n    def foo(inputs):\n\n        logits, logits_2, c_outputs, r_outputs = predictor(inputs=inputs, output_attentions=True, output_hidden_states=True, return_dict=False)\n        return logits\n\n    logits = foo(inputs)\n    print(logits)","execution_count":null,"outputs":[]},{"metadata":{"id":"QSUmOK6p11tz"},"cell_type":"markdown","source":"## Learing rate with warmup"},{"metadata":{"id":"jQ9dlVXl11t2","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\n        # To avoid overfitting\n        # _lr = self.lr\n        _lr = tf.math.minimum(self.lr, 1 * self.end_lr)\n        \n        decay_lr = (1 - is_warmup) * (_lr - (_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":{"id":"EDS76NlPIhXe","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        'shifted_user_answer': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='shifted_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        'abs_pos_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='abs_pos_ids'),\n        'shifted_abs_pos_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='shifted_abs_pos_ids'),\n        'task_container_pos_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='task_container_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        'd_ans_input_ids': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='d_ans_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        'lag_time': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='lag_time'),\n        'n_questions_answered': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_questions_answered'),  \n        'n_lectures_watched': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_lectures_watched'),  \n        'n_questions_answered_correctly': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_questions_answered_correctly'),  \n        'n_questions_answered_scaled': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='n_questions_answered_scaled'),\n        'n_lectures_watched_scaled': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='n_lectures_watched_scaled'),\n        'answered_correctly_ratio': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='answered_correctly_ratio'),\n        'part_count_scaled': tf.TensorSpec(shape=[batch_size, seq_len, PART_VOCAB_SIZE - 2], dtype=DTYPE, name='part_count_scaled'),\n        'part_correctness_ratio': tf.TensorSpec(shape=[batch_size, seq_len, PART_VOCAB_SIZE - 2], dtype=DTYPE, name='part_correctness_ratio'),\n        'current_part_count_scaled': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_correccurrent_part_count_scaledt_answer_count_scaled'),\n        'current_part_correctness_ratio': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_part_correctness_ratio'),\n        'correct_answer_count_scaled': tf.TensorSpec(shape=[batch_size, seq_len, ANSWER_3_ID - ANSWER_0_ID + 1], dtype=DTYPE, name='correct_answer_count_scaled'),\n        'correct_answer_correctness_ratio': tf.TensorSpec(shape=[batch_size, seq_len, ANSWER_3_ID - ANSWER_0_ID + 1], dtype=DTYPE, name='correct_answer_correctness_ratio'),            \n        'current_correct_answer_count_scaled': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_correct_answer_count_scaled'),\n        'current_correct_answer_correctness_ratio': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_correct_answer_correctness_ratio'),\n        'n_prev_seen': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_prev_seen'),\n        'n_prev_correctness': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='n_prev_correctness'),\n        'current_question_count_scaled': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_question_count_scaled'),\n        'current_question_correctness_ratio': tf.TensorSpec(shape=[batch_size, seq_len], dtype=DTYPE, name='current_question_correctness_ratio'),\n        'correct_answer_id': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='correct_answer_id'),       \n        'target': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='target'),\n        'answer_target': tf.TensorSpec(shape=[batch_size, seq_len], dtype=tf.int32, name='answer_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        # input_signatures['valid_blocks_start_pos'] =\n        # input_signatures['valid_blocks_end_pos'] =\n        # input_signatures['valid_block_pos'] =\n\n    return [input_signatures]","execution_count":null,"outputs":[]},{"metadata":{"id":"srxuqrpW11ur"},"cell_type":"markdown","source":"## Train Manager"},{"metadata":{"id":"s0MnJ1m111uu","trusted":true},"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        max_n_contents_per_user_for_sampling_prob=MAX_N_CONTENTS_PER_USER_FOR_SAMPLING_PROB,\n        valid_fold=VALID_FOLD\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        self.max_n_contents_per_user_for_sampling_prob = max_n_contents_per_user_for_sampling_prob\n        self.valid_fold = valid_fold\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            ### n_training_examples += 1\n        del n_contents_dict\n\n        if self.valid_fold not in [0, 1, 2, 3]:\n            valid_info = {}\n            n_valid_examples = 0\n        else:\n            valid_info = load_data(valid_info_paths[self.valid_fold])\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        # Disabled - Let's make each epoch longer\n        n_training_steps_per_epoch = 1 * 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        n_contents_dict = load_data(n_contents_dict_path)\n        user_id_list = []\n        n_contents_list = []\n        \n        if self.train_config.valid_fold in [0, 1, 2, 3]:\n            # Get the splitted training dataset.\n            train_valid_split_indices = load_data(train_valid_split_indices_paths[self.train_config.valid_fold])\n            self.valid_user_ids = list(train_valid_split_indices.keys())\n        else:\n            train_valid_split_indices = {}\n            self.valid_user_ids = []\n        \n        for user_id, n_contents_orig in n_contents_dict.items():\n            \n            if self.train_config.valid_fold in [0, 1, 2, 3] and user_id in train_valid_split_indices:\n                split_index = train_valid_split_indices[user_id]\n                \n            else:\n                split_index = n_contents_orig\n\n            n_contents_in_train = split_index\n\n            user_id_list.append(user_id)\n            if n_contents_in_train <= 32:\n                n_contents_in_train = 32\n            n_contents_list.append(n_contents_in_train)\n\n        user_id_tensor = tf.constant(user_id_list, dtype=tf.int64)\n        n_contents_tensor = tf.constant(n_contents_list, dtype=tf.int32)\n        n_contents_tensor = tf.math.minimum(n_contents_tensor, self.train_config.max_n_contents_per_user_for_sampling_prob)\n        sample_prob_tensor = tf.cast(n_contents_tensor, dtype=tf.float32) / tf.cast(tf.math.reduce_sum(n_contents_tensor), dtype=tf.float32)    \n        # Use `0.9` for some more randomness\n        sample_prob_tensor = 0.9 * sample_prob_tensor / tf.math.reduce_max(sample_prob_tensor)\n\n        del user_id_list\n        del n_contents_list\n        del n_contents_dict\n\n        initializer = tf.lookup.KeyValueTensorInitializer(user_id_tensor, sample_prob_tensor)\n        training_sample_prob_table = tf.lookup.StaticHashTable(initializer, default_value=0.0)\n        self.training_sample_prob_table = training_sample_prob_table\n\n        self.train_ds = strategy.experimental_distribute_dataset(self.get_train_ds(from_valid=False))\n        self.train_ds_from_valid = strategy.experimental_distribute_dataset(self.get_train_ds(from_valid=True))\n        self.train_ds_from_valid_only_last = strategy.experimental_distribute_dataset(self.get_train_ds(from_valid=True, only_last=True))\n\n        self.valid_ds = strategy.experimental_distribute_dataset(self.get_valid_ds(debug=False))\n        \n        self.valid_ds_debug = self.get_valid_ds(debug=True)\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, from_valid=False, only_last=False):\n\n        only_last = False\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        # add aggregated historical information\n        train_raw_ds = train_raw_ds.map(lambda example: add_aggregated_historical_information(example), num_parallel_calls=self.train_config.num_parallel_calls, deterministic=self.train_config.deterministic)\n\n        if self.train_config.valid_fold in [0, 1, 2, 3]:\n            # Get the splitted training dataset.\n            train_valid_split_indices = load_data(train_valid_split_indices_paths[self.train_config.valid_fold])            \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        else:\n            reduced_raw_ds = train_raw_ds\n            \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),\n                    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=only_last, 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, debug=False):\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_user_id_tensor = tf.constant(self.valid_user_ids, dtype=tf.int64)\n        repeat_tensor = tf.ones_like(valid_user_id_tensor, dtype=tf.int32)\n        initializer = tf.lookup.KeyValueTensorInitializer(valid_user_id_tensor, repeat_tensor)\n        valid_users_table = tf.lookup.StaticHashTable(\n            initializer, default_value=tf.constant(0, dtype=tf.int32), name=None\n        )        \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        # don't include examples in other folds\n        valid_raw_ds.filter(lambda x: tf.math.equal(valid_users_table.lookup(x['user_id']), 1))\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        valid_ds = valid_ds.prefetch(8)\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            _lr = LinearLR(total_steps=self.train_config.n_training_steps, 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            if self.config.use_softmax:\n                loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(\n                    from_logits=True, reduction=tf.keras.losses.Reduction.NONE,\n                    name='sparse_categorical_crossentropy_for_correctness'\n                )\n            else:\n                loss_obj = tf.keras.losses.BinaryCrossentropy(\n                    from_logits=True, label_smoothing=0, reduction=tf.keras.losses.Reduction.NONE,\n                    name='binary_crossentropy_for_correctness'\n                )\n\n            loss_obj_answer = tf.keras.losses.SparseCategoricalCrossentropy(\n                from_logits=True, reduction=tf.keras.losses.Reduction.NONE,\n                name='sparse_categorical_crossentropy_for_answer'\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            loss_metric_answer = tf.keras.metrics.Sum()\n            acc_metric_answer = tf.keras.metrics.SparseCategoricalAccuracy()\n\n            # --------------------------------------------------\n\n            metrics = (loss_metric, acc_metric, auc_metric, loss_metric_answer, acc_metric_answer)\n        \n        return (predictor, optimizer,  loss_obj, loss_obj_answer, metrics)\n\n    def train_valid(\n            self, predictor, optimizer, ckpt_manager,\n            loss_obj, loss_obj_answer, metrics,\n            last_epoch=0,\n            from_valid=False, valid_epochs=None, only_valid=False\n        ):\n\n        use_user_answer_loss = tf.cast(tf.constant(self.config.use_user_answer_loss), dtype=DTYPE)\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, loss_metric_answer, acc_metric_answer = 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            answer_targets = train_batch['answer_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            selected_answer_targets = tf.gather_nd(answer_targets, pred_indices)\n            \n            # Compute the gradients for a list of variables.\n            with tf.GradientTape() as tape:\n\n                # `logits`: shape = [batch_size, seq_len, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n                (logits, answer_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                \n                # shape = [n_selected_places, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n                selected_logits = tf.gather_nd(logits, pred_indices)\n                \n                # shape = [n_selected_places, 4]\n                selected_answer_logits = tf.gather_nd(answer_logits, pred_indices)\n\n                # From the doc, `tf.keras.losses.SparseCategoricalCrossentropy` should use `selected_targets`, but it seems ok to have the last dimension.\n                # For `BinaryCrossentropy`, we need to use `selected_targets[:, tf.newaxis]` to have the 2nd dimension so the losses are not averaged.\n                losses = loss_obj(selected_targets[:, tf.newaxis], selected_logits)\n                \n                answer_losses = loss_obj_answer(selected_answer_targets, selected_answer_logits)\n                \n                total_loss = tf.math.reduce_sum(losses)\n                total_answer_loss = tf.math.reduce_sum(answer_losses)\n\n                # `train_batch['nb_pred_places'][0]` is the total number of places used for calculating loss\n                loss = (total_loss + total_answer_loss * use_user_answer_loss) / tf.cast(train_batch['nb_pred_places'][0], dtype=DTYPE)\n                            \n            grads = tape.gradient(loss, predictor.trainable_variables)\n\n            # Process the gradients, for example cap them, etc.\n            # capped_grads = [MyCapper(g) for g in grads]\n            ### processed_grads = [process_gradient(g) for g in grads]\n            ### processed_grads = grads\n\n            # Ask the optimizer to apply the processed gradients.\n            optimizer.apply_gradients(zip(grads, predictor.trainable_variables))\n            \n            if self.config.use_softmax:\n                # shape = [batch_size, seq_len, 2]\n                preds = tf.math.softmax(logits)\n                # shape = [n_selected_places, 2]\n                selected_preds = tf.math.softmax(selected_logits)\n            else:\n                # shape = [batch_size, seq_len, 1]\n                preds = tf.math.sigmoid(logits)\n                # shape = [n_selected_places, 1]\n                selected_preds = tf.math.sigmoid(selected_logits)\n            \n            answer_preds = tf.math.softmax(answer_logits, axis=-1)\n            selected_answer_preds = tf.math.softmax(selected_answer_logits, axis=-1)\n            \n            loss_metric.update_state(total_loss)\n            \n            # Use `selected_preds[:, -1:]` to get the probabilities for class `1`, with the 2nd dim\n            acc_metric.update_state(selected_targets[:, tf.newaxis], selected_preds[:, -1:])\n            # Use `selected_preds[:, -1]` to get the probabilities for class `1`, without the 2nd dim\n            auc_metric.update_state(selected_targets, selected_preds[:, -1])\n\n            loss_metric_answer.update_state(total_answer_loss)\n            acc_metric_answer.update_state(selected_answer_targets, selected_answer_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            # 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, answer_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'] * tf.cast(valid_batch['pred_time_mask'] != PAD_TOKEN, dtype=tf.int32)\n            targets = valid_batch['target']\n            answer_targets = valid_batch['answer_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            answer_targets = answer_targets[:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:]\n            logits = logits[:, -MAX_PRED_TIME_QUESTION_BUNDLE_LEN:, :]\n            answer_logits = answer_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            answer_targets = answer_targets * pred_time_mask + NON_TARGET_ID * (1 - pred_time_mask)\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\n            if self.config.use_softmax:\n                # shape = [batch_size, seq_len, 2]\n                preds = tf.math.softmax(logits)\n                # shape = [n_selected_places, 2]\n                selected_preds = tf.math.softmax(selected_logits)\n            else:\n                # shape = [batch_size, seq_len, 1]\n                preds = tf.math.sigmoid(logits)\n                # shape = [n_selected_places, 1]\n                selected_preds = tf.math.sigmoid(selected_logits)            \n\n            answer_preds = tf.math.softmax(answer_logits, axis=-1)\n\n            selected_answer_targets = tf.gather_nd(answer_targets, pred_indices)\n            selected_answer_logits = tf.gather_nd(answer_logits, pred_indices)\n            selected_answer_preds = tf.math.softmax(selected_answer_logits, axis=-1)\n\n            # From the doc, `tf.keras.losses.SparseCategoricalCrossentropy` should use `selected_targets`, but it seems ok to have the last dimension.\n            # For `BinaryCrossentropy`, we need to use `selected_targets[:, tf.newaxis]` to have the 2nd dimension so the losses are not averaged.\n            losses = loss_obj(selected_targets[:, tf.newaxis], selected_logits)\n            \n            answer_losses = loss_obj_answer(selected_answer_targets, selected_answer_logits)\n\n            total_loss = tf.math.reduce_sum(losses)\n            total_answer_loss = tf.math.reduce_sum(answer_losses)\n\n            # `train_batch['nb_pred_places'][0]` is the total number of places used for calculating loss\n            loss = (total_loss + total_answer_loss * use_user_answer_loss) / tf.cast(tf.math.reduce_mean(valid_batch['nb_pred_places']), dtype=DTYPE)\n\n            loss_metric.update_state(total_loss)\n            \n            # Use `selected_preds[:, -1:]` to get the probabilities for class `1`, with the 2nd dim\n            acc_metric.update_state(selected_targets[:, tf.newaxis], selected_preds[:, -1:])\n            # Use `selected_preds[:, -1]` to get the probabilities for class `1`, without the 2nd dim\n            auc_metric.update_state(selected_targets, selected_preds[:, -1])\n\n            loss_metric_answer.update_state(total_answer_loss)\n            acc_metric_answer.update_state(selected_answer_targets, selected_answer_preds)   \n\n            return targets, preds, row_ids, user_ids, pred_mask, answer_targets, answer_preds\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, answer_targets, answer_preds = strategy.run(valid_step, args=(dist_valid_batch,))\n \n            return targets, preds, row_ids, user_ids, pred_mask, answer_targets, answer_preds\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            t_arr_6 = tf.TensorArray(tf.int32, size=0, dynamic_size=True, clear_after_read=False)\n            t_arr_7 = tf.TensorArray(DTYPE, 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, _answer_targets, _answer_preds = 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                    _answer_targets = tf.concat(_answer_targets.values, axis=0)\n                    _answer_preds = tf.concat(_answer_preds.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                total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\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                t_arr_6 = t_arr_6.write(step_idx, _answer_targets)\n                t_arr_7 = t_arr_7.write(step_idx, _answer_preds)\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            answer_targets = t_arr_6.concat()\n            answer_preds = t_arr_7.concat()\n\n            return targets, preds, row_ids, user_ids, pred_mask, total_nb_pred_places, answer_targets, answer_preds\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            t_arr_6 = tf.TensorArray(tf.int32, size=0, dynamic_size=True, clear_after_read=False)\n            t_arr_7 = tf.TensorArray(DTYPE, 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, _answer_targets, _answer_preds = 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                    _answer_targets = tf.concat(_answer_targets.values, axis=0)\n                    _answer_preds = tf.concat(_answer_preds.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                total_nb_pred_places += tf.cast(nb_pred_places, dtype=DTYPE)\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                t_arr_6 = t_arr_6.write(step_idx, _answer_targets)\n                t_arr_7 = t_arr_7.write(step_idx, _answer_preds)\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            answer_targets = t_arr_6.concat()\n            answer_preds = t_arr_7.concat()\n\n            return targets, preds, row_ids, user_ids, pred_mask, total_nb_pred_places, answer_targets, answer_preds\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        # TODO: Not use hard coded numbers\n        if from_valid:\n            n_epochs = 3\n            n_training_examples = int(2380588 / self.train_config.window_size)\n            n_training_steps_per_epoch = int(n_training_examples / self.train_config.batch_size)\n            n_training_steps = n_epochs * n_training_steps_per_epoch\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, from_valid=False, 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            if from_valid:\n                end_epoch = start_epoch + 3\n                train_ds = self.train_ds_from_valid\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                    'answer_loss': [],\n                    'answer_acc': [],\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                    avg_answer_loss = (loss_metric_answer.result() / total_nb_pred_places).numpy()\n                    answer_acc = acc_metric_answer.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                        training_history[epoch]['answer_loss'].append(float(avg_answer_loss))\n                        training_history[epoch]['answer_acc'].append(float(answer_acc))\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'answer_loss: {avg_answer_loss}')\n                        print(f'answer_acc: {answer_acc}')\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                training_history[epoch]['train_answer_loss'] = float(avg_answer_loss)\n                training_history[epoch]['train_answer_acc'] = float(answer_acc)\n\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                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(f'train answer loss: {avg_answer_loss}')\n                print(f'train answer acc: {answer_acc}')\n                print('-' * 80)\n\n                loss_metric.reset_states()\n                acc_metric.reset_states()\n                auc_metric.reset_states()\n                loss_metric_answer.reset_states()\n                acc_metric_answer.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                    'answer_loss': [],\n                    'answer_acc': [],\n                    'timing': []        \n                }                \n\n            valid_user_ids = []\n            valid_row_ids = []\n            valid_targets = []\n            valid_preds = []\n            valid_answer_targets = []\n            valid_answer_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            n_valid_calls = self.train_config.n_valid_calls\n            if self.train_config.n_steps_in_last_valid_call > 0:\n                n_valid_calls -= 1\n\n            for call_idx in range(n_valid_calls):\n\n                start_t = datetime.datetime.now()\n\n                targets, preds, row_ids, user_ids, pred_mask, nb_pred_places, answer_targets, answer_preds = 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                \n                # shape = [n_selected_places, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n                selected_preds = tf.gather_nd(preds, pred_indices)\n\n                selected_answer_targets = tf.gather_nd(answer_targets, pred_indices)\n                selected_answer_preds = tf.gather_nd(answer_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                avg_answer_loss = (loss_metric_answer.result() / total_nb_pred_places).numpy()\n                answer_acc = acc_metric_answer.result().numpy()\n\n                elapsed = (datetime.datetime.now() - start_t).total_seconds()\n\n                if call_idx % 1 == 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'valid auc: {auc}')\n                    print(f'valid answer loss: {avg_answer_loss}') \n                    print(f'valid answer acc: {answer_acc}')        \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                \n                # Use `selected_preds[:, -1]` to get the probabilities for class `1`, without the 2nd dim\n                valid_preds.extend(selected_preds[:, -1].numpy().tolist())\n                \n                valid_answer_targets.extend(selected_answer_targets.numpy().tolist())\n                valid_answer_preds.extend(selected_answer_preds.numpy().tolist())\n\n            # ----------------------------------------------------------------------------------------------------\n            # Last TPU call\n\n            if self.train_config.n_steps_in_last_valid_call > 0:\n\n                start_t = datetime.datetime.now()\n\n                targets, preds, row_ids, user_ids, pred_mask, nb_pred_places, answer_targets, answer_preds = dist_valid_multi_steps_last_call(dist_valid_iter)\n                n_steps += self.train_config.n_steps_in_last_valid_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_answer_targets = tf.gather_nd(answer_targets, pred_indices)\n                selected_answer_preds = tf.gather_nd(answer_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                avg_answer_loss = (loss_metric_answer.result() / total_nb_pred_places).numpy()\n                answer_acc = acc_metric_answer.result().numpy()\n\n                elapsed = (datetime.datetime.now() - start_t).total_seconds()\n\n                valid_row_ids.extend(selected_row_ids.numpy())\n                valid_user_ids.extend(selected_user_ids.numpy())\n                \n                # Use `selected_preds[:, -1]` to get the probabilities for class `1`, without the 2nd dim\n                valid_preds.extend(selected_preds[:, -1].numpy())\n\n                valid_answer_targets.extend(selected_answer_targets.numpy())\n                valid_answer_preds.extend(selected_answer_preds.numpy())\n\n            # ----------------------------------------------------------------------------------------------------\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            training_history[epoch]['valid_answer_loss'] = float(avg_answer_loss)\n            training_history[epoch]['valid_answer_acc'] = float(answer_acc)\n\n            print(f'epoch: {epoch} - valid')\n            print(f'valid loss: {avg_loss}')\n            print(f'valid acc: {acc}')\n            print(f'valid auc: {auc}')\n            print(f'valid answer loss: {avg_answer_loss}') \n            print(f'valid answer acc: {answer_acc}')                     \n            print('=' * 80)\n\n            loss_metric.reset_states()\n            acc_metric.reset_states()\n            auc_metric.reset_states()\n            loss_metric_answer.reset_states()\n            acc_metric_answer.reset_states()\n            \n            # --------------------------------------------------\n\n            valid_submission_dict = {\n                'row_ids': valid_row_ids,\n                'user_ids': valid_user_ids,\n                'targets': valid_targets,\n                'preds': valid_preds\n            }\n\n            if self.config.use_user_answer_loss:\n\n                valid_answer_preds = np.array(valid_answer_preds)\n\n                valid_submission_dict['answer_targets'] = valid_answer_targets\n                valid_submission_dict['answer_preds_0'] = valid_answer_preds[:, 0]\n                valid_submission_dict['answer_preds_1'] = valid_answer_preds[:, 1]\n                valid_submission_dict['answer_preds_2'] = valid_answer_preds[:, 2]\n                valid_submission_dict['answer_preds_3'] = valid_answer_preds[:, 3]\n\n\n            # valid_submission = pd.DataFrame.from_dict(valid_submission_dict)\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, from_valid, valid_epochs)","execution_count":null,"outputs":[]},{"metadata":{"id":"HiusgP5gPlNg"},"cell_type":"markdown","source":"#### check"},{"metadata":{"id":"rCf0qm5pPkQA","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n    \n    config = EdFormerConfig(model_type='ed', model_desc='dummy')\n    train_config = TrainConfig()\n    train_manager = Train_Manager(config, train_config)","execution_count":null,"outputs":[]},{"metadata":{"id":"WsH3YIL4ZDUo","outputId":"c79c209c-615e-4378-f92d-63fabe09e964","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n    train_manager.train_ds","execution_count":null,"outputs":[]},{"metadata":{"id":"Pu5X9U1QZBUU","outputId":"b99f8e52-5350-47f1-8803-0b468ff77baa","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n    train_manager.valid_ds","execution_count":null,"outputs":[]},{"metadata":{"id":"8pZoTc8CfRHn","outputId":"06538a0a-945b-4abd-efa4-913f305e431a","trusted":true},"cell_type":"code","source":"if not IS_KAGGLE:\n    train_manager.valid_ds_debug","execution_count":null,"outputs":[]},{"metadata":{"id":"kQH29iRS11u9"},"cell_type":"markdown","source":"# Iter Manager"},{"metadata":{"id":"b8YFZCamtnrv"},"cell_type":"markdown","source":"## Utilitites"},{"metadata":{"id":"In05zzkbth7C","trusted":true},"cell_type":"code","source":"def convert_dt(_dt, pred=False, unique_question_id_train=None, unique_lecture_id_train=None):\n    \"\"\"\n    Change column type, deal with NaN value. If it is a `datatable.Frame` from the test dataset,\n    change it to a format suitable for prediction.\n\n    Args:\n        _dt: `datatable.Frame`, representing a block of the training dataset or a test batch given by `env.iter_test`.\n    \"\"\"\n\n    _dt[dt.f.prior_question_elapsed_time] = dt.float32\n    _dt[dt.f.prior_question_elapsed_time == None, 'prior_question_elapsed_time'] = -1.0\n    _dt[dt.f.prior_question_had_explanation] = dt.int8\n    _dt[dt.f.prior_question_had_explanation == None, 'prior_question_had_explanation'] = -1\n    _dt[dt.f.content_type_id] = dt.int8\n\n    if pred:\n\n        _dt['answered_correctly'] = MASK_TOKEN\n        _dt['user_answer'] = MASK_TOKEN\n\n        try:\n            del _dt['prior_group_answers_correct']\n            del _dt['prior_group_responses']\n            # del _dt['group_num']\n        except:\n            pass\n\n    if DEBUG:\n        \n        if unique_question_id_train is not None:\n            # All test questions must have been seen in training time.\n            assert set(_dt[dt.f.content_type_id == 0, 'content_id'].to_list()[0]).issubset(unique_question_id_train)\n        if unique_lecture_id_train is not None:    \n            # All test lectures must have been seen in training time.\n            assert set(_dt[dt.f.content_type_id == 1, 'content_id'].to_list()[0]).issubset(unique_lecture_id_train)\n            \n            \ndef convert_df_to_dt(df, test=False, unique_question_id_train=None, unique_lecture_id_train=None):\n    \"\"\"Convert a `pandas.DataFrame` to `datatable.Frame` with some extra processing.\n\n    Args:\n        df: `pandas.DataFrame`, representing a block of the training dataset or a test batch given by `env.iter_test()`.    \n    \"\"\"\n\n    if DEBUG:\n\n        _question_ids = set(df[df['content_type_id'] == 0]['content_id'].values.tolist())\n        _lecture_ids = set(df[df['content_type_id'] == 1]['content_id'].values.tolist())\n\n        if unique_question_id_train is not None:\n            # All test questions must have been seen in training time.\n            assert _question_ids.issubset(unique_question_id_train)\n\n        if unique_lecture_id_train is not None:    \n            # All test lectures must have been seen in training time.\n            assert _lecture_ids.issubset(unique_lecture_id_train)\n\n    _dt = dt.Frame(df.astype({\"prior_question_had_explanation\": float}))\n\n    if test:\n\n        prior_group_answers_correct = eval(_dt[0, 'prior_group_answers_correct'])\n        prior_group_responses = eval(_dt[0, 'prior_group_responses'])\n\n        if prior_group_answers_correct is None:\n            prior_group_answers_correct = []\n\n        if prior_group_responses is None:\n            prior_group_responses = []\n\n        if DEBUG:\n\n            assert type(prior_group_answers_correct) == list\n            assert type(prior_group_responses) == list\n\n    convert_dt(_dt, pred=test, unique_question_id_train=unique_question_id_train, unique_lecture_id_train=unique_lecture_id_train)\n\n    if test:\n        return _dt, prior_group_answers_correct, prior_group_responses\n    else:\n        return _dt","execution_count":null,"outputs":[]},{"metadata":{"id":"pxczjuZp11vH"},"cell_type":"markdown","source":"## Pred Iter Manager"},{"metadata":{"id":"wJM1Gciy11vK","trusted":true},"cell_type":"code","source":"class Pred_Iter_Manager:\n    \n    def __init__(self, unique_question_id_train=None, unique_lecture_id_train=None):\n                \n        self.current_batch_no = 0\n        \n        self._unique_question_id_train = unique_question_id_train\n        self._unique_lecture_id_train = unique_lecture_id_train\n                \n    def get_split_index(self, user_id):\n\n        # `-1` means that `user_id` is not in `valid_info.json`.\n        # For the test dataset, `valid_info.json` is not used, and any `user_id`, if it is in the original training dataset,\n        # the model is trained for this user.\n        return -1 \n        \n    def _step(self, n_users=None):\n        \n        raise NotImplementedError\n        \n        # return pred_df, pred_dt, prior_group_answers_correct, prior_group_responses, n_blocks_in_batch\n        \n    def iter_dataset(self):\n        \n        raise NotImplementedError","execution_count":null,"outputs":[]},{"metadata":{"id":"0u1kvYYv11vo"},"cell_type":"markdown","source":"## Valid Iter Manager"},{"metadata":{"id":"FVbRM7WUt47p"},"cell_type":"markdown","source":"### Function to sample sequence lengths"},{"metadata":{"id":"pMcCHE2Mt159","trusted":true},"cell_type":"code","source":"def sample_sequence_lengths(upper_bound, sample_size, expected_mean, seed):\n    \n    seq_lengths = tf.range(1, upper_bound+1)\n    seq_len_mean = tf.math.reduce_mean(tf.cast(seq_lengths, dtype=DTYPE))\n    \n    N = sample_size\n    mu = expected_mean\n    sigma = mu / 2.5\n    \n    normal_dist = tfp.distributions.Normal(\n        loc=mu, scale=sigma\n    )\n\n    # But the values might not in seq_lengths.\n    # Extreme case: seq_lengths is a constant --> only 1 value is allowed ...\n    seq_len_sample = normal_dist.sample(sample_shape=[N], seed=seed)\n    \n    avg_selected_seq_len = tf.math.reduce_mean(seq_len_sample)\n    \n    mask = tf.math.logical_and(2 * mu >= seq_len_sample, seq_len_sample >= 1)\n\n    # About 95 % will be kept\n    seq_len_sample = tf.gather_nd(seq_len_sample, indices=tf.where(mask))\n    \n    seq_len_sample = tf.cast(tf.math.round(seq_len_sample), dtype=tf.int32)\n    \n    uniform_dist = tfp.distributions.Uniform(\n        low=0, high=len(seq_lengths)\n    )\n\n    # For the remaining 5 %, let's take the uniform distribution\n    selected_indices = tf.cast(uniform_dist.sample(sample_shape=[N - seq_len_sample.shape[0]], seed=seed), dtype=tf.int32)\n\n    seq_len_sample_2 = tf.gather(seq_lengths, indices=selected_indices)\n\n    seq_len_sample_2 = seq_len_sample_2\n    \n    seq_len_sample = tf.concat([seq_len_sample, seq_len_sample_2], axis=0)\n        \n    return seq_len_sample","execution_count":null,"outputs":[]},{"metadata":{"id":"uneL278qt_3t"},"cell_type":"markdown","source":"### Valid Iter Manager"},{"metadata":{"id":"sRyzZLqw11vs","trusted":true},"cell_type":"code","source":"class User_Interaction_Manager(collections.abc.Iterator):\n    \n    def __init__(self, user_id, user_valid_info, train_dt):\n                \n        self.user_id = user_id\n        self.user_valid_info = user_valid_info\n        \n        self._n_remaining_blocks = self.n_blocks\n        self._current_block_starting_index = 0\n        \n        self._train_dt = train_dt\n        \n    @property\n    def n_blocks(self):\n        \n        return self.user_valid_info['bundle_info']['n_blocks']\n    \n    @property\n    def n_remaining_blocks(self):\n        \n        return self._n_remaining_blocks\n    \n    def _get_next_block_starting_index(self):\n        \n        if self._current_block_starting_index not in self.user_valid_info['bundle_info']['block_starting_index_dict']:\n            next_index = self._current_block_starting_index + 1\n        else:\n            next_index = self.user_valid_info['bundle_info']['block_starting_index_dict'][self._current_block_starting_index]\n        \n        if DEBUG:\n            assert 0 <= next_index <= self.user_valid_info['seq_len']\n        \n        return next_index\n        \n    def _get_user_dt(self, next_block_starting_index):\n        \n        row_start = self.user_valid_info['row_start']\n        \n        block_row_start = row_start + self._current_block_starting_index\n        block_row_end = row_start + next_block_starting_index - 1\n        \n        if DEBUG:\n            assert block_row_end <= self.user_valid_info['row_end']\n        \n        user_dt = self._get_user_dt_from_row_indices(block_row_start, block_row_end)\n        \n        return user_dt\n    \n    def _get_user_dt_from_row_indices(self, row_start, row_end):\n    \n        user_dt = self._train_dt[row_start:(row_end + 1), :]\n        \n        return user_dt\n    \n    def has_next(self):\n        \n        return not (self._n_remaining_blocks <= 0 or self._current_block_starting_index >= self.user_valid_info['seq_len'])\n    \n    def __next__(self):\n        \n        if self._n_remaining_blocks <= 0:\n            \n            # raise ValueError('User has no more remaining block!')\n            raise StopIteration\n            \n        if self._current_block_starting_index >= self.user_valid_info['seq_len']:\n\n            # raise ValueError('Block starting index >= seq_len, cannot step anymore!')\n            raise StopIteration\n            \n        next_block_starting_index = self._get_next_block_starting_index()\n\n        user_dt = self._get_user_dt(next_block_starting_index)\n        \n        # reset the state\n        self._n_remaining_blocks -= 1        \n        self._current_block_starting_index = next_block_starting_index\n        \n        if DEBUG:\n            \n            # sanity check\n            if self._n_remaining_blocks <= 0:            \n                assert self._current_block_starting_index == self.user_valid_info['seq_len']\n            if self._current_block_starting_index >= self.user_valid_info['seq_len']:\n                assert self._n_remaining_blocks == 0\n\n            assert user_dt.nrows > 0    \n            \n        return user_dt\n        \n\nclass Valid_Iter_Manager(Pred_Iter_Manager):\n    \n    def __init__(self, valid_info, train_dt, unique_question_id_train=None, unique_lecture_id_train=None):\n\n        super(Valid_Iter_Manager, self).__init__(unique_question_id_train, unique_lecture_id_train)\n        \n        self._valid_info = valid_info\n        \n        self._n_users = len(set(self._valid_info.keys()))\n\n        _n_blocks = 0\n        for user_id, user_valid_info in self._valid_info.items():\n            _n_blocks += user_valid_info['bundle_info']['n_blocks']\n        self._n_blocks = _n_blocks\n\n        self._remaining_users = set(self._valid_info.keys())\n        self._n_remaining_blocks = self._n_blocks\n                \n        self._current_answered_correctly = None\n        self._current_user_answer = None\n        \n        self._user_manager_dict = dict()\n        for user_id, user_valid_info in self._valid_info.items():\n            self._user_manager_dict[user_id] = User_Interaction_Manager(user_id, user_valid_info, train_dt)\n\n        self.train_dt = train_dt\n        \n        self._split_indices = {user_id: valid_info_for_user['split_index'] for user_id, valid_info_for_user in self._valid_info.items()}\n        \n    @classmethod\n    def create(cls, valid_info_path, train_dt_path_or_obj, unique_question_id_train_path=None, unique_lecture_id_train_path=None):\n\n        args = []\n        if valid_info_path is None:\n            valid_info = {}\n        else:\n            valid_info = load_data(valid_info_path)\n        args.append(valid_info)\n            \n        for p in [unique_question_id_train_path, unique_lecture_id_train_path]:\n            if p is not None:\n                args.append(load_data(p))\n                \n        if type(train_dt_path_or_obj) == str:\n            train_dt_path_or_obj = load_data(train_dt_path_or_obj)\n            \n        args = args[:1] + [train_dt_path_or_obj] + args[1:]\n        \n        return cls(*args)\n                \n    def get_split_index(self, user_id):\n        \n        # This means that no interaction for `user_id` is included in the validation dataset.\n        if user_id not in self._split_indices:\n            return -1\n                    \n        if DEBUG:\n            # `self._split_indices[user_id]` could never be `seq_len`.\n            assert 0 <= self._split_indices[user_id] < self._valid_info[user_id]['orig_seq_len']\n        \n        # If `self._split_indices[user_id]` is `0`, the user with `user_id` is a user unseen in the splitted training dataset.\n        return self._split_indices[user_id]\n        \n    def iter_dataset(self, max_steps=None):\n        \n        print('entering iter_dataset()')\n\n        expected_mean_blocks_per_batch = 30.0\n        expected_n_batches = int(self._n_blocks / expected_mean_blocks_per_batch) + 1\n        \n        \n        ### sampled_block_size = sample_sequence_lengths(upper_bound=40, sample_size=expected_n_batches, expected_mean=expected_mean_blocks_per_batch, seed=SEED)\n        sampled_block_size = [30] * 100000\n        \n        ### sample_mean_blocks_per_batch = tf.math.reduce_mean(tf.cast(sampled_block_size, dtype=DTYPE))\n        sample_mean_blocks_per_batch = 30.0\n        \n        ### sampled_block_size = tf.random.shuffle(sampled_block_size, seed=SEED)\n        ### sampled_block_size = sampled_block_size.numpy().tolist()\n        \n        if DEBUG:\n            assert len(sampled_block_size) >= expected_n_batches\n            \n        print(f'expected_mean_blocks_per_batch_in_pred: {expected_mean_blocks_per_batch}')\n        print(f'sample_mean_blocks_per_batch_in_pred: {sample_mean_blocks_per_batch}')\n        print(f'expected_n_batches_in_pred: {expected_n_batches}')\n        \n        total_blocks_yield = 0\n        \n        print('start looping ...')\n        while True:\n            \n            if max_steps is not None and self.current_batch_no >= max_steps:\n                break            \n            \n            try:\n                \n                n_blocks_to_select = 30.0\n                if self.current_batch_no < len(sampled_block_size):\n                    n_blocks_to_select = sampled_block_size[self.current_batch_no]    \n                    \n                valid_df, valid_dt, prior_group_answers_correct, prior_group_responses, n_blocks_in_batch = self._step(n_blocks_to_select)\n\n                if DEBUG:\n                    assert n_blocks_in_batch > 0\n                \n                total_blocks_yield += n_blocks_in_batch\n\n                # sanity check\n                if total_blocks_yield > self._n_blocks:\n                    raise ValueError('total_blocks_yield > self._n_blocks! The code has some bugs to fix!')\n                \n                if DEBUG:\n                    assert total_blocks_yield == self._n_blocks - self._n_remaining_blocks\n\n                self.current_batch_no += 1\n\n                if self.current_batch_no % PRINTING_STEPS == 0:\n\n                    print(f'current_batch_no: {self.current_batch_no}')\n                    print(f'n_blocks processed in current batch: {n_blocks_in_batch}')\n                    print(f'n_blocks processed: {total_blocks_yield}')\n                    print(f'averaged n_blocks processed per batch: {total_blocks_yield / self.current_batch_no}')\n                    print(f'n_users_completed: {self._n_users - len(self._remaining_users)}')\n                    \n                yield valid_df, valid_dt, prior_group_answers_correct, prior_group_responses\n    \n            except StopIteration:\n                break    \n    \n        # sanity check\n        if DEBUG:\n            assert len(self._user_manager_dict) == 0\n            assert len(self._remaining_users) == 0\n            assert self._n_remaining_blocks == 0    \n\n        print('finished looping.')\n\n    def _step(self, n_users):\n        \n        return self._blocks_from_random_users(n_users)\n            \n    def _blocks_from_random_users(self, n_users):\n        \n        if DEBUG:\n            assert len(self._remaining_users) >= 0\n        \n        if len(self._remaining_users) == 0:\n            raise StopIteration\n        \n        n_users_to_select = min(n_users, len(self._remaining_users))\n                \n        # randomly select `n_users`.\n        selected_users = random.sample(self._remaining_users, k=n_users_to_select)\n        \n        if DEBUG:\n            assert len(selected_users) == n_users_to_select\n        \n        # For each selected user, get its 1st remained block (in `dt` format) and reset its states.\n        user_dts = []\n        for user_id in selected_users:\n            \n            if self._user_manager_dict[user_id].has_next():\n\n                user_dts.append(next(self._user_manager_dict[user_id]))\n\n                if not self._user_manager_dict[user_id].has_next():\n                    del self._user_manager_dict[user_id]\n                    self._remaining_users.remove(user_id)\n            \n        # Combine dt Frames\n        valid_dt = dt.rbind(user_dts)\n         \n        n_blocks_selected = n_users_to_select\n        self._n_remaining_blocks -= n_blocks_selected\n        \n        # --------------------------------------------------------------------------------\n        \n        prior_group_answers_correct = self._current_answered_correctly\n        prior_group_responses = self._current_user_answer\n\n        if prior_group_answers_correct is None:\n            prior_group_answers_correct = []\n            \n        if prior_group_responses is None:\n            prior_group_responses = []\n\n        if DEBUG:\n            \n            assert type(prior_group_answers_correct) == list\n            assert type(prior_group_responses) == list\n            \n        # --------------------------------------------------------------------------------\n        \n        self._current_answered_correctly = valid_dt[:, 'answered_correctly'].to_list()[0]\n        self._current_user_answer = valid_dt[:, 'user_answer'].to_list()[0]\n        \n        # --------------------------------------------------------------------------------        \n        \n        convert_dt(valid_dt, pred=True, unique_question_id_train=self._unique_question_id_train, unique_lecture_id_train=self._unique_lecture_id_train)\n            \n        valid_df = valid_dt.to_pandas()\n\n        return valid_df, valid_dt, prior_group_answers_correct, prior_group_responses, n_blocks_selected","execution_count":null,"outputs":[]},{"metadata":{"id":"MnKxCXRZ11v4"},"cell_type":"markdown","source":"## Test Iter Manager"},{"metadata":{"id":"ti4Xru0D11v6","trusted":true},"cell_type":"code","source":"class Test_Iter_Manager(Pred_Iter_Manager):\n    \n    def __init__(self, unique_question_id_train=None, unique_lecture_id_train=None):\n        \n        super(Test_Iter_Manager, self).__init__(unique_question_id_train, unique_lecture_id_train)\n\n        self.env = riiideducation.make_env()\n        self._iter_test = iter(self.env.iter_test())\n        \n    @classmethod\n    def create(cls, unique_question_id_train_path=None, unique_lecture_id_train_path=None):\n        \n        args = []\n        for p in [unique_question_id_train_path, unique_lecture_id_train_path]:\n            if p is not None:\n                args.append(load_data(p))\n        \n        return cls(*args)\n      \n    def _has_next(self):\n        \n        return self._has_next\n        \n    def iter_dataset(self, max_steps=None):\n        \n        print('entering iter_dataset()')\n        \n        total_blocks_yield = 0\n        \n        print('start looping ...')\n        while True:\n            \n            try:\n                \n                test_df, test_dt, prior_group_answers_correct, prior_group_responses, n_blocks_in_batch = self._step(n_users=None)\n                assert n_blocks_in_batch > 0\n                total_blocks_yield += n_blocks_in_batch\n                self.current_batch_no += 1\n                \n                if self.current_batch_no % PRINTING_STEPS == 0:\n                \n                    print(f'current_batch_no: {self.current_batch_no}')\n                    print(f'n_blocks processed in current batch: {n_blocks_in_batch}')\n                    print(f'n_blocks processed: {total_blocks_yield}')\n                    print(f'averaged n_blocks processed per batch: {total_blocks_yield / self.current_batch_no}')                    \n\n                yield test_df, test_dt, prior_group_answers_correct, prior_group_responses\n                \n            except StopIteration:\n                break\n            \n        print('finished looping.')\n                \n    def _step(self, n_users=None):\n        \n        test_df, _ = next(self._iter_test)\n        \n        # conversion\n        test_dt, prior_group_answers_correct, prior_group_responses = convert_df_to_dt(\n            test_df, test=True,\n            unique_question_id_train=self._unique_question_id_train,\n            unique_lecture_id_train=self._unique_lecture_id_train\n        )        \n\n        n_blocks = len(set(test_dt[:, 'user_id'].to_list()[0]))\n        \n        return test_df, test_dt, prior_group_answers_correct, prior_group_responses, n_blocks","execution_count":null,"outputs":[]},{"metadata":{"id":"0aKwbtV-11wE"},"cell_type":"markdown","source":"# Prediction Manager"},{"metadata":{"id":"vP4-KLsy11wI","trusted":true},"cell_type":"code","source":"def get_user_dt(user_id, _dt, user_id_to_row_id, split_index=-1):\n    \"\"\"Get the partial `datatable.Frame` in `_dt` containing only `user_id`.\n    \n    Args:\n        user_id: `int`. It must in `user_id_to_row_id`, which should be precomputed from `_dt`.\n        _dt: `datatable.Frame`.\n        user_id_to_row_id: `dict`. See the above markdown cell for the format. \n    \"\"\"\n\n#     assert user_id in user_id_to_row_id\n    \n    (row_id_start, row_id_end) = user_id_to_row_id[user_id]\n#     assert split_index == -1 or (0 <= split_index <= (row_id_end - row_id_start) + 1)\n    \n    if split_index > -1:\n        row_id_end = row_id_start + split_index - 1 \n    \n    user_dt = _dt[row_id_start:row_id_end + 1, :]\n\n    return user_dt\n\n\nattrs = [\n    'user_id',\n    'row_id',\n    'timestamp',\n    'content_id',\n    'content_type_id',\n    'task_container_id',\n    'user_answer',\n    'answered_correctly',\n    'prior_question_elapsed_time',\n    'prior_question_had_explanation'\n]\n\nextended_attrs = ['user_id'] + ['seq_len', 'prev_seq_len', 'start', 'end'] + \\\n    [k for k in attrs if k not in ['user_id']] + \\\n    ['pred_time_mask', 'abs_pos', 'shifted_abs_pos', 'shifted_answered_correctly', 'shifted_user_answer'] + \\\n    ['lag_time']\n    \naggregated_attrs_less = [\n    'n_questions_answered',\n    'n_questions_answered_correctly',\n    'n_lectures_watched'\n]\nfor part_idx in range(2, 9):\n        \n    aggregated_attrs_less.append(f'part_{part_idx}_count')\n    aggregated_attrs_less.append(f'part_{part_idx}_correctness_count')\n\nfor correct_answer_idx in range(ANSWER_0_ID, ANSWER_3_ID + 1):\n        \n\taggregated_attrs_less.append(f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_count')\n\taggregated_attrs_less.append(f'correct_answer_{correct_answer_idx - ANSWER_0_ID}_answered_correctly_count')\n    \n# containing 'n_prev_seen' and 'n_prev_correctness'\naggregated_attrs = aggregated_attrs_less + ['n_prev_seen', 'n_prev_correctness']    \n    \n    \nattr_dtypes = {}\nfor k in extended_attrs:\n    \n    if k in ['row_id', 'user_id', 'timestamp', 'lag_time']:\n        attr_dtypes[k] = tf.int64\n    elif k in ['prior_question_elapsed_time']:\n        attr_dtypes[k] = DTYPE\n    else:\n        attr_dtypes[k] = tf.int32\n\nfor k in aggregated_attrs:\n    attr_dtypes[k] = tf.int32\n\n        \nclass User_Record:\n    \"\"\"\n    Sequences (per attribute) of records for a single user.\n    \"\"\"\n\n    def __init__(self, user_dt=None, user_id=None):\n        \"\"\"Exactly one argument should be `None`.\n        \n        Args:\n            user_dt: A `datatable.Frame` object for a single user.\n            user_id: int.\n        \"\"\"\n\n        ### assert (user_dt is not None and user_id is None) or (user_id is not None and user_dt is None)\n        \n        user_ids = user_dt[:, 'user_id'].to_list()[0] if user_dt is not None else []\n        \n#         # single value - Each record is for a single user.\n#         if user_dt is not None:\n            \n#             assert user_dt.nrows > 0\n#             assert len(set(user_ids)) == 1\n        \n        self.user_id = user_ids[0] if user_dt is not None else user_id\n    \n        self.row_id = user_dt[:, 'row_id'].to_list()[0] if user_dt is not None else []\n        self.timestamp = user_dt[:, 'timestamp'].to_list()[0] if user_dt is not None else []\n        self.content_id = user_dt[:, 'content_id'].to_list()[0] if user_dt is not None else []\n        self.content_type_id = user_dt[:, 'content_type_id'].to_list()[0] if user_dt is not None else []\n        self.task_container_id = user_dt[:, 'task_container_id'].to_list()[0] if user_dt is not None else []\n        self.user_answer = user_dt[:, 'user_answer'].to_list()[0] if user_dt is not None else []\n        self.answered_correctly = user_dt[:, 'answered_correctly'].to_list()[0] if user_dt is not None else []\n        self.prior_question_elapsed_time = user_dt[:, 'prior_question_elapsed_time'].to_list()[0] if user_dt is not None else []\n        self.prior_question_had_explanation = user_dt[:, 'prior_question_had_explanation'].to_list()[0] if user_dt is not None else []\n        \n        # This informatioin should be set externally\n        for aggregated_feature_idx, aggregated_feature_name in enumerate(aggregated_attrs):\n            setattr(self, aggregated_feature_name, [])\n        \n#         if DEBUG:\n#             # make sure the timestamp is always in order.\n#             assert self.timestamp == sorted(self.timestamp)\n            \n    def extend(self, other):\n        \"\"\"\n        Add the content of another recocrd `other` to the record `self`.\n        \"\"\"        \n        \n#         assert (self.user_id == other.user_id)\n        \n#         if DEBUG:\n#             # The `timestamp` should be in order while adding new entries to existing record.\n#             if len(self.timestamp) > 0:\n#                 # Don't allow a non-empty record to be extended by an empty record. \n#                 assert len(other.timestamp) > 0\n#                 assert self.timestamp[-1] <= other.timestamp[0]\n                \n        for k in attrs + aggregated_attrs:\n            if k != 'user_id':\n                getattr(self, k).extend(getattr(other, k))\n\n    def update_answer_results(self, prior_correctnesses, prior_answers, prior_user_info_for_post_update, question_history_at_training_end):\n        \"\"\"\n        Update the answers and their correctnesses in a record which was previously unknown in the last test batch.\n        \"\"\"\n        \n        \"\"\"\n        'user_id': user_id,\n        'pred_record': _pred_record,\n        'n_pred_time_steps': n_pred_time_steps,\n        'aggregated_info_to_update': aggregated_info_to_update        \n        \"\"\"\n        \n        n_pred_time_steps = prior_user_info_for_post_update['n_pred_time_steps']\n        aggregated_info_to_update = prior_user_info_for_post_update['aggregated_info_to_update']\n        \n#         assert n_pred_time_steps == len(prior_correctnesses)\n#         assert n_pred_time_steps == len(prior_answers)\n        \n#         if DEBUG:\n            \n#             # sanity check\n#             assert len(prior_correctnesses) == len(prior_answers)\n\n#             assert len(self.answered_correctly) >= len(prior_correctnesses)\n#             assert len(self.user_answer) >= len(prior_answers)\n\n#             # the places to be updated should contain only `MASK_TOKEN` (i.e. unknown results)\n#             assert set(self.answered_correctly[-len(prior_correctnesses):]) == {MASK_TOKEN}\n#             assert set(self.user_answer[-len(prior_answers):]) == {MASK_TOKEN}\n                \n        self.answered_correctly = self.answered_correctly[:-len(prior_correctnesses)] + prior_correctnesses\n        self.user_answer = self.user_answer[:-len(prior_answers)] + prior_answers\n        \n        # Update aggregated correction information\n        for idx in range(n_pred_time_steps):\n            \n            index_in_pred_record = - (n_pred_time_steps - idx)\n            \n            # ------------------------------------------------------------------------\n            # Update `n_pre_correctness` in `question_history_at_training_end`\n            \n            # At the current step, it is a question and answered correctly\n            if prior_correctnesses[idx] == 1:\n                \n                # Get the current question id\n                current_question_id_str = str(self.content_id[index_in_pred_record])\n                # update `n_pre_correctness` in `question_history_at_training_end`.\n                # The 'n_prev_seen' is already update in `question_history_at_training_end`.\n                # `self.n_prev_seen` and `self.n_prev_correctness` will be updated later.\n                question_history_at_training_end[str(self.user_id)][current_question_id_str][1] += 1\n            # ------------------------------------------------------------------------\n            # Update `aggregated_info_to_update`, which are used to update aggregated attributes in `self`.\n                \n            # No op to perform if we are at the 1st step in the current pred batch (of a particular user).\n            if idx > 0:\n                \n                for k in aggregated_attrs_less:\n                    aggregated_info_to_update[k][idx] = aggregated_info_to_update[k][idx-1]                 \n                \n                # We need to check the previous correctness\n                prev_correctness = prior_correctnesses[idx - 1]\n                # So  the previous step is a question and being answered correctly\n                if prev_correctness == 1:\n                    \n                    aggregated_info_to_update['n_questions_answered_correctly'][idx] = aggregated_info_to_update['n_questions_answered_correctly'][idx-1] + 1\n                    \n                    prev_part_key = aggregated_info_to_update['prev_part_key'][idx]\n\n                    prev_correct_answer_key = aggregated_info_to_update['prev_correct_answer_key'][idx]\n\n#                     assert prev_part_key is not None\n#                     assert prev_correct_answer_key is not None\n                    \n                    aggregated_info_to_update[prev_part_key][idx] = aggregated_info_to_update[prev_part_key][idx-1] + 1\n                    aggregated_info_to_update[prev_correct_answer_key][idx] = aggregated_info_to_update[prev_correct_answer_key][idx-1] + 1\n                    \n        # ------------------------------------------------------------------------                    \n                    \n        for k in aggregated_attrs_less:\n            if k == 'n_questions_answered_correctly' or 'correctness_count' in k or 'answered_correctly_count' in k:\n                setattr(self, k, getattr(self, k)[:-n_pred_time_steps] + aggregated_info_to_update[k])\n                    \n#         for k in aggregated_attrs:\n#             assert len(getattr(self, k)) == len(self)      \n                    \n        # ------------------------------------------------------------------------\n                    \n    def toJSON(self):\n        \n        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=False, indent=4)\n    \n    def __str__(self):\n        \n        return self.toJSON()\n    \n    def __len__(self):\n        \n        return len(self.row_id)\n\n\nclass Record_Buffer:\n    \"\"\"\n    A dictionary like buffer to store and manage (i.e. updating) records.\n    \"\"\"\n    \n    def __init__(self, record_dict=None):\n        \"\"\"\n        `record_dict`: A `dict` mapping user ids (`str`) to their records (`User_Record`).\n        \"\"\"\n        \n        if record_dict is None:\n            self.buffer = {}\n        else:\n            self.buffer = record_dict\n           \n    def __contains__(self, x):\n        \n        return x in self.buffer\n        \n    def __getitem__(self, x):\n        \n        if x not in self.buffer:\n            raise KeyError(str(x))\n        \n        return self.buffer[x]\n    \n    def __len__(self):\n        \n        return len(self.buffer)\n    \n    def __del__(self):\n        \n        del self.buffer\n                \n    def update(self, record):\n        \"\"\"\n        Add a record to the buffer. If its user_id already exists, find and update the existing record.\n        \"\"\"\n\n        if record.user_id not in self.buffer:\n            self.buffer[record.user_id] = User_Record(user_id=record.user_id)\n            \n        self.buffer[record.user_id].extend(record)\n\n    def update_answer_results(self, user_id, prior_correctnesses, prior_answers, prior_user_info_for_post_update, question_history_at_training_end):\n        \"\"\"\n        Update the answers and their correctnesses for a single user which was previously unknown in the last test batch.\n        \"\"\"\n\n#         assert user_id in self.buffer\n#         assert user_id == prior_user_info_for_post_update['user_id']\n\n        record = self.buffer[user_id]\n        record.update_answer_results(prior_correctnesses, prior_answers, prior_user_info_for_post_update, question_history_at_training_end)\n                \n    def toJSON(self):\n        \n        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=False, indent=4)\n    \n    def __str__(self):\n        \n        return self.toJSON()\n        \n        \ndef combine_user_record(record_1, record_2, check=False):\n    \"\"\"\n    Returns:\n        A new record that combines the two `User_Record` objects `record_1` and  `record_2`.\n        The arguments are not modified.\n    \"\"\"\n\n#     assert record_1.user_id == record_2.user_id\n\n#     if check:\n#         # This check should be only used for validation!\n#         # For submission, we don't know what the row ids would be!\n#         if len(record_1.row_id) > 0 and len(record_2.row_id) > 0:\n#             assert record_1.row_id[-1] == record_2.row_id[0] - 1   \n    \n    record = User_Record(user_id=record_1.user_id)\n\n    record.extend(record_1)\n    record.extend(record_2)\n\n    return record    \n    \n\nclass Pred_Manager:\n    \"\"\"\n    See `update()` for the description.\n    \"\"\"\n    \n    def __init__(\n        self, config, train_config, pred_iter_manager, train_dt, user_id_to_row_id_train,\n        question_history_at_training_end, single_question_history_at_training_end_optimized,\n        max_train_buffer_size=30000, probe=False\n    ):\n        \n        self.config = config\n        self.train_config = train_config\n        \n        self.pred_iter_manager = pred_iter_manager\n            \n        self.train_dt = train_dt\n        self.user_id_to_row_id_train = user_id_to_row_id_train\n\n        self.question_history_at_training_end = question_history_at_training_end\n        self.single_question_history_at_training_end_optimized = single_question_history_at_training_end_optimized\n\n        self.train_record_buffer = Record_Buffer()\n        self.pred_record_buffer = Record_Buffer()\n        \n        self.current_batch_row_ids = None        \n        self.current_batch_users = None\n        \n        # Used to avoid memory error - not sure if it is necessary.\n        self.max_train_buffer_size = max_train_buffer_size\n\n        self.users_in_pred = set()\n        self.common_users_in_train_and_pred = set()\n        self.new_users_in_pred = set()\n        \n        self.n_blocks_in_pred = 0\n        self.n_blocks_from_common_users_in_train_and_pred = 0\n        self.n_blocks_from_new_users_in_pred = 0\n        \n        self.avg_n_blocks_in_pred = 0\n        self.avg_n_blocks_from_common_users_in_train_and_pred = 0\n        self.avg_n_blocks_from_new_users_in_pred = 0\n        \n        # The user ids in the original (before splitted into training and validation datasets)\n        self.users_in_train = set(user_id_to_row_id_train.keys())\n        \n        # The user ids in the actual training dataset (i.e. excluded the validation dataset)\n        self.users_trained = set()\n        \n        for user_id in self.users_in_train:\n            \n            split_index = self.pred_iter_manager.get_split_index(user_id)\n            \n            if split_index == -1 or split_index > 0:\n                self.users_trained.add(user_id)\n        \n        self.probe = probe\n        \n        self.user_performance_hdf5_fp = h5py.File(user_performance_hdf5_path, \"r\")\n        self.user_performance_hdf5 = self.user_performance_hdf5_fp['user_performance']\n        \n    @classmethod\n    def create(cls, config, train_config, pred_iter_manager_class,\n               train_dt_path_or_obj,               \n               user_id_to_row_id_train_path,\n               unique_question_id_train_path,\n               unique_lecture_id_train_path,\n               question_history_at_training_end_path,\n               single_question_history_at_training_end_optimized_path,\n               valid_info_path=None,\n               max_train_buffer_size=30000,\n               probe=False,\n               debug=False\n        ):\n        \n        if pred_iter_manager_class == Valid_Iter_Manager:\n            \n            data_paths = [valid_info_path, train_dt_path_or_obj]\n            \n        elif pred_iter_manager_class == Test_Iter_Manager:\n            data_paths = []\n            \n        assert type(unique_question_id_train_path) == str and type(unique_lecture_id_train_path) == str\n        data_paths.extend([unique_question_id_train_path, unique_lecture_id_train_path])      \n        pred_iter_manager = pred_iter_manager_class.create(*data_paths)\n        \n        user_id_to_row_id_train = load_data(user_id_to_row_id_train_path)\n\n        \n        start_memeory = psutil.virtual_memory().available / 1024.0 / 1024.0\n        print(f'start_memeory before loading question_history_at_training_end.json: {start_memeory}')\n\n        question_history_at_training_end = load_data(question_history_at_training_end_path)\n\n        end_memeory = psutil.virtual_memory().available / 1024.0 / 1024.0\n        print(f'end_memeory after loading question_history_at_training_end.json: {end_memeory}')\n\n        used_memory = start_memeory - end_memeory\n        print(f'used_memory for loading question_history_at_training_end.json: {used_memory}')\n        print('----------------------------------------')\n        \n        start_memeory = psutil.virtual_memory().available / 1024.0 / 1024.0\n        print(f'start_memeory before loading single_question_history_at_training_end_optimized.json: {start_memeory}')\n        \n        single_question_history_at_training_end_optimized = load_data(single_question_history_at_training_end_optimized_path)\n        \n        end_memeory = psutil.virtual_memory().available / 1024.0 / 1024.0\n        print(f'end_memeory after loading single_question_history_at_training_end_optimized.json: {end_memeory}')\n\n        used_memory = start_memeory - end_memeory\n        print(f'used_memory for loading single_question_history_at_training_end_optimized.json: {used_memory}')\n        print('----------------------------------------')\n        \n        if pred_iter_manager_class == Valid_Iter_Manager:\n            train_dt = pred_iter_manager.train_dt\n        elif type(train_dt_path_or_obj) != str:\n            train_dt = train_dt_path_or_obj\n        else:\n            train_dt = load_data(train_dt_path_or_obj)\n\n        args = [\n                config, train_config, pred_iter_manager, train_dt, user_id_to_row_id_train,\n                question_history_at_training_end, single_question_history_at_training_end_optimized,\n                max_train_buffer_size, probe\n        ]\n        \n        return cls(*args)\n        \n#     def reset_train_record_buffer(self):\n#         \"\"\"\n#         \"\"\"\n        \n#         del self.train_record_buffer\n#         self.train_record_buffer = Record_Buffer()\n        \n    \n    def reset_train_record_buffer(self):\n        \"\"\"\n        \"\"\"\n        \n        user_ids = list(self.train_record_buffer.buffer.keys())\n        for user_id in user_ids:\n            del self.train_record_buffer.buffer[user_id]\n        \n        del self.train_record_buffer\n\n        self.train_record_buffer = Record_Buffer()    \n    \n    def update_batch_users(self, user_ids):\n        \"\"\"\n        Store the user ids in the current test batch.\n        \"\"\"\n        \n        self.current_batch_users = user_ids\n\n    def update_answer_results(self, prev_batch_users, prior_group_answers_correct, prior_group_responses, prior_info_for_post_update):\n        \"\"\"\n        When we get a new test batch, we also get the `prior_group_answers_correct` and `prior_group_responses`.\n        We use these information to update the `answered_correctly` and `user_anser` fields of the records of\n        the users in the previous test batch.\n        \"\"\"\n        \n#         if DEBUG:\n        \n#             # sanity check: try to make sure the answer results are for `prev_batch_users` by verifying their lengths.\n#             assert len(prev_batch_users) == len(prior_group_answers_correct)\n#             assert len(prior_group_answers_correct) == len(prior_group_responses)\n        \n        d1 = defaultdict(list)\n        d2 = defaultdict(list)\n        \n        for user_id, prior_ans_correct, prior_ans in zip(prev_batch_users, prior_group_answers_correct, prior_group_responses):\n            d1[user_id].append(prior_ans_correct)\n            d2[user_id].append(prior_ans)\n            \n        for user_id in d1:\n                        \n            self.pred_record_buffer.update_answer_results(\n                user_id,\n                d1[user_id], d2[user_id],\n                prior_info_for_post_update[user_id],\n                self.question_history_at_training_end\n            )\n                    \n    def update_aggregated_info_buffer(self, user_id, pred_record, record_to_watch, index_to_watch, n_pred_time_steps, pred_time_step_index, buffer):\n                \n        user_id_str = str(user_id)\n        \n        # ====================================================================================================\n        # This should be done for all `current` positions - dealing with `n_prev_seen` and `n_prev_correctness`\n                \n        current_index_in_pred_record = - (n_pred_time_steps - pred_time_step_index)\n        \n        n_prev_seen, n_prev_correctness = None, None\n        \n        # OP need to be done only if the current position is a question\n        if pred_record.content_type_id[current_index_in_pred_record] == 0:\n            \n            # current is a question\n            current_question_id = pred_record.content_id[current_index_in_pred_record]\n            current_question_id_str = str(current_question_id)\n            \n            # -----------------------------------------------------------------------------------------------\n            # During the prediction time, we make sure all the questions' info about `n_prev_seen` and `n_prev_correctness` are\n            # added to `question_history_at_training_end`.\n            if user_id_str not in self.question_history_at_training_end:\n                self.question_history_at_training_end[user_id_str] = dict()\n            if current_question_id_str not in self.question_history_at_training_end[user_id_str]:\n                # Here we use `[None, None]` as a placeholder. The `None` values should be replaced by some integers\n                # after the computation is done below.\n                # The idea is to make sure we have the necessary keys - so we don't have `KeyError` error.\n                self.question_history_at_training_end[user_id_str][current_question_id_str] = [None, None]\n            else:\n                # We have the values for the current question id, that we can put into the buffer directly.\n                n_prev_seen, n_prev_correctness = self.question_history_at_training_end[user_id_str][current_question_id_str]\n                    \n            # We can't find information in `question_history_at_training_end`,\n            # so check with `single_question_history_at_training_end_optimized`.\n            if n_prev_seen is None:\n                \n                ### assert n_prev_correctness is None\n                \n                if user_id_str in self.single_question_history_at_training_end_optimized:\n                    \n                    # Change the lists to sets on the fly.\n                    # We don't do this op during the loading, because it takes too many memory.\n                    # We only do this op for the users with question seen during the prediction.\n                    if type(self.single_question_history_at_training_end_optimized[user_id_str]['correct']) == list:\n                        self.single_question_history_at_training_end_optimized[user_id_str]['correct'] = set(self.single_question_history_at_training_end_optimized[user_id_str]['correct'])\n                    if type(self.single_question_history_at_training_end_optimized[user_id_str]['incorrect']) == list:\n                        self.single_question_history_at_training_end_optimized[user_id_str]['incorrect'] = set(self.single_question_history_at_training_end_optimized[user_id_str]['incorrect'])\n                    \n                    # We remove the information from `single_question_history_at_training_end_optimized`,\n                    # because we will put the information to `question_history_at_training_end`.\n                    if current_question_id_str in self.single_question_history_at_training_end_optimized[user_id_str]['correct']:\n                        n_prev_seen, n_prev_correctness = [1, 1]\n                        self.single_question_history_at_training_end_optimized[user_id_str]['correct'].remove(current_question_id_str)\n                    elif current_question_id_str in self.single_question_history_at_training_end_optimized[user_id_str]['incorrect']:\n                        n_prev_seen, n_prev_correctness = [1, 0]\n                        self.single_question_history_at_training_end_optimized[user_id_str]['incorrect'].remove(current_question_id_str)\n                        \n            if n_prev_seen is None:\n                \n                ### assert n_prev_correctness is None\n                n_prev_seen, n_prev_correctness = [0, 0]\n                \n            # Unlike `buffer`, `question_history_at_training_end` should be updated only if the current position is a question.\n            # Update `question_history_at_training_end` by combining `n_prev_seen` and the current question information.\n            self.question_history_at_training_end[user_id_str][current_question_id_str][0] = n_prev_seen + 1\n            # We don't know yet the correctness. This should be update after.\n            # The current value is not the correct value (it is the value in the state one time step before) \n            self.question_history_at_training_end[user_id_str][current_question_id_str][1] =  n_prev_correctness\n            \n            # -----------------------------------------------------------------------------------------------\n            \n        # The current position is a lecture. The vaules should be all `0`.\n        if n_prev_seen is None:\n\n            ### assert n_prev_correctness is None\n            n_prev_seen, n_prev_correctness = [0, 0]            \n            \n        # This is the value to be used.\n        # We need to use `append` because the current position doesn't exist in these two lists yet.            \n        # Be careful, unlike the update for `question_history_at_training_end` above with `n_prev_seen + 1`,\n        # we need to use `n_prev_seen` here!\n        buffer['n_prev_seen'].append(n_prev_seen)\n        buffer['n_prev_correctness'].append(n_prev_correctness)            \n            \n        # ====================================================================================================\n        # For attributes in `aggregated_attrs` other than `n_prev_seen` and `n_prev_correctness`.\n        # We need to look the values in the previous step in the history.\n            \n        # ------------------------------------------------------------\n        # No previous history, we still need to assign value of `0` and put them to the lists.\n        \n        if record_to_watch is None:\n            # `n_prev_seen` and `n_prev_correctness` is done already.\n            for k in aggregated_attrs_less:             \n                buffer[k].append(0)\n            for k in ['prev_part_key', 'prev_correct_answer_key']:\n                # We assign `None` value.\n                # In post-update, these values shoudn't be used, otherwise some logic is wrong, and we need to debug.\n                buffer[k].append(None)\n                            \n            return\n        # ------------------------------------------------------------\n        # There is previous hisotry, either in training, or in previous predictions.\n        \n        # Here we just copy the values (of aggregated-attrs, other than `n_prev_seen` and `n_prev_correctness`)\n        # from the previous step in the history. We need to update these values later by watching the \n        # values of non-aggregated-attrs from the previous step in the history.\n        \n        # i.e. we are at the first interaction (of a particular user) in the current pred batch.\n        if pred_time_step_index == 0:\n            # `n_prev_seen` and `n_prev_correctness` is done already.\n            for k in aggregated_attrs_less:\n                # Here, we can use `-1`\n                # If `record_to_watch` is `train_record`, `-1` is of course fine.\n                # However, if `record_to_watch` is `pred_record` - for attributes in `aggregated_attrs`, the lists haven't been extended and updated,\n                # therefore, `-1` allows us to get the previous value.\n                buffer[k].append(getattr(record_to_watch, k)[-1])\n        # i.e. we are not at the first interaction (of a particular user) in the current pred batch.\n        # `buffer[k]` should already have values, and we are safe to use `-1` as index.\n        # (for `n_prev_seen` and `n_prev_correctness`, we sholud no longer operate on them), because they must have been done above.\n        else:\n            for k in aggregated_attrs_less:\n                # Use the previous values in `buffer[k]` - and to update later.\n                buffer[k].append(buffer[k][-1])\n                \n        # Just placeholders - if the previous step is a question\n        for k in ['prev_part_key', 'prev_correct_answer_key']:\n            buffer[k].append(None)\n            \n        # ------------------------------------------------------------\n        # We use `index_to_watch` to watch the previous values (of non aggregated-attributes) in `record_to_watch`.\n        # We shouldn't watch the `aggregated-attributes` by using `index_to_watch`, because they don't exist yet!\n        \n        content_id = record_to_watch.content_id[index_to_watch]\n        content_type_id = record_to_watch.content_type_id[index_to_watch]\n    \n        # calculated only for question\n        if content_type_id == 0:\n\n            content_input_id = question_id_to_input_id_dict[content_id]\n            # This is in [2, 8]\n            part_id = c_inputs_ids_to_part_dict[content_input_id] + 1\n            # This is in [0, 3]\n            correct_answer = c_inputs_ids_to_correct_answer_id_dict[content_input_id] - ANSWER_0_ID\n\n            # need to be a string\n            question_id = str(content_id)\n            \n            answered_correctly = (record_to_watch.answered_correctly[index_to_watch] == 1)\n\n            # We can use `-1` as index here, because we already extend the list with the value from the previous step.\n            # If the previous step is not a question, the values remain `None`.\n            buffer['prev_part_key'][-1] = f'part_{part_id}_correctness_count'\n            buffer['prev_correct_answer_key'][-1] = f'correct_answer_{correct_answer}_answered_correctly_count'\n\n        is_lecture = int(content_type_id == 1)\n\n        # ------------------------------------------------------------\n        # We only update the involved aggregated-attributes\n\n        buffer['n_lectures_watched'][-1] += is_lecture        \n\n        # i.e. the prev step is a question\n        if is_lecture == 0:\n            \n            buffer['n_questions_answered'][-1] += 1\n\n            k = f'part_{part_id}_count'\n            buffer[k][-1] += 1\n            k = f'correct_answer_{correct_answer}_count'\n            buffer[k][-1] += 1\n            \n            # Update the correctness if `pred_time_step_index = 0` and the previous correctness is `True`.\n            # This should only occur for `pred_time_step_index = 0`.\n            # Because the correctness information in the current pred batch is not available, and they have `MASK_TOKEN` as values.\n            # `pred_record` will have the values at the first (current) prediction time question's place.\n            if answered_correctly is True:\n                \n                ### assert pred_time_step_index == 0\n                \n                for k in [\n                    f'n_questions_answered_correctly',\n                    f'part_{part_id}_correctness_count',\n                    f'correct_answer_{correct_answer}_answered_correctly_count'\n                ]:\n                \n                    buffer[k][-1] += 1\n                    \n    def get_aggregated_info_to_update(self, user_id, train_record, pred_record, n_pred_time_steps):\n            \n        if len(pred_record) > n_pred_time_steps:\n            # i.e. we have already some prediction time history for that user\n            record_to_watch = pred_record\n            index_to_watch = - (n_pred_time_steps + 1)\n        elif len(train_record) > 0:\n            # i.e. the first pred batch, but the training history is not empty\n            record_to_watch = train_record\n            index_to_watch = -1\n        else:\n            # i.e. the first pred batch, but no history used for traininig\n            record_to_watch = None\n            index_to_watch = None\n        \n        aggregated_info_buffer = {\n            # `aggregated_attrs` contains `n_prev_seen` and `n_prev_correctness`\n            # we also add `prev_part_key` and `prev_correct_answer_key` so we don't need to reculate the part and correct answer\n            # in post-updating.\n            k: [] for k in aggregated_attrs + ['prev_part_key', 'prev_correct_answer_key']\n        }\n        \n        for pred_time_step_index in range(n_pred_time_steps):\n            if pred_time_step_index >= 1:\n                record_to_watch = pred_record\n                index_to_watch = - (n_pred_time_steps - pred_time_step_index) - 1\n            self.update_aggregated_info_buffer(user_id, pred_record, record_to_watch, index_to_watch, n_pred_time_steps, pred_time_step_index, aggregated_info_buffer)\n        \n        return aggregated_info_buffer\n        \n    def update(self, pred_dt, prior_group_answers_correct, prior_group_responses, prior_info_for_post_update):\n        \"\"\"\n        For a prediction batch `pred_df` (`pandas.DataFrame`) given by `env.iter_test`, this method performs:\n            1. update the `answered_correctly` and `user_anser` information in the previous test batch\n            2. update the user records (only in the test time) by appending the information in the current batch\n            3. get the user records in the training time\n            4. combine the user records in the training time and test time - so we have a full history and the current batch to predict\n        \"\"\"\n            \n        user_ids = pred_dt['user_id'].to_list()[0]\n        prev_batch_users = self.current_batch_users\n        \n        if prev_batch_users is not None:\n            ### assert prior_info_for_post_update is not None\n            self.update_answer_results(prev_batch_users, prior_group_answers_correct, prior_group_responses, prior_info_for_post_update)\n        self.update_batch_users(user_ids)\n        \n        if DEBUG:\n            \n            row_ids = pred_dt['row_id'].to_list()[0]\n            prev_batch_row_ids = self.current_batch_row_ids\n            \n            if isinstance(self.pred_iter_manager, Test_Iter_Manager):\n                debug_pred_batch_row_ids(row_ids, prev_batch_row_ids)\n        \n            self.current_batch_row_ids = row_ids\n        \n        # To `User_Record`. Here, each record contains exactly one user interaction.\n        # But the same user might appear several times in the list.\n        record_batch = [User_Record(user_dt=pred_dt[idx, :]) for idx in range(pred_dt.nrows)]\n        \n        if DEBUG:\n            debug_record_batch(record_batch)\n                        \n        # update test buffer - add info about the new test batch\n        # at this moment, for each attr in `aggregated_attrs`, each record is an empty list,\n        # and the corresponding records in the buffer, for `aggregated_attrs`, are not updated\n        for record in record_batch:\n            self.pred_record_buffer.update(record)\n        \n        # ----------------------------------------------------\n        \n        self.users_in_pred.update(set(user_ids))\n\n        n_new_users_in_batch = 0\n        for x in set(user_ids):\n            if x in self.users_trained:\n                self.common_users_in_train_and_pred.add(x)\n            else:\n                self.new_users_in_pred.add(x)\n                n_new_users_in_batch += 1\n            \n        # ----------------------------------------------------            \n            \n        nb_common_users = len(self.common_users_in_train_and_pred)\n        nb_new_users = len(self.new_users_in_pred)        \n        \n        common_users_in_batch = set(user_ids).intersection(self.users_trained)\n        new_users_in_batch = set(user_ids).difference(self.users_trained)\n        \n        self.n_blocks_in_pred += len(set(user_ids))\n        self.n_blocks_from_common_users_in_train_and_pred += len(common_users_in_batch)\n        self.n_blocks_from_new_users_in_pred += len(new_users_in_batch)\n                \n        self.avg_n_blocks_in_pred = self.n_blocks_in_pred / len(self.users_in_pred)\n        \n        if nb_common_users > 0:\n            self.avg_n_blocks_from_common_users_in_train_and_pred = self.n_blocks_from_common_users_in_train_and_pred / nb_common_users       \n        \n        if nb_new_users > 0:\n            self.avg_n_blocks_from_new_users_in_pred = self.n_blocks_from_new_users_in_pred / nb_new_users\n      \n        # ----------------------------------------------------    \n    \n#         if self.probe:\n        \n#             # ----------------------------------------------------\n#             # Probe: The number of users in the test dataset.\n\n#             # At most N_USERS users in the test dataset.       \n#             # N_USERS = 20992  # False\n#             N_USERS = 21056  # True\n\n#             assert len(self.pred_record_buffer) <= N_USERS        \n\n#             # ----------------------------------------------------\n#             # Probe: The number of new users in the test dataset.\n\n#             # N_NEW_USERS = 5904 # False\n#             N_NEW_USERS = 5936 # True\n\n#             assert nb_new_users <= N_NEW_USERS\n\n        # ----------------------------------------------------\n        \n        # get the updated record from test buffer for each user in `user_ids`.\n        # The same user might appear several times in the list, however, they get the same user history (test time) sequence.\n        ### _pred_record_batch = [self.pred_record_buffer[x] for x in user_ids]\n        \n        # get the updated record from test buffer for each user in `user_ids` without duplication.\n        # The same user can only appear at most 1 time in the list.\n        _pred_record_batch = []\n        n_pred_time_steps_dict = dict()\n        \n        for x in user_ids:\n        \n            if x not in n_pred_time_steps_dict:\n                _pred_record_batch.append(self.pred_record_buffer[x])\n                n_pred_time_steps_dict[x] = 1\n            else:\n                n_pred_time_steps_dict[x] += 1\n                continue\n                        \n        user_ids_without_duplication = [x.user_id for x in _pred_record_batch]        \n        n_pred_time_steps_batch = [n_pred_time_steps_dict[x.user_id] for x in _pred_record_batch]\n                \n        ### assert 0 <= min(n_pred_time_steps_batch) \n        ### assert max(n_pred_time_steps_batch) <= self.train_config.window_size\n        \n        # get the record from train_dt or train record buffer for each user in `user_ids`\n        train_record_batch = [self.get_training_record(x) for x in user_ids_without_duplication]\n    \n        # obtain the full history (in training time + the previous batches in test time) and the current batch to predict\n        check = isinstance(self.pred_iter_manager, Valid_Iter_Manager)    \n    \n        info_for_post_update = {}\n        pred_record_batch = []\n        for user_id, train_record, _pred_record, n_pred_time_steps in zip(\n            user_ids_without_duplication,\n            train_record_batch,\n            _pred_record_batch,\n            n_pred_time_steps_batch\n        ):\n                   \n            # A dictionary of lists, key are in `aggregated_attrs` and values are all of length `n_pred_time_steps`\n            aggregated_info_to_update = self.get_aggregated_info_to_update(user_id, train_record, _pred_record, n_pred_time_steps)\n                            \n            pred_record = combine_user_record(train_record, _pred_record, check=check)\n            pred_record_batch.append(pred_record)\n            \n            info_for_post_update[user_id] = {\n                'user_id': user_id,\n                'pred_record': _pred_record,\n                'n_pred_time_steps': n_pred_time_steps,\n                'aggregated_info_to_update': aggregated_info_to_update\n            }            \n            \n            for k in aggregated_attrs:\n                \n                ### assert len(getattr(_pred_record, k)) == len(_pred_record) - n_pred_time_steps\n                \n                # `pred_record` is used for model prediction\n                if k in ['n_lectures_watched', 'n_prev_seen', 'n_prev_correctness']:\n                    getattr(pred_record, k).extend(aggregated_info_to_update[k])\n                else:\n                    # only use the information before the current pred batch\n                    getattr(pred_record, k).extend([aggregated_info_to_update[k][0]] * n_pred_time_steps)\n                    \n                # We need to update those correct information later!!!\n                getattr(_pred_record, k).extend(aggregated_info_to_update[k])\n                                     \n#             for k in aggregated_attrs:\n\n#                 assert len(getattr(_pred_record, k)) == len(_pred_record)   \n#                 assert len(getattr(pred_record, k)) == len(pred_record)\n\n#         if len(self.train_record_buffer) >= self.max_train_buffer_size:\n#             self.reset_train_record_buffer()\n            \n        if len(self.train_record_buffer) >= 3000:\n            self.reset_train_record_buffer()\n            gc.collect()\n    \n        return pred_record_batch, n_pred_time_steps_batch, info_for_post_update\n\n    def get_training_record(self, user_id):\n        \"\"\"\n        Get the training time history of the user with `user_id`.\n        \"\"\"\n               \n        # Already cached the result in the buffer\n        if user_id in self.train_record_buffer:\n            return self.train_record_buffer[user_id]\n\n        # unseen user in the test time\n        if user_id not in self.user_id_to_row_id_train:\n            \n            assert isinstance(self.pred_iter_manager, Test_Iter_Manager)\n        \n        else:\n            \n            split_index = self.pred_iter_manager.get_split_index(user_id)        \n            \n            # user appeared in the original (unsplitted) training dataset\n            \n            user_dt = get_user_dt(user_id, self.train_dt, self.user_id_to_row_id_train, split_index)\n                    \n            # During the test time, if an user seen in training time, its whole training history should be used!\n            # During the validatioin time, it is possible not to use any of its (original) training examples (i.e. new user for the validation).\n            if isinstance(self.pred_iter_manager, Test_Iter_Manager):\n                assert user_dt.nrows > 0\n            \n            if user_dt.nrows > 0:\n                \n                convert_dt(user_dt, pred=False)\n\n                user_record = User_Record(user_dt=user_dt)\n                self.train_record_buffer.update(user_record)\n                \n                # ----------------------------------------\n                # From now, we should use `self.train_record_buffer[user_id]` (the one stored in the buffer), not `user_record`!\n                \n                # set the training user performance information\n                row_ids = self.train_record_buffer[user_id].row_id\n                ### assert len(row_ids) > 0\n                row_start, row_end = row_ids[0], row_ids[-1]\n                seq_len = row_end - row_start + 1\n                ### assert len(row_ids) == seq_len\n                \n                # array of shape [row_end - row_start + 1, len(aggregated_attrs)]\n                user_performance = self.user_performance_hdf5[row_start:row_end + 1, :]\n                ### assert user_performance.shape == (seq_len, len(aggregated_attrs))\n\n                for feature_idx, k in enumerate(aggregated_attrs):\n                    \n                    feature_performance = user_performance[:, feature_idx].tolist()\n                    getattr(self.train_record_buffer[user_id], k).extend(feature_performance)\n                    ### assert len(getattr(self.train_record_buffer[user_id], k)) == seq_len\n                \n                # ----------------------------------------\n                \n                return self.train_record_buffer[user_id]\n\n        # Empty training (new users in prediction time)\n        user_record = User_Record(user_id=user_id)\n        self.train_record_buffer.update(user_record)\n        \n        return self.train_record_buffer[user_id]\n\n    def convert_to_pred_batch(self, pred_record_batch, n_pred_time_steps_batch):\n\n        batch_record = self.convert_record_batch(pred_record_batch, n_pred_time_steps_batch)\n  \n        # A single dict of `tf.RaggedTensor` objects                                                                \n        batch = Pred_Manager.batch_record_to_tensor(batch_record)\n\n        training = tf.constant(0, dtype=tf.int32)\n        \n        # A single dict of `tf.Tensor` objects\n        ###pred_batch = add_input_ids_and_targets(batch, training, generative=self.config.generative, use_abs_pos=self.config.use_abs_pos)\n                        \n        ###return pred_batch\n        return batch\n            \n    def record_to_dict(self, record, n_pred_time_steps):\n                            \n        record_as_dict = {}\n        \n        orig_seq_len = len(record.row_id)\n        used_seq_len = orig_seq_len\n        \n        pad_seq = []\n        if orig_seq_len < self.train_config.window_size:\n            pad_seq = [PAD_TOKEN] * (self.train_config.window_size - orig_seq_len)\n            \n        if orig_seq_len > self.train_config.window_size:\n            used_seq_len = self.train_config.window_size\n        \n        record_as_dict['user_id'] = record.user_id\n    \n        # adding information\n        record_as_dict['seq_len'] = used_seq_len\n        record_as_dict['prev_seq_len'] = orig_seq_len\n        record_as_dict['start'] = max(orig_seq_len - self.train_config.window_size, 0)\n        record_as_dict['end'] = orig_seq_len - 1\n        \n        # Be careful: we need this to compute `shifted_answered_correctly` correctly later.\n        # Otherwise, we will get `START_TOKEN` before `PAD_TOKEN`.\n        answered_correctly = getattr(record, 'answered_correctly')\n        user_answer = getattr(record, 'user_answer')\n\n        for k in attrs + aggregated_attrs:\n            if k not in ['user_id']:\n                record_as_dict[k] = (pad_seq + getattr(record, k))[- self.train_config.window_size:]\n#                 assert len(record_as_dict[k]) == self.train_config.window_size\n                    \n        # We need to use `START_TOKEN` rather than `PAD_TOKEN` here.\n        shifted_answered_correctly = [START_TOKEN] + answered_correctly[:-1]\n        record_as_dict['shifted_answered_correctly'] = (pad_seq + shifted_answered_correctly)[- self.train_config.window_size:]\n#         assert len(record_as_dict['shifted_answered_correctly']) == self.train_config.window_size\n\n        shifted_user_answer = [START_TOKEN] + user_answer[:-1]\n        record_as_dict['shifted_user_answer'] = (pad_seq + shifted_user_answer)[- self.train_config.window_size:]\n#         assert len(record_as_dict['shifted_user_answer']) == self.train_config.window_size\n\n#         assert self.train_config.window_size >= n_pred_time_steps\n                \n        pred_time_mask = [0] * orig_seq_len\n        pred_time_mask = (pad_seq + pred_time_mask)[- self.train_config.window_size:]\n        \n        # Be careful: `n_pred_time_steps` might be `0`, and using `pred_time_mask[:- n_pred_time_steps]` gives wrong results.\n        pred_time_mask = pred_time_mask[:self.train_config.window_size - n_pred_time_steps] + [1] * n_pred_time_steps\n        \n        record_as_dict['pred_time_mask'] = pred_time_mask\n#         assert len(record_as_dict['pred_time_mask']) == self.train_config.window_size\n\n        abs_pos = list(range(orig_seq_len))\n        shifted_abs_pos = [START_TOKEN] + abs_pos[:-1]\n\n        record_as_dict['abs_pos'] = (pad_seq + abs_pos)[- self.train_config.window_size:]\n        record_as_dict['shifted_abs_pos'] = (pad_seq + shifted_abs_pos)[- self.train_config.window_size:]\n        \n#         assert len(record_as_dict['abs_pos']) == self.train_config.window_size\n#         assert len(record_as_dict['shifted_abs_pos']) == self.train_config.window_size\n        \n        # --------------------------------------------------------------------------------\n        # compute `lag time`\n        \n        timestamps = getattr(record, 'timestamp')\n        task_container_ids = getattr(record, 'task_container_id')\n        nb_actual_elts = min(self.train_config.window_size, len(timestamps))\n        \n        # try to get the (ending) subsequence of length `window_size + `MAX_PRED_TIME_QUESTION_BUNDLE_LEN`, but it might be shorter\n        _start_index = max(0, len(timestamps) - (self.train_config.window_size + MAX_PRED_TIME_QUESTION_BUNDLE_LEN))\n        # truncated: length = len(...) - `_start_index`\n        timestamps = timestamps[_start_index:]\n        task_container_ids = task_container_ids[_start_index:]\n        \n        # add leading `0`s if the lenght is not enough\n        nb_elt_to_add = (self.train_config.window_size + MAX_PRED_TIME_QUESTION_BUNDLE_LEN) - len(timestamps)\n        if nb_elt_to_add > 0:\n            timestamps = [0] * nb_elt_to_add + timestamps\n            task_container_ids = [PAD_TOKEN] * nb_elt_to_add + task_container_ids\n\n#         assert len(timestamps) == self.train_config.window_size + MAX_PRED_TIME_QUESTION_BUNDLE_LEN\n#         assert len(task_container_ids) == self.train_config.window_size + MAX_PRED_TIME_QUESTION_BUNDLE_LEN\n\n        # loop\n        lag_time = []\n        current_lag_time = 0\n        for _t, _prev_t, _bundle, _prev_bundle in zip(timestamps[1:], timestamps[:-1], task_container_ids[1:], task_container_ids[:-1]):\n            _diff = _t - _prev_t\n            if _diff == 0 and _bundle == _prev_bundle:\n                # In the same bundle\n                _diff = current_lag_time\n            elif _diff == 0 and _bundle != _prev_bundle:\n                # problematic data point - but we can use `0` and post-process it below\n                pass\n            else:\n                # update the latest lag time\n                # assert _diff > 0\n                current_lag_time = _diff\n            lag_time.append(_diff)\n            \n#         assert len(lag_time) == self.train_config.window_size + MAX_PRED_TIME_QUESTION_BUNDLE_LEN - 1\n        \n        # remove the leading part\n        lag_time = lag_time[-nb_actual_elts:]\n        timestamps = timestamps[1:][-nb_actual_elts:]\n        \n#         assert len(lag_time) == nb_actual_elts\n\n        # set to 1 second if `lag_time == 0` but they are not in the same bunlde\n        _mask = tf.cast(tf.math.logical_and(tf.constant(lag_time, dtype=tf.int64) == 0, tf.constant(timestamps, dtype=tf.int64) > 0), dtype=tf.int64) \n        lag_time = tf.constant(lag_time, dtype=tf.int64) * (1 - _mask) + tf.constant(1000, dtype=tf.int64) * _mask          \n        lag_time = lag_time.numpy().tolist()\n        \n        lag_time = (pad_seq + lag_time)[- self.train_config.window_size:]\n#         assert len(lag_time) == self.train_config.window_size\n        \n        record_as_dict['lag_time'] = lag_time\n        \n        # --------------------------------------------------------------------------------\n\n        return record_as_dict\n    \n    def convert_record_batch(self, record_batch, n_pred_time_steps_batch):\n        \"\"\"A format change for a list of records.\n\n        Args:\n            record_batch: A list of `User_Record`.\n\n        Returns:\n            batch_record: A `dict` with keys in `attrs`. For each key `k`, the value\n            is a list of sequences [x.k for x in record_batch]\n        \"\"\"\n\n        # time consuming\n        records_as_dicts = [Pred_Manager.record_to_dict(self, record, n_pred_time_steps) for record, n_pred_time_steps in zip(record_batch, n_pred_time_steps_batch)]\n        batch_record = {k: [x[k] for x in records_as_dicts] for k in extended_attrs + aggregated_attrs}\n\n        return batch_record\n\n\n    @staticmethod\n    def batch_record_to_tensor(batch_record):\n        \"\"\"Convert a `batch_record` obtained from `convert_record_batch` to `tf.Tensor`.\n        \"\"\"\n        \n        t = {\n            k: tf.constant(batch_record[k], dtype=attr_dtypes[k]) for k in extended_attrs + aggregated_attrs\n        }    \n        \n#         assert t['timestamp'].dtype == tf.int64\n        \n        return t","execution_count":null,"outputs":[]},{"metadata":{"id":"UG4riKib11wR"},"cell_type":"markdown","source":"### Debug functions"},{"metadata":{"id":"XZdATfXD11wU","trusted":true},"cell_type":"code","source":"def debug_pred_batch_row_ids(row_ids, prev_batch_row_ids):\n    \"\"\"Verify the properties of row ids in a single and across pred batch(es).\n    \n    Args:\n        row_ids: The row ids in a batch during the pred time given by `env.iter_test()`\n        prev_batch_row_ids: The row ids in the batch just before the batch of `row_ids` during the test time given by `env.iter_test()`\n    \"\"\"\n\n    # sanity check\n    # row ids must be distinct\n    assert len(set(row_ids)) == len(row_ids)\n    \n    # row ids must be in sorted order in a batch\n    assert row_ids == sorted(row_ids)\n    \n    # row ids must be in sorted order across all batch during the pred time.\n    if prev_batch_row_ids is not None:\n        assert row_ids[0] > prev_batch_row_ids[-1]\n\n\ndef debug_record_batch(record_batch):\n    \"\"\"Verify the properties of the question bundle for a single user in a pred batch.\n    \n    Args:\n        record_batch: A list. Each element (`User_Record`) should contain only one record in a single timestamp.\n    \"\"\"\n\n    _tmp = defaultdict(list)\n    for record in record_batch:\n        # each record here contains only one entry.\n        assert len(record.row_id) == 1\n        _tmp[record.user_id].append(record)\n\n    for user_id, records in _tmp.items():\n\n        task_container_ids_for_questions = [x.task_container_id[0] for x in records if x.content_type_id[0] == 0]\n\n        # If there is any question for a user\n        if len(task_container_ids_for_questions) > 0:\n\n            # There must be exactly one question bundle.\n            # This is `True`.\n            assert len(set(task_container_ids_for_questions)) == 1\n\n            row_ids_for_questions = [x.row_id[0] for x in records if x.content_type_id[0] == 0]\n            \n            # This is `False`: the question bundle must be in a consecutive block with continuous row ids.\n            # assert row_ids_for_questions == list(range(row_ids_for_questions[0], row_ids_for_questions[-1] + 1))\n\n            records_in_between = [x for x in records if row_ids_for_questions[0] <= x.row_id[0] <= row_ids_for_questions[-1]]\n            row_ids_in_between = [x.row_id[0] for x in records_in_between]\n            \n            # This is `True`: the question bundle must be in a consecutive block (but the row ids may jump).\n            assert row_ids_for_questions == row_ids_in_between\n\n            # The question bundle must be at the end of the sequence (for a single user) in a test batch.\n            assert row_ids_for_questions == [x.row_id[0] for x in records][-len(row_ids_for_questions):]","execution_count":null,"outputs":[]},{"metadata":{"id":"UaXqg-O011wn"},"cell_type":"markdown","source":"## Prediction routine"},{"metadata":{"trusted":true},"cell_type":"code","source":"max_buffer_size = 20\nbuffer = []\n\n\ndef run_pred(pred_manager, predictor, max_steps=None, ckpt_no=None):\n    \n    batch_size = None\n    seq_len = pred_manager.train_config.window_size\n\n    encoder_decoder = tf.constant(pred_manager.config.model_type=='ed', dtype=tf.int32)\n    training = tf.constant(0, dtype=tf.int32)\n    generative = tf.constant(pred_manager.config.generative, dtype=tf.int32)\n    allow_bundle_atten = tf.constant(pred_manager.config.allow_bundle_atten, dtype=tf.int32)\n    \n    # Only used for validation\n    valid_user_ids = []\n    valid_row_ids = []\n    valid_targets = []\n    valid_preds = []\n\n    if pred_manager.config.use_softmax:\n        loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE,\n            name='sparse_categorical_crossentropy_for_correctness'\n        )\n    else:\n        loss_obj = tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, label_smoothing=0, reduction=tf.keras.losses.Reduction.NONE,\n            name='binary_crossentropy_for_correctness'\n        )\n\n    loss_metric = tf.keras.metrics.Mean()    \n    acc_metric = tf.keras.metrics.BinaryAccuracy()\n    auc_metric = tf.keras.metrics.AUC(num_thresholds=2000)    \n    \n#     @tf.function(\n#         input_signature=get_input_signatures(batch_size, seq_len, valid=False)\n#     )\n\n\n    @tf.function(\n        input_signature=get_input_signatures_2()\n    )\n    def predict(batch):\n\n        pred_batch = add_input_ids_and_targets_2(batch)\n        \n        c_mask, r_mask, r_c_mask, c_r_mask = get_attention_masks(pred_batch, training, encoder_decoder, generative, allow_bundle_atten)\n\n        # This is the mask where only the interactions in the current pred batch will be `1`.\n        # shape = [batch_size, window_size]\n        pred_time_mask = pred_batch['pred_time_mask'] * tf.cast(pred_batch['pred_time_mask'] != PAD_TOKEN, dtype=tf.int32)\n\n        logits, answer_logits, c_outputs, r_outputs = predictor(\n            pred_batch, c_mask, r_mask, r_c_mask, c_r_mask,\n            output_attentions=False, output_hidden_states=False, training=False\n        )\n            \n        # preds = tf.math.sigmoid(logits)\n\n        if pred_manager.config.use_softmax:\n            # shape = [batch_size, seq_len, 2]\n            preds = tf.math.softmax(logits)\n        else:\n            # shape = [batch_size, seq_len, 1]\n            preds = tf.math.sigmoid(logits)\n\n        # The places where in the current prediction batch.\n        pred_indices = tf.where(pred_time_mask == 1)\n\n        # shape = [n_selected_places, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n        selected_logits = tf.gather_nd(logits, pred_indices)\n        # shape = [n_selected_places, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n        selected_preds = tf.gather_nd(preds, pred_indices)\n\n        selected_row_ids = tf.gather_nd(pred_batch['row_id'], pred_indices)\n        \n        return logits, preds, selected_logits, selected_preds, selected_row_ids\n    \n    # ----------------------------------------------------\n    \n    # ----------------------------------------------------\n    \n    start_memeory_global = psutil.virtual_memory().available / 1024.0 / 1024.0\n    start = datetime.datetime.now()\n    start_global = start    \n    \n    pred_history = []\n\n    n_pred_batches = 0\n    n_interactions = 0\n    \n    prior_info_for_post_update = None\n    \n    if isinstance(pred_manager.pred_iter_manager, Test_Iter_Manager):\n        if max_steps is not None:\n            print(f'For `Test_Iter_Manager`, {max_steps} shoulde be `None`, and we set it for you here!')\n            max_steps = None\n        \n    total = 0.0\n    for pred_df, pred_dt, prior_group_answers_correct, prior_group_responses in pred_manager.pred_iter_manager.iter_dataset(max_steps=max_steps):\n\n        n_pred_batches += 1\n    \n        try:\n            \n            # ---------------------------------------------------- \n            \n            if DEBUG:\n                \n                pass\n            \n#                 _question_ids = set(pred_df[pred_df['content_type_id'] == 0]['content_id'].values.tolist())\n#                 _lecture_ids = set(pred_df[pred_df['content_type_id'] == 1]['content_id'].values.tolist())              \n            \n#                 # All pred questions must have been seen in training time.\n#                 assert _question_ids.issubset(pred_manager.pred_iter_manager._unique_question_id_train)\n\n#                 # All pred lectures must have been seen in training time.\n#                 assert _lecture_ids.issubset(pred_manager.pred_iter_manager._unique_lecture_id_train)\n\n            pred_record_batch, n_pred_time_steps_batch, info_for_post_update = pred_manager.update(\n                pred_dt, prior_group_answers_correct, prior_group_responses, prior_info_for_post_update\n            )\n            elapsed = (datetime.datetime.now() - start).total_seconds()\n\n            # Update\n            prior_info_for_post_update = info_for_post_update\n            \n            pred_batch = pred_manager.convert_to_pred_batch(pred_record_batch, n_pred_time_steps_batch)\n                                    \n            logits, preds, selected_logits, selected_preds, selected_row_ids = predict(pred_batch)\n                \n            # ----------------------------------------------------            \n            # Compute metrics\n            \n            # shape = [n_selected_places]\n            predictions = selected_preds[:, -1].numpy().tolist()\n            \n            # sanity check\n            assert len(predictions) == len(pred_df['content_type_id'])\n            \n            user_ids = pred_manager.current_batch_users\n            n_interactions += len(user_ids)\n\n            if isinstance(pred_manager.pred_iter_manager, Valid_Iter_Manager):\n\n                row_ids = selected_row_ids.numpy().tolist()\n                \n                # Be careful, the `targets` here is not the same as in `pred_batch['target']`.\n                # In `pred_batch`, we will get `-100` (i.e. `NON_TARGET_ID`) for the places not to be predicted (i.e. padding/lectures),\n                # while in `targets`, we will get `-1`.\n                targets = pred_manager.pred_iter_manager._current_answered_correctly\n                # user_answers = pred_manager.pred_iter_manager._current_user_answer\n\n                # sanity check\n                assert len(set([len(user_ids), len(row_ids), len(targets), len(predictions)])) == 1\n                \n                targets_t = tf.constant(targets, dtype=tf.int32)\n                # Be careful: Here we use `-1` rather than `NON_TARGET_ID`.\n                pred_mask = targets_t != -1\n                pred_indices = tf.where(pred_mask)\n\n                selected_targets_for_loss = tf.gather_nd(targets_t, pred_indices)\n                \n                # shape = [n_selected_places_for_loss, n_targets], `n_targets = 2` if `use_softmax`, otherwise `1`.\n                selected_logits_for_loss = tf.gather_nd(selected_logits, pred_indices)\n                selected_preds_for_loss = tf.gather_nd(selected_preds, pred_indices)\n                \n                # From the doc, `tf.keras.losses.SparseCategoricalCrossentropy` should use `selected_targets_for_loss`, but it seems ok to have the last dimension.\n                # For `BinaryCrossentropy`, we need to use `selected_targets_for_loss[:, tf.newaxis]` to have the 2nd dimension so the losses are not averaged.                \n                losses = loss_obj(selected_targets_for_loss[:, tf.newaxis], selected_logits_for_loss)\n                \n                loss_metric.update_state(losses)\n                # Use `selected_preds_for_loss[:, -1:]` to get the probabilities for class `1`, with the 2nd dim\n                acc_metric.update_state(selected_targets_for_loss[:, tf.newaxis], selected_preds_for_loss[:, -1:])\n                # Use `selected_preds_for_loss[:, -1]` to get the probabilities for class `1`, without the 2nd dim\n                auc_metric.update_state(selected_targets_for_loss, selected_preds_for_loss[:, -1])\n\n                valid_user_ids.extend(user_ids)\n                valid_row_ids.extend(row_ids)\n                valid_targets.extend(targets)\n                valid_preds.extend(predictions)       \n\n            elif isinstance(pred_manager.pred_iter_manager, Test_Iter_Manager):\n\n                pred_df['answered_correctly'] = predictions\n                pred_manager.pred_iter_manager.env.predict(pred_df.loc[pred_df['content_type_id'] == 0, ['row_id', 'answered_correctly']])\n\n            else:\n\n                 raise ValueError('`pred_manager.pred_iter_manager` must be an instance of either `Valid_Iter_Manager` or `Test_Iter_Manager`')\n            \n            # ----------------------------------------------------        \n        \n#             if len(buffer) < max_buffer_size:\n#                 buffer.append({'pred_batch': pred_batch, 'logits': logits, 'preds': preds, 'selected_preds': selected_preds})\n        \n            # ----------------------------------------------------        \n            \n#             if pred_manager.probe:\n\n#                 # ----------------------------------------------------\n#                 # Probe: The number of pred batches\n\n#                 # N_PRED_BATCHES = 57104 # False\n#                 N_PRED_BATCHES = 57120 # True (something strange after API update)\n#                 N_PRED_BATCHES = 60000 # True (need to check again)\n                \n#                 # Disable for now\n#                 ### assert n_pred_batches <= N_PRED_BATCHES\n\n#                 # ----------------------------------------------------\n#                 # Probe: the number of blocks / the number of blocks given by new users\n\n#                 if n_pred_batches > 57100:\n\n#                     ### assert 64 <= pred_manager.avg_n_blocks_from_new_users # False\n#                     assert 62 <= pred_manager.avg_n_blocks_from_new_users_in_pred # True\n\n#                     ### assert 88 <= pred_manager.avg_n_blocks_from_common_users_in_train_and_pred # False\n#                     assert 83 <= pred_manager.avg_n_blocks_from_common_users_in_train_and_pred # True\n                    \n            # ----------------------------------------------------\n\n            if n_pred_batches % 100 == 0:\n                \n                end = datetime.datetime.now()\n                elapsed = (end - start).total_seconds()\n                elapsed_global = (end - start_global).total_seconds()\n                avg_timing_per_batch = elapsed / PRINTING_STEPS\n                avg_timing_per_batch_global = elapsed_global / n_pred_batches\n\n                end_memeory = psutil.virtual_memory().available / 1024.0 / 1024.0\n                used_memory_global = start_memeory_global - end_memeory                \n                \n                print(f'used memory: {used_memory_global}')\n                \n                print(f'pred_manager.n_blocks_in_pred: {pred_manager.n_blocks_in_pred}')\n                print(f'pred_manager.n_blocks_from_common_users_in_train_and_pred: {pred_manager.n_blocks_from_common_users_in_train_and_pred}')\n                print(f'pred_manager.n_blocks_from_new_users_in_pred: {pred_manager.n_blocks_from_new_users_in_pred}')\n\n                print(f'pred_manager.avg_n_blocks_in_pred: {pred_manager.avg_n_blocks_in_pred}')\n                print(f'pred_manager.avg_n_blocks_from_common_users_in_train_and_pred: {pred_manager.avg_n_blocks_from_common_users_in_train_and_pred}')\n                print(f'pred_manager.avg_n_blocks_from_new_users_in_pred: {pred_manager.avg_n_blocks_from_new_users_in_pred}')\n                \n                print(f'number of user interactions processed in current batch: {len(user_ids)}')\n                print(f'total number of user interactions processed: {n_interactions}')\n                print(f'average number of user interactions per batch: {n_interactions / n_pred_batches}')\n                \n                print(f'avg_timing_per_batch: {avg_timing_per_batch}')\n                print(f'avg_timing_per_batch_global: {avg_timing_per_batch_global}')\n                \n                if isinstance(pred_manager.pred_iter_manager, Valid_Iter_Manager):\n                    \n                    valid_loss = loss_metric.result().numpy()\n                    valid_acc = acc_metric.result().numpy()\n                    valid_auc = auc_metric.result().numpy()\n                    \n                    print(f'valid_loss: {valid_loss}')\n                    print(f'valid_acc: {valid_acc}')\n                    print(f'valid_auc: {valid_auc}')\n                \n                start = datetime.datetime.now()\n                \n                print('-' * 32)\n\n            # ----------------------------------------------------        \n\n            # Save a few buffer status to check things are expected.\n            if n_pred_batches <= 4:\n                pred_history.append(str(pred_manager.pred_record_buffer))\n\n        except AssertionError as e:\n            \n            print('some assertions are wrong, breaking the loop')\n            raise e\n\n#     if DEBUG:\n#         assert pred_manager.pred_iter_manager.current_batch_no == n_pred_batches\n    \n#     print(f'total n_pred_batches: {pred_manager.pred_iter_manager.current_batch_no}')\n\n#     print(f'pred_manager.n_blocks_in_pred: {pred_manager.n_blocks_in_pred}')\n#     print(f'pred_manager.n_blocks_from_common_users_in_train_and_pred: {pred_manager.n_blocks_from_common_users_in_train_and_pred}')\n#     print(f'pred_manager.n_blocks_from_new_users_in_pred: {pred_manager.n_blocks_from_new_users_in_pred}')\n\n#     print(f'pred_manager.avg_n_blocks_in_pred: {pred_manager.avg_n_blocks_in_pred}')\n#     print(f'pred_manager.avg_n_blocks_from_common_users_in_train_and_pred: {pred_manager.avg_n_blocks_from_common_users_in_train_and_pred}')\n#     print(f'pred_manager.avg_n_blocks_from_new_users_in_pred: {pred_manager.avg_n_blocks_from_new_users_in_pred}')\n\n#     print(f'number of user interactions processed in current batch: {len(user_ids)}')\n#     print(f'total number of user interactions processed: {n_interactions}')\n#     print(f'average number of user interactions per batch: {n_interactions / n_pred_batches}')    \n    \n    if isinstance(pred_manager.pred_iter_manager, Valid_Iter_Manager):\n        \n        if MAX_VALID_ITER_STEPS is not None:\n            assert not DEBUG\n        else:\n            print(f'pred_manager.pred_iter_manager: {pred_manager.pred_iter_manager._user_manager_dict}')\n        \n        end = datetime.datetime.now()\n        elapsed_global = (end - start_global).total_seconds()\n        avg_timing_per_batch_global = elapsed_global / n_pred_batches\n        print(f'avg_timing_per_batch_global: {avg_timing_per_batch_global}')        \n        \n        valid_loss = float(loss_metric.result().numpy())\n        valid_acc = float(acc_metric.result().numpy())\n        valid_auc = float(auc_metric.result().numpy())\n\n        print(f'valid_loss: {valid_loss}')\n        print(f'valid_acc: {valid_acc}')\n        print(f'valid_auc: {valid_auc}')\n        \n        valid_results = {\n            'valid_loss': valid_loss,\n            'valid_acc': valid_acc,\n            'valid_auc': valid_auc\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        \n        if ckpt_no is None:\n            ckpt_no = 0\n        \n        valid_submission.to_csv(f'valid_submission_ckpt_{ckpt_no}.csv', index=False)\n        with open(f'valid_results_ckpt_{ckpt_no}.json', 'w', encoding='UTF-8') as fp:\n            json.dump(valid_results, fp, ensure_ascii=False, indent=4)\n\n        if not IS_KAGGLE:\n            !gsutil cp -r './valid_results_ckpt_{ckpt_no}.json' '{CKPT_PRED_PATH}'\n            !gsutil cp -r './valid_submission_ckpt_{ckpt_no}.csv' '{CKPT_PRED_PATH}'\n\n        del pred_manager.pred_iter_manager\n        del pred_manager\n        gc.collect()\n            \n    tf.keras.backend.clear_session()","execution_count":null,"outputs":[]},{"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        r = predictor(inputs=inputs)\n\n        return r\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        d_ans_input_ids = tf.constant(1, shape=[3, 5])\n        pos_ids = tf.constant(1, shape=[3, 5])\n        shifted_pos_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        lag_time = tf.constant(1.0, shape=[3, 5])\n        abs_pos_ids = tf.constant(1.0, shape=[3, 5])\n        shifted_abs_pos_ids = tf.constant(1.0, shape=[3, 5])\n        task_container_pos_ids = tf.constant(1.0, shape=[3, 5])\n        correct_answer_id = tf.constant(1, shape=[3, 5])\n        n_questions_answered_scaled = tf.constant(1.0, shape=[3, 5])\n        n_lectures_watched_scaled = tf.constant(1.0, shape=[3, 5])\n        answered_correctly_ratio = tf.constant(1.0, shape=[3, 5])\n        part_correctness_ratio = tf.constant(1.0, shape=[3, 5, PART_VOCAB_SIZE - 2])\n        part_count_scaled = tf.constant(1.0, shape=[3, 5, PART_VOCAB_SIZE - 2])\n        correct_answer_count_scaled = tf.constant(1.0, shape=[3, 5, ANSWER_3_ID - ANSWER_0_ID + 1])\n        correct_answer_correctness_ratio = tf.constant(1.0, shape=[3, 5, ANSWER_3_ID - ANSWER_0_ID + 1])\n        current_part_count_scaled = tf.constant(1.0, shape=[3, 5])\n        current_part_correctness_ratio = tf.constant(1.0, shape=[3, 5])\n        current_correct_answer_count_scaled = tf.constant(1.0, shape=[3, 5])\n        current_correct_answer_correctness_ratio = tf.constant(1.0, shape=[3, 5])\n        current_question_count_scaled = tf.constant(1.0, shape=[3, 5])\n        current_question_correctness_ratio = 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': d_input_ids,\n            'd_ans_input_ids': d_ans_input_ids,\n            'pos_ids': pos_ids,\n            'shifted_pos_ids': shifted_pos_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            'lag_time': lag_time,\n            'abs_pos_ids': abs_pos_ids,\n            'shifted_abs_pos_ids': shifted_abs_pos_ids,\n            'task_container_pos_ids': task_container_pos_ids,\n            'correct_answer_id': correct_answer_id,\n            'n_questions_answered_scaled': n_questions_answered_scaled,\n            'n_lectures_watched_scaled': n_lectures_watched_scaled,\n            'answered_correctly_ratio': answered_correctly_ratio,\n            'part_correctness_ratio': part_correctness_ratio,\n            'part_count_scaled': part_count_scaled,\n            'correct_answer_count_scaled': correct_answer_count_scaled,\n            'correct_answer_correctness_ratio': correct_answer_correctness_ratio,\n            'current_part_count_scaled': current_part_count_scaled,\n            'current_part_correctness_ratio': current_part_correctness_ratio,\n            'current_correct_answer_count_scaled': current_correct_answer_count_scaled,\n            'current_correct_answer_correctness_ratio': current_correct_answer_correctness_ratio,\n            'current_question_count_scaled': current_question_count_scaled,\n            'current_question_correctness_ratio': current_question_correctness_ratio,          \n        }\n\n        r1 = foo(inputs)\n    \n        return r1","execution_count":null,"outputs":[]},{"metadata":{"id":"Ltf4Pc_z11w7","trusted":true},"cell_type":"code","source":"def load_ckpt(predictor, optimizer, ckpt_path, ckpt_no=-1):\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        # ----------------------------------------------------------------------\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        # Load ckpt\n\n        print(f'try to load {ckpt_file}')\n        if ckpt_no_to_load > 0:\n            ### pass\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        # ----------------------------------------------------------------------\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":"9CH2gEjp11yP","trusted":true},"cell_type":"code","source":"# Only for a quick change, shouldn't use in real commit or submission\n\nDUMMY = True\n\nif SUBMISSION:\n    DUMMY = False","execution_count":null,"outputs":[]},{"metadata":{"id":"Ti-aQwLD11yb","trusted":true},"cell_type":"code","source":"# Only for a quick change, shouldn't use in real commit or submission\n\nif DUMMY:\n\n    # --------------------------------------------------\n\n    MODEL_TYPE = 'ed'\n    MODEL_SIZE = 'b-2'\n\n    # Our base config\n    # MODEL_DESC = 'lag-user-corr-ans-enc-loss'\n\n    # Addon\n    MODEL_DESC = 'master-train-all'\n\n    ACTIVATION = 'gelu'\n    USE_PRE_CLASSIFIER = False\n    USE_SOFTMAX = False\n    USE_USER_ANSWER = True\n    USE_USER_ANSWER_LOSS = True\n    USE_CORRECT_ANSWER_FOR_ENCODER = True\n    USE_CORRECT_ANSWER_FOR_DECODER = False\n    USE_ABS_POS = False\n    USE_TASK_CONTAINER_POS = False\n    SHARE_POS_EMBEDDING = True\n    USE_TAGS = True\n    USE_PART = True\n    USE_PRIOR_EXPLANATION = True\n    USE_PRIOR_QUESTION_ELAPSED_TIME_INPUT = True\n    USE_LAG_TIME = True\n    USE_LAG_TIME_FOR_ENCODER=False\n    USE_USER_LEVEL_AGGREGATED_HISTORICAL_INFO = True\n    USE_PART_AGGREGATED_HISTORICAL_INFO = True\n    USE_CORRECT_ANSWER_AGGREGATED_HISTORICAL_INFO = True\n    USE_QUESTION_LEVEL_AGGREGATED_HISTORICAL_INFO = True\n    ALLOW_BUNDLE_ATTEN = True\n    GENERATIVE = False\n\n    VALID_FOLD = 1\n    \n    WINDOW_SIZE = 128\n    LOSS_WEIGHT_WINDOW_SIZE = None\n    BATCH_SIZE = 16 * strategy.num_replicas_in_sync\n    if not tpu:\n        BATCH_SIZE = 8 * BATCH_SIZE\n    PRED_BATCH_SIZE = 256 * strategy.num_replicas_in_sync\n    if tpu:\n        assert BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\n    assert PRED_BATCH_SIZE // strategy.num_replicas_in_sync != WINDOW_SIZE\n    N_EPOCHS = 100\n    STEPS_PER_CALL = 1000\n    MAX_N_CONTENTS_PER_USER_FOR_SAMPLING_PROB = 512\n\n    LR = 2e-4\n    END_LR = 2e-5\n    WARMUP_STEPS = 80000\n\n    DETERMINISTIC = False\n\n    if DETERMINISTIC:\n\n        SEED = 2021\n        N_PARALLEL_READS = None  # 1\n        N_PARALLEL_CALLS = None  # 1\n        SHUFFLE_BUFFER_SIZE = 1\n\n    else:\n        \n        SEED = None\n        N_PARALLEL_READS = 16\n        N_PARALLEL_CALLS = tf.data.experimental.AUTOTUNE\n        SHUFFLE_BUFFER_SIZE = 65536\n\n    MAX_TRAIN_ITER_STEPS = None\n    MAX_VALID_ITER_STEPS = None\n\n    PRINTING_STEPS = 1000\n\n    CKPT_DIR = f'{MODEL_TYPE}-{MODEL_SIZE}-{MODEL_DESC}/'\n\n    # --------------------------------------------------\n\n    TRAIN = False\n    VALID = False\n    PRED = True\n\n    RESUME_TRAINING = True\n\n    N_FILES = 6\n    SUBMISSION = False\n\n    if IS_KAGGLE:\n        N_FILES = len(os.listdir('/kaggle/input/riiid-test-answer-prediction'))\n        SUBMISSION = (N_FILES != 6)\n    else:\n        PRED = False\n\n    if SUBMISSION:\n        \n        TRAIN = False\n        VALID = False\n        PRED = True\n\n    if not TRAIN:\n        RESUME_TRAINING = False\n\n    DEBUG = False\n    PROBE = False\n\n    CKPT_TRAIN_PATH = None\n    CKPT_PRED_PATH = None\n    if 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\n    if 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":{"id":"JUgpMNIuhn5D","trusted":true},"cell_type":"code","source":"FFN_FACTOR = 4\n\nif MODEL_SIZE == 'baby':\n    # Only for debug\n\n    n_layers = 1\n    n_heads = 2\n    dim = 4\n    hidden_dim = dim\n\nelif MODEL_SIZE == 't-0':\n\n    n_layers = 2\n    n_heads = 4\n    dim = 32\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 't-1':\n\n    n_layers = 2\n    n_heads = 4\n    dim = 64\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 't-2':\n\n    n_layers = 2\n    n_heads = 4\n    dim = 128\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 't-3':\n\n    n_layers = 2\n    n_heads = 4\n    dim = 256\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 's-0':\n\n    n_layers = 4\n    n_heads = 4\n    dim = 32\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 's-1':\n\n    n_layers = 4\n    n_heads = 4\n    dim = 64\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 's-2':\n\n    n_layers = 4\n    n_heads = 4\n    dim = 128\n    hidden_dim = FFN_FACTOR * dim                                                \n\nelif MODEL_SIZE == 's-3':\n\n    n_layers = 4\n    n_heads = 4\n    dim = 256\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'b-0':\n\n    n_layers = 4\n    n_heads = 8\n    dim = 64\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'b-1':\n\n    n_layers = 4\n    n_heads = 8\n    dim = 128\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'b-2':\n\n    n_layers = 4\n    n_heads = 8\n    dim = 256\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'b-3':\n\n    n_layers = 4\n    n_heads = 8\n    dim = 512\n    hidden_dim = FFN_FACTOR * dim // 2\n\nelif MODEL_SIZE == 'b-4':\n\n    n_layers = 4\n    n_heads = 16\n    dim = 512\n    hidden_dim = FFN_FACTOR * dim  // 2\n\nelif MODEL_SIZE == 'm-0':\n\n    n_layers = 6\n    n_heads = 8\n    dim = 64\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'm-1':\n\n    n_layers = 6\n    n_heads = 8\n    dim = 128\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'm-2':\n\n    n_layers = 6\n    n_heads = 8\n    dim = 256\n    hidden_dim = FFN_FACTOR * dim\n\nelif MODEL_SIZE == 'm-3':\n\n    n_layers = 6\n    n_heads = 8\n    dim = 512\n    hidden_dim = FFN_FACTOR * dim // 2\n\nelif MODEL_SIZE == 'm-4':\n\n    n_layers = 6\n    n_heads = 16\n    dim = 512\n    hidden_dim = FFN_FACTOR * dim // 2\n\n# --------------------------------------------------\n\nassert MODEL_TYPE is not None and MODEL_TYPE != ''\nassert MODEL_DESC is not None and MODEL_DESC != ''","execution_count":null,"outputs":[]},{"metadata":{"id":"fWVxPT8E2Vpb","trusted":true},"cell_type":"code","source":"# `0` is used for padding, the real position is from `1` to `WINDOW_SIZE`.\nmax_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_user_answer=USE_USER_ANSWER,\n    use_user_answer_loss=USE_USER_ANSWER_LOSS,\n    use_correct_answer_for_encoder=USE_CORRECT_ANSWER_FOR_ENCODER,\n    use_correct_answer_for_decoder=USE_CORRECT_ANSWER_FOR_DECODER,    \n    use_abs_pos=USE_ABS_POS,\n    use_task_container_pos=USE_TASK_CONTAINER_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_lag_time=USE_LAG_TIME,\n    use_lag_time_for_encoder=USE_LAG_TIME_FOR_ENCODER,\n    use_user_level_aggregated_historical_info=USE_USER_LEVEL_AGGREGATED_HISTORICAL_INFO,\n    use_part_aggregated_historical_info=USE_PART_AGGREGATED_HISTORICAL_INFO,\n    use_correct_answer_aggregated_historical_info=USE_CORRECT_ANSWER_AGGREGATED_HISTORICAL_INFO,\n    use_question_level_aggregated_historical_info=USE_QUESTION_LEVEL_AGGREGATED_HISTORICAL_INFO,\n    allow_bundle_atten=ALLOW_BUNDLE_ATTEN,\n    generative=GENERATIVE,\n    use_pre_classifier=USE_PRE_CLASSIFIER,\n    use_softmax=USE_SOFTMAX,\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    max_n_contents_per_user_for_sampling_prob=MAX_N_CONTENTS_PER_USER_FOR_SAMPLING_PROB,\n    valid_fold=VALID_FOLD\n)","execution_count":null,"outputs":[]},{"metadata":{"id":"fQ2AaCkjWjq9"},"cell_type":"markdown","source":"### check tensor input is good"},{"metadata":{"id":"AxvDlySG11yq","trusted":true},"cell_type":"code","source":"if TRAIN:\n    \n    from_valid = False\n    only_valid = False\n    ckpts = [0]\n    valid_epochs = [\n        5, 10, 15, 20,\n        21, 22, 23, 24, 25,\n        26, 27, 28, 29, 30,\n        31, 32, 33, 34, 35,\n        36, 37, 38, 39, 40,\n        41, 42, 43, 44, 45,\n        46, 47, 48, 49, 50,\n        51, 52, 53, 54, 55,\n        56, 57, 58, 59, 60,               \n    ]\n\n    train_manager = Train_Manager(config, train_config)\n    predictor, optimizer, loss_obj, loss_obj_answer, metrics = train_manager.get_train_objs(\n        lr=LR, end_lr=END_LR, warmup_steps=WARMUP_STEPS\n    )\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(\n                predictor, optimizer,\n                CKPT_TRAIN_PATH, ckpt_no=ckpt_no\n            )\n\n        else:\n\n            with strategy.scope():\n\n                loaded_ckpt_no = 0\n\n                checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=predictor)\n                ckpt_path = CKPT_TRAIN_PATH\n                ckpt_manager = tf.train.CheckpointManager(checkpoint, directory=ckpt_path, max_to_keep=None)\n\n        train_manager.train_valid(\n            predictor, optimizer, ckpt_manager,\n            loss_obj, loss_obj_answer, metrics,\n            last_epoch=loaded_ckpt_no,\n            from_valid=from_valid,\n            valid_epochs=valid_epochs,\n            only_valid=only_valid\n        )","execution_count":null,"outputs":[]},{"metadata":{"id":"2COvA_dof18z","outputId":"7a8e3a31-f124-403d-9e66-cb4fee539e5a","trusted":true},"cell_type":"code","source":"CKPT_TRAIN_PATH","execution_count":null,"outputs":[]},{"metadata":{"id":"s06dw6EBR2ko","outputId":"d0412c6a-1ba3-4dd5-919c-7c03287eb8df","trusted":true},"cell_type":"code","source":"!ls -l","execution_count":null,"outputs":[]},{"metadata":{"id":"kyBdTym011y9","outputId":"dafaa14d-f456-428d-d1be-728a2f4cde19","trusted":true},"cell_type":"code","source":"if CKPT_TRAIN_PATH.startswith('gs://'):\n    !gsutil ls '{CKPT_TRAIN_PATH}'\nelse:\n    !ls -l '{CKPT_TRAIN_PATH}'","execution_count":null,"outputs":[]},{"metadata":{"id":"m1hq4KTh11zH","trusted":true},"cell_type":"code","source":"predictor = TFEdFormerAnswerPredictionModel(config)\n\nif ((VALID or PRED) and not DUMMY) or SUBMISSION:\n    assert CKPT_PRED_PATH is not None and CKPT_PRED_PATH != './'","execution_count":null,"outputs":[]},{"metadata":{"id":"dvRJTq1J3-Wt","outputId":"68c82831-7094-4f09-c1e1-3f27a04499ab","trusted":true},"cell_type":"code","source":"CKPT_PRED_PATH","execution_count":null,"outputs":[]},{"metadata":{"id":"96vKNooU11zR","outputId":"8152077b-dfa3-4d7d-a0e6-b22959c356ec","trusted":true},"cell_type":"code","source":"if VALID:\n    \n    valid_results = {'ckpt_no': [], 'valid_loss': [], 'valid_acc': [], 'valid_auc': []}\n\n    ckpts_to_run = [60]\n                    \n    print(f'ckpts_to_run: {ckpts_to_run}')\n        \n    for ckpt_no in ckpts_to_run:\n            \n        # reload later ckpt\n        if ckpt_no > 0:\n\n            ckpt_manager, loaded_ckpt_no = load_ckpt(\n                predictor=predictor, optimizer=None,\n                ckpt_path=CKPT_PRED_PATH, ckpt_no=ckpt_no,\n            )\n            assert loaded_ckpt_no > 0\n        \n        pred_iter_manager_class = Valid_Iter_Manager\n\n        pred_manager = Pred_Manager.create(\n            config=config,\n            train_config=train_config,\n            pred_iter_manager_class=pred_iter_manager_class,\n            train_dt_path_or_obj=train_dt,               \n            user_id_to_row_id_train_path=user_id_to_row_id_train_path,\n            unique_question_id_train_path=unique_question_id_splitted_train_path,\n            unique_lecture_id_train_path=unique_lecture_id_splitted_train_path,\n            # ------------------------------------------------------------------------------------------\n            # this should be replaced with the corresponding files for validation\n#             question_history_at_training_end_path=question_history_at_training_end_path,\n#             single_question_history_at_training_end_optimized_path=single_question_history_at_training_end_optimized_path,\n          \n            question_history_at_training_end_path=question_history_at_training_end_for_valid_path,\n            single_question_history_at_training_end_optimized_path=single_question_history_at_training_end_optimized_for_valid_path,\n\n            # ------------------------------------------------------------------------------------------\n            valid_info_path=(valid_info_paths[VALID_FOLD] if VALID_FOLD in [0, 1, 2, 3] else None),\n            max_train_buffer_size=30000,\n            probe=False,\n            debug=False\n        )\n\n        run_pred(pred_manager, predictor, max_steps=MAX_VALID_ITER_STEPS, ckpt_no=ckpt_no)\n        \n        with open(f'valid_results_ckpt_{ckpt_no}.json', 'r', encoding='UTF-8') as fp:\n            \n            valid_results_ckpt = json.load(fp)\n            \n            valid_results['ckpt_no'].append(ckpt_no)\n            valid_results['valid_loss'].append(valid_results_ckpt['valid_loss'])\n            valid_results['valid_acc'].append(valid_results_ckpt['valid_acc'])\n            valid_results['valid_auc'].append(valid_results_ckpt['valid_auc'])\n            \n        print('=' * 80)\n        \n    with open(f'valid_results.json', 'w', encoding='UTF-8') as fp:\n        json.dump(valid_results, fp, ensure_ascii=False, indent=True)\n\n    if not IS_KAGGLE:\n        !gsutil cp -r './valid_results.json' '{CKPT_PRED_PATH}'","execution_count":null,"outputs":[]},{"metadata":{"id":"D3DJBpGA11zf","trusted":true},"cell_type":"code","source":"!ls -l","execution_count":null,"outputs":[]},{"metadata":{"id":"VhSrNX2U11zz","trusted":true},"cell_type":"code","source":"valid_df = None\nif VALID:\n\n    valid_df = pd.read_csv(f'valid_submission_ckpt_{ckpts_to_run[-1]}.csv')\n    pd.set_option('display.max_rows', 200)\n\nvalid_df","execution_count":null,"outputs":[]},{"metadata":{"id":"wI0VmveY110E","trusted":true},"cell_type":"code","source":"valid_results = {}\nif VALID:\n\n    with open('valid_results.json', 'r', encoding='UTF-8') as fp:\n        valid_results = json.load(fp)\n\nvalid_results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CKPT_PRED_PATH","execution_count":null,"outputs":[]},{"metadata":{"id":"puFpAPGw110X","trusted":true},"cell_type":"code","source":"pred_start = datetime.datetime.now()\nelapsed = (pred_start - kernel_start).total_seconds()\nprint(f'time before prediction: {elapsed} seconds')\n\nif PRED:\n    \n    ckpt_no = 60\n    \n    # reload ckpt\n    ckpt_manager, loaded_ckpt_no = load_ckpt(\n        predictor=predictor, optimizer=None,\n        ckpt_path=CKPT_PRED_PATH, ckpt_no=ckpt_no,\n    )\n    \n    if SUBMISSION:\n        assert loaded_ckpt_no > 0\n    \n    pred_iter_manager_class = Test_Iter_Manager\n\n    pred_manager = Pred_Manager.create(\n        config=config,\n        train_config=train_config,\n        pred_iter_manager_class=pred_iter_manager_class,\n        train_dt_path_or_obj=train_dt,  \n        user_id_to_row_id_train_path=user_id_to_row_id_train_path,\n        unique_question_id_train_path=unique_question_id_train_path,\n        unique_lecture_id_train_path=unique_lecture_id_train_path,\n        question_history_at_training_end_path=question_history_at_training_end_path,\n        single_question_history_at_training_end_optimized_path=single_question_history_at_training_end_optimized_path,\n        valid_info_path=None,   \n        max_train_buffer_size=30000,\n        probe=PROBE,\n        debug=False        \n    )\n\n    run_pred(pred_manager, predictor, max_steps=None)","execution_count":null,"outputs":[]},{"metadata":{"id":"ZN8mrPjvJOuR","trusted":true},"cell_type":"code","source":"!ls -l","execution_count":null,"outputs":[]}],"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}