{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Training on GCP + TPU"},{"metadata":{},"cell_type":"markdown","source":"# Important\n\n* Please change `FLAGS.model_name` from `distilbert-base-uncased-distilled-squad` to the one you want to use.\n\n* Please change `FLAGS.epochs` from `0` to the no. of epochs you want.\n\n* Please change `FLAGS.train_batch_size` from  `64 * 8` to `8 * 8` if you decide to use `bert-large-uncased-whole-word-masking-finetuned-squad`\n\n* Change values for `TPU_WORKER`, `ZONE`, `PROJECT`, `TRAIN_TF_RECORD` and `CHECKPOINT_DIR`."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport json\nimport tensorflow as tf\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport absl\nimport datetime\nfrom tensorflow.keras.optimizers import Adam\nfrom adamw_optimizer import AdamW\nfrom tensorflow.python.lib.io.file_io import recursive_create_dir\n\nprint(tf.__version__)\n\n\n# # ----------------------------------------------------------------------------------------\n# Ref: https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/custom_training.ipynb#scrollTo=jwJtsCQhHK-E\n\n\n# Your TPU node internal ip\nTPU_WORKER = 'grpc://XXX.XXX.XXX.XXX:8470'\n\n# Your TPU Zone, for example 'europe-west4-a'\nZONE = ''\n\n# Your project name, for example, 'kaggle-nq-123456'\nPROJECT = ''\n\n# Your training tf record file on Google Storage bucket. For example, gs://kaggle-my-nq-competition/nq_train.tfrecord\nTRAIN_TF_RECORD = ''\n\n# Your checkpoint dir on Google Storage bucket. For example, \"gs://kaggle-my-nq-competition/checkpoints/\"\nCHECKPOINT_DIR = ''\n\n\ntf.keras.backend.clear_session()\n\n# # ----------------------------------------------------------------------------------------\n\n\nIS_KAGGLE = True\nINPUT_DIR = \"/kaggle/input/\"\n\n# The original Bert Joint Baseline data.\nBERT_JOINT_BASE_DIR = os.path.join(INPUT_DIR, \"bertjointbaseline\")\n\n# This nq dir contains all files for publicly use.\nNQ_DIR = os.path.join(INPUT_DIR, \"nq-competition\")\n\n\nMY_OWN_NQ_DIR = NQ_DIR\n\n# For local usage.\nif not os.path.isdir(INPUT_DIR):\n    IS_KAGGLE = False\n    INPUT_DIR = \"./\"\n    NQ_DIR = \"./\"\n    MY_OWN_NQ_DIR = \"./\"\n\n\nfor dirname, _, filenames in os.walk(INPUT_DIR):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n\n# NQ_DIR contains some packages / modules\nsys.path.append(NQ_DIR)\nsys.path.append(os.path.join(NQ_DIR, \"transformers\"))\n\nfrom nq_flags import DEFAULT_FLAGS as FLAGS\nfrom nq_flags import del_all_flags\nfrom nq_dataset_utils import *\n\nimport sacremoses as sm\nimport transformers\nfrom adamw_optimizer import AdamW\n\nfrom transformers import TFBertModel\nfrom transformers import TFBertMainLayer, TFBertPreTrainedModel\nfrom transformers.modeling_tf_utils import get_initializer\n\n\nfrom transformers import BertTokenizer\nfrom transformers import TFBertModel, TFDistilBertModel\nfrom transformers import TFBertMainLayer, TFDistilBertMainLayer, TFBertPreTrainedModel, TFDistilBertPreTrainedModel\nfrom transformers.modeling_tf_utils import get_initializer\n\n\nPRETRAINED_MODELS = {\n    \"BERT\": [\n        'bert-base-uncased',\n        'bert-large-uncased-whole-word-masking-finetuned-squad',\n    ],\n    \"DISTILBERT\": [\n        'distilbert-base-uncased-distilled-squad'\n    ]\n}\n\n\nflags = absl.flags\ndel_all_flags(flags.FLAGS)\n\nflags.DEFINE_bool(\n    \"do_lower_case\", True,\n    \"Whether to lower case the input text. Should be True for uncased \"\n    \"models and False for cased models.\")\n\nvocab_file = os.path.join(NQ_DIR, \"vocab-nq.txt\")\n\nflags.DEFINE_string(\"vocab_file\", vocab_file,\n                    \"The vocabulary file that the BERT model was trained on.\")\n\nflags.DEFINE_integer(\n    \"max_seq_length_for_training\", 512,\n    \"The maximum total input sequence length after WordPiece tokenization for training examples. \"\n    \"Sequences longer than this will be truncated, and sequences shorter \"\n    \"than this will be padded.\")\n\nflags.DEFINE_integer(\n    \"max_seq_length\", 512,\n    \"The maximum total input sequence length after WordPiece tokenization. \"\n    \"Sequences longer than this will be truncated, and sequences shorter \"\n    \"than this will be padded.\")\n\nflags.DEFINE_integer(\n    \"doc_stride\", 128,\n    \"When splitting up a long document into chunks, how much stride to \"\n    \"take between chunks.\")\n\nflags.DEFINE_float(\n    \"include_unknowns_for_training\", 0.02,\n    \"If positive, for converting training dataset, probability of including answers of type `UNKNOWN`.\")\n\nflags.DEFINE_float(\n    \"include_unknowns\", -1.0,\n    \"If positive, probability of including answers of type `UNKNOWN`.\")\n\nflags.DEFINE_boolean(\n    \"skip_nested_contexts\", True,\n    \"Completely ignore context that are not top level nodes in the page.\")\n\nflags.DEFINE_integer(\"max_contexts\", 48,\n                     \"Maximum number of contexts to output for an example.\")\n\nflags.DEFINE_integer(\n    \"max_position\", 50,\n    \"Maximum context position for which to generate special tokens.\")\n\nflags.DEFINE_integer(\n    \"max_query_length\", 64,\n    \"The maximum number of tokens for the question. Questions longer than \"\n    \"this will be truncated to this length.\")\n\n\n\n    \nflags.DEFINE_string(\"train_tf_record\", TRAIN_TF_RECORD,\n                    \"Precomputed tf records for training dataset.\")\n\nflags.DEFINE_bool(\"do_train\", False, \"Whether to run training dataset.\")\n\n\nflags.DEFINE_string(\n    \"input_checkpoint_dir\", CHECKPOINT_DIR,\n    \"The root directory that contains checkpoints to be loaded of all trained models.\")\n\nflags.DEFINE_string(\"model_dir\", NQ_DIR, \"Root dir of all Hugging Face's models\")\n\nflags.DEFINE_string(\"model_name\", \"distilbert-base-uncased-distilled-squad\", \"Name of Hugging Face's model to use.\")\n\nflags.DEFINE_integer(\"epochs\", 0, \"Total epochs for training.\")\n\nflags.DEFINE_integer(\"train_batch_size\", 64 * 8, \"Batch size for training.\")\n\nflags.DEFINE_integer(\"shuffle_buffer_size\", 100000, \"Shuffle buffer size for training.\")\n\nflags.DEFINE_float(\"init_learning_rate\", 5e-5, \"The initial learning rate for AdamW optimizer.\")\n\nflags.DEFINE_bool(\"cyclic_learning_rate\", True, \"If to use cyclic learning rate.\")\n\nflags.DEFINE_float(\"init_weight_decay_rate\", 0.01, \"The initial weight decay rate for AdamW optimizer.\")\n\nflags.DEFINE_integer(\"num_warmup_steps\", 0, \"Number of training steps to perform linear learning rate warmup.\")\n\nflags.DEFINE_integer(\"num_train_examples\", None, \"Number of precomputed training steps in 1 epoch.\")\n\n# Make the default flags as parsed flags\nFLAGS.mark_as_parsed()\n\nNB_SHORT_ANSWER_TYPES = 5\n\n# ----------------------------------------------------------------------------------------\n\n\ndef get_dataset(tf_record_file, seq_length, batch_size=1, shuffle_buffer_size=0, is_training=False):\n\n    if is_training:\n        features = {\n            \"unique_ids\": tf.io.FixedLenFeature([], tf.int64),\n            \"input_ids\": tf.io.FixedLenFeature([seq_length], tf.int64),\n            \"input_mask\": tf.io.FixedLenFeature([seq_length], tf.int64),\n            \"segment_ids\": tf.io.FixedLenFeature([seq_length], tf.int64),\n            \"start_positions\": tf.io.FixedLenFeature([], tf.int64),\n            \"end_positions\": tf.io.FixedLenFeature([], tf.int64),\n            \"answer_types\": tf.io.FixedLenFeature([], tf.int64)\n        }\n    else:\n        features = {\n            \"unique_ids\": tf.io.FixedLenFeature([], tf.int64),\n            \"input_ids\": tf.io.FixedLenFeature([seq_length], tf.int64),\n            \"input_mask\": tf.io.FixedLenFeature([seq_length], tf.int64),\n            \"segment_ids\": tf.io.FixedLenFeature([seq_length], tf.int64)\n        }        \n\n    def decode_record(record, features):\n        \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n        example = tf.io.parse_single_example(record, features)\n\n        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n        # So cast all int64 to int32.\n        for name in list(example.keys()):\n            t = example[name]\n            if t.dtype == tf.int64:\n                t = tf.cast(t, tf.int32)\n            example[name] = t\n        return example\n\n    def select_data_from_record(record):\n        \n        x = {\n            'unique_ids': record['unique_ids'],\n            'input_ids': record['input_ids'],\n            'input_mask': record['input_mask'],\n            'segment_ids': record['segment_ids']\n        }\n\n        if is_training:\n            y = {\n                'short_start_positions': record['start_positions'],\n                'short_end_positions': record['end_positions'],\n                'short_answer_types': record['answer_types']\n            }\n\n            return (x, y)\n        \n        return x\n\n    dataset = tf.data.TFRecordDataset(tf_record_file)\n    \n    dataset = dataset.map(lambda record: decode_record(record, features))\n    dataset = dataset.map(select_data_from_record)\n    \n    if shuffle_buffer_size > 0:\n        dataset = dataset.shuffle(shuffle_buffer_size)\n    \n    dataset = dataset.batch(batch_size, drop_remainder=True)\n    \n    return dataset\n\n\n# ----------------------------------------------------------------------------------------\n\n\ncluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER, zone=ZONE, project=PROJECT)\ntf.config.experimental_connect_to_cluster(cluster_resolver)\ntf.tpu.experimental.initialize_tpu_system(cluster_resolver)\ntpu_strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)\n\n\n# ----------------------------------------------------------------------------------------\n\n\nif FLAGS.num_train_examples is None:\n    FLAGS.num_train_examples = 494670\n\n\n# ----------------------------------------------------------------------------------------\n\n\nclass TFNQModel:\n    \n    def __init__(self, config, *inputs, **kwargs):\n        \"\"\"\n        \n        Subclasses of this class are different in self.backend,\n        which should be a model that outputs a tensor of shape (batch_size, hidden_dim), and the\n        `backend_call()` method.\n        \n        We will use Hugging Face Bert/DistilBert as backend in this notebook.\n        \"\"\"\n\n        self.backend = None\n        \n        self.seq_output_dropout = tf.keras.layers.Dropout(kwargs.get('seq_output_dropout_prob', 0.05))\n        self.pooled_output_dropout = tf.keras.layers.Dropout(kwargs.get('pooled_output_dropout_prob', 0.05))\n    \n        self.short_pos_classifier = tf.keras.layers.Dense(2,\n                                        kernel_initializer=get_initializer(config.initializer_range),\n                                        name='pos_classifier')       \n\n        self.short_answer_type_classifier = tf.keras.layers.Dense(NB_SHORT_ANSWER_TYPES,\n                                        kernel_initializer=get_initializer(config.initializer_range),\n                                        name='answer_type_classifier')        \n                \n    def backend_call(self, inputs, **kwargs):\n        \"\"\"This method should be implemented by subclasses.\n           \n           The implementation should take into account the (somehow) different input formats of Hugging Face's\n           models.\n           \n           For example, the `TFDistilBert` model, unlike `Bert` model, doesn't have segment_id as input.\n           \n           Then it calls `self.backend_call()` to get the outputs from Bert's model, which is used in self.call().\n        \"\"\"\n        \n        raise NotImplementedError\n\n    \n    def call(self, inputs, **kwargs):\n        \n        # sequence / [CLS] outputs from original bert\n        sequence_output, pooled_output = self.backend_call(inputs, **kwargs)  # shape = (batch_size, seq_len, hidden_dim) / (batch_size, hidden_dim)\n        \n        # dropout\n        sequence_output = self.seq_output_dropout(sequence_output, training=kwargs.get('training', False))\n        pooled_output = self.pooled_output_dropout(pooled_output, training=kwargs.get('training', False))\n    \n        short_pos_logits = self.short_pos_classifier(sequence_output)  # shape = (batch_size, seq_len, 2)\n        \n        short_start_pos_logits = short_pos_logits[:, :, 0]  # shape = (batch_size, seq_len)\n        short_end_pos_logits = short_pos_logits[:, :, 1]  # shape = (batch_size, seq_len)\n        \n        short_answer_type_logits = self.short_answer_type_classifier(pooled_output)  # shape = (batch_size, NB_SHORT_ANSWER_TYPES)\n\n        outputs = (short_start_pos_logits, short_end_pos_logits, short_answer_type_logits)\n\n        return outputs  # logits\n    \n    \nclass TFBertForNQ(TFNQModel, TFBertPreTrainedModel):\n    \n    def __init__(self, config, *inputs, **kwargs):\n        \n        TFBertPreTrainedModel.__init__(self, config, *inputs, **kwargs)  # explicit calls without super\n        TFNQModel.__init__(self, config)\n\n        self.bert = TFBertMainLayer(config, name='bert')\n    \n    def backend_call(self, inputs, **kwargs):\n        \n        outputs = self.bert(inputs, **kwargs)\n        sequence_output, pooled_output = outputs[0], outputs[1]  # shape = (batch_size, seq_len, hidden_dim) / (batch_size, hidden_dim)\n        \n        return sequence_output, pooled_output\n\n\nclass TFDistilBertForNQ(TFNQModel, TFDistilBertPreTrainedModel):\n    \n    def __init__(self, config, *inputs, **kwargs):\n        \n        TFDistilBertPreTrainedModel.__init__(self, config, *inputs, **kwargs)  # explicit calls without super\n        TFNQModel.__init__(self, config)\n\n        self.backend = TFDistilBertMainLayer(config, name=\"distilbert\")\n        \n    def backend_call(self, inputs, **kwargs):\n        \n        if isinstance(inputs, tuple):\n            # Distil bert has no segment_id (i.e. `token_type_ids`)\n            inputs = inputs[:2]\n        else:\n            inputs = inputs\n        \n        outputs = self.backend(inputs, **kwargs)\n        \n        # TFDistilBertModel's output[0] is of shape (batch_size, sequence_length, hidden_size)\n        # We take only for the [CLS].\n        \n        sequence_output = outputs[0]  # shape = (batch_size, seq_len, hidden_dim)\n        pooled_output = sequence_output[:, 0, :]  # shape = (batch_size, hidden_dim)\n        \n        return sequence_output, pooled_output\n\n\nmodel_mapping = {\n    \"bert\": TFBertForNQ,\n    \"distilbert\": TFDistilBertForNQ\n}\n\n\ndef get_pretrained_model(model_name):\n    \n    pretrained_path = os.path.join(FLAGS.model_dir, model_name)\n    \n    tokenizer = BertTokenizer.from_pretrained(pretrained_path)\n    \n    model_type = model_name.split(\"-\")[0]\n    if model_type not in model_mapping:\n        raise ValueError(\"Model definition not found.\")\n    \n    model_class = model_mapping[model_type]\n    model = model_class.from_pretrained(pretrained_path)\n    \n    return tokenizer, model\n\n\ndef get_metrics(name):\n\n    loss = tf.keras.metrics.Mean(name=f'{name}_loss')\n\n    loss_short_start_pos = tf.keras.metrics.Mean(name=f'{name}_loss_short_start_pos')\n    loss_short_end_pos = tf.keras.metrics.Mean(name=f'{name}_loss_short_end_pos')\n    loss_short_ans_type = tf.keras.metrics.Mean(name=f'{name}_loss_short_ans_type')\n    \n    acc = tf.keras.metrics.SparseCategoricalAccuracy(name=f'{name}_acc')\n    \n    acc_short_start_pos = tf.keras.metrics.SparseCategoricalAccuracy(name=f'{name}_acc_short_start_pos')\n    acc_short_end_pos = tf.keras.metrics.SparseCategoricalAccuracy(name=f'{name}_acc_short_end_pos')\n    acc_short_ans_type = tf.keras.metrics.SparseCategoricalAccuracy(name=f'{name}_acc_short_ans_type')\n    \n    return loss, loss_short_start_pos, loss_short_end_pos, loss_short_ans_type, acc, acc_short_start_pos, acc_short_end_pos, acc_short_ans_type\n\n\nclass CustomSchedule(tf.keras.optimizers.schedules.PolynomialDecay):\n    \n    def __init__(self,\n      initial_learning_rate,\n      decay_steps,\n      end_learning_rate=0.0001,\n      power=1.0,\n      cycle=False,\n      name=None,\n      num_warmup_steps=1000):\n        \n        # Since we have a custom __call__() method, we pass cycle=False when calling `super().__init__()` and\n        # in self.__call__(), we simply do `step = step % self.decay_steps` to have cyclic behavior.\n        super(CustomSchedule, self).__init__(initial_learning_rate, decay_steps, end_learning_rate, power, cycle=False, name=name)\n        \n        self.num_warmup_steps = num_warmup_steps\n        \n        self.cycle = tf.constant(cycle, dtype=tf.bool)\n        \n    def __call__(self, step):\n        \"\"\" `step` is actually the step index, starting at 0.\n        \"\"\"\n        \n        # For cyclic behavior\n        step = tf.cond(self.cycle and step >= self.decay_steps, lambda: step % self.decay_steps, lambda: step)\n        \n        learning_rate = super(CustomSchedule, self).__call__(step)\n\n        # Copy (including the comments) from original bert optimizer with minor change.\n        # Ref: https://github.com/google-research/bert/blob/master/optimization.py#L25\n        \n        # Implements linear warmup: if global_step < num_warmup_steps, the\n        # learning rate will be `global_step / num_warmup_steps * init_lr`.\n        if self.num_warmup_steps > 0:\n            \n            steps_int = tf.cast(step, tf.int32)\n            warmup_steps_int = tf.constant(self.num_warmup_steps, dtype=tf.int32)\n\n            steps_float = tf.cast(steps_int, tf.float32)\n            warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)\n\n            # The first training step has index (`step`) 0.\n            # The original code use `steps_float / warmup_steps_float`, which gives `warmup_percent_done` being 0,\n            # and causing `learning_rate` = 0, which is undesired.\n            # For this reason, we use `(steps_float + 1) / warmup_steps_float`.\n            # At `step = warmup_steps_float - 1`, i.e , at the `warmup_steps_float`-th step, \n            #`learning_rate` is `self.initial_learning_rate`.\n            warmup_percent_done = (steps_float + 1) / warmup_steps_float\n            \n            warmup_learning_rate = self.initial_learning_rate * warmup_percent_done\n\n            is_warmup = tf.cast(steps_int < warmup_steps_int, tf.float32)\n            learning_rate = ((1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)\n                        \n        return learning_rate\n\n\nnum_train_steps = int(FLAGS.epochs * FLAGS.num_train_examples / FLAGS.train_batch_size)\nprint(f'num_train_steps: {num_train_steps}')\n\n\n# ----------------------------------------------------------------------------------------\n\nwith tpu_strategy.scope():\n\n    # Model\n    bert_tokenizer, bert_nq = get_pretrained_model(FLAGS.model_name)\n\n    # Metric\n    train_loss, train_loss_short_start_pos, train_loss_short_end_pos, train_loss_short_ans_type, train_acc, train_acc_short_start_pos, train_acc_short_end_pos, train_acc_short_ans_type = get_metrics(\"train\")\n\n    # Loss\n    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n\n\n    def loss_function(nq_labels, nq_logits):\n        (short_start_pos_labels, short_end_pos_labels, short_answer_type_labels) = nq_labels\n        (short_start_pos_logits, short_end_pos_logits, short_answer_type_logits) = nq_logits\n\n        loss_short_start_pos = loss_object(short_start_pos_labels, short_start_pos_logits)\n        loss_short_end_pos = loss_object(short_end_pos_labels, short_end_pos_logits)\n        loss_short_ans_type = loss_object(short_answer_type_labels, short_answer_type_logits)\n\n        loss_short_start_pos = tf.nn.compute_average_loss(loss_short_start_pos, global_batch_size=FLAGS.train_batch_size)\n        loss_short_end_pos = tf.nn.compute_average_loss(loss_short_end_pos, global_batch_size=FLAGS.train_batch_size)\n        loss_short_ans_type = tf.nn.compute_average_loss(loss_short_ans_type, global_batch_size=FLAGS.train_batch_size)\n\n        loss = (loss_short_start_pos + loss_short_end_pos + loss_short_ans_type) / 3.0\n\n        return loss, loss_short_start_pos, loss_short_end_pos, loss_short_ans_type\n\n\n    learning_rate = CustomSchedule(\n        initial_learning_rate=FLAGS.init_learning_rate,\n        decay_steps=num_train_steps,\n        end_learning_rate=FLAGS.init_learning_rate,\n        power=1.0,\n        cycle=FLAGS.cyclic_learning_rate,\n        num_warmup_steps=FLAGS.num_warmup_steps\n    )\n\n    decay_var_list = []\n    for i in range(len(bert_nq.trainable_variables)):\n        name = bert_nq.trainable_variables[i].name\n        if any(x in name for x in [\"LayerNorm\", \"layer_norm\", \"bias\"]):\n            decay_var_list.append(name)\n\n    optimizer = AdamW(weight_decay=FLAGS.init_weight_decay_rate, learning_rate=learning_rate, beta_1=0.9, beta_2=0.999, epsilon=1e-6, decay_var_list=decay_var_list)\n\n    input_signature = [\n        tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n        tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n        tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n        tf.TensorSpec(shape=(None,), dtype=tf.int32),\n        tf.TensorSpec(shape=(None,), dtype=tf.int32),\n        tf.TensorSpec(shape=(None,), dtype=tf.int32)\n    ]\n\n\n    @tf.function(input_signature=input_signature)\n    def train_step(input_ids, input_masks, segment_ids, short_start_pos_labels, short_end_pos_labels, short_answer_type_labels):\n\n        nq_inputs = (input_ids, input_masks, segment_ids)\n        nq_labels = (short_start_pos_labels, short_end_pos_labels, short_answer_type_labels)\n\n        with tf.GradientTape() as tape:\n\n            nq_logits = bert_nq(nq_inputs, training=True)\n            loss, loss_short_start_pos, loss_short_end_pos, loss_short_ans_type = loss_function(nq_labels, nq_logits)\n\n        gradients = tape.gradient(loss, bert_nq.trainable_variables)\n\n        (short_start_pos_logits, short_end_pos_logits, short_answer_type_logits) = nq_logits\n\n        train_acc.update_state(short_start_pos_labels, short_start_pos_logits)\n        train_acc.update_state(short_end_pos_labels, short_end_pos_logits)\n        train_acc.update_state(short_answer_type_labels, short_answer_type_logits)\n\n        train_acc_short_start_pos.update_state(short_start_pos_labels, short_start_pos_logits)\n        train_acc_short_end_pos.update_state(short_end_pos_labels, short_end_pos_logits)\n        train_acc_short_ans_type.update_state(short_answer_type_labels, short_answer_type_logits)\n\n        optimizer.apply_gradients(zip(gradients, bert_nq.trainable_variables))\n\n        train_loss(loss)\n\n        train_loss_short_start_pos(loss_short_start_pos)\n        train_loss_short_end_pos(loss_short_end_pos)\n        train_loss_short_ans_type(loss_short_ans_type)\n\n\n    # `experimental_run_v2` replicates the provided computation and runs it with the distributed input.\n    @tf.function\n    def distributed_train_step(dataset_inputs):\n\n        features, targets = dataset_inputs\n        (input_ids, input_masks, segment_ids) = (features['input_ids'], features['input_mask'], features['segment_ids'])\n        (short_start_pos_labels, short_end_pos_labels, short_answer_type_labels) = (targets['short_start_positions'], targets['short_end_positions'], targets['short_answer_types'])\n\n        tpu_strategy.experimental_run_v2(train_step, args=(input_ids, input_masks, segment_ids, short_start_pos_labels, short_end_pos_labels, short_answer_type_labels))\n\ncheckpoint_path = FLAGS.input_checkpoint_dir + FLAGS.model_name + \"/\"\nckpt = tf.train.Checkpoint(model=bert_nq)\nckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=10000)\n\n# if a checkpoint exists, restore the latest checkpoint.\nif ckpt_manager.latest_checkpoint:\n    ckpt.restore(ckpt_manager.latest_checkpoint)\n    last_epoch = int(ckpt_manager.latest_checkpoint.split(\"-\")[-1])\n    print (f'Latest BertNQ checkpoint restored -- Model trained for {last_epoch} epochs')\nelse:\n    print('Checkpoint not found. Train BertNQ from scratch')\n    last_epoch = 0\n\n\ntrain_start_time = datetime.datetime.now()\n\nepochs = FLAGS.epochs\nfor epoch in range(epochs):\n\n    print(\"Epoch = {}\".format(epoch))\n\n    train_dataset = get_dataset(\n        FLAGS.train_tf_record,\n        FLAGS.max_seq_length_for_training,\n        FLAGS.train_batch_size,\n        FLAGS.shuffle_buffer_size,\n        is_training=True\n    )\n\n    print(\"train_dataset is OK.\")\n\n    train_dist_dataset = tpu_strategy.experimental_distribute_dataset(train_dataset)\n\n    print(\"train_dist_dataset is OK.\")\n\n    train_loss.reset_states()\n    \n    train_loss_short_start_pos.reset_states()\n    train_loss_short_end_pos.reset_states()\n    train_loss_short_ans_type.reset_states()           \n    \n    train_acc.reset_states()\n    \n    train_acc_short_start_pos.reset_states()\n    train_acc_short_end_pos.reset_states()\n    train_acc_short_ans_type.reset_states()\n    \n    epoch_start_time = datetime.datetime.now()\n\n    print(\"start iterating over train_dist_dataset ...\")\n\n    for (batch_idx, dataset_inputs) in enumerate(train_dist_dataset):\n\n        batch_start_time = datetime.datetime.now()\n\n        distributed_train_step(dataset_inputs)\n\n        batch_end_time = datetime.datetime.now()\n        batch_elapsed_time = (batch_end_time - batch_start_time).total_seconds()\n        \n        if (batch_idx + 1) % 100 == 0:\n            print('Epoch {} | Batch {} | Elapsed Time {}'.format(\n                epoch + 1,\n                batch_idx + 1,\n                batch_elapsed_time\n            ))\n            print('Loss {:.6f} | Loss_SS {:.6f} | Loss_SE {:.6f} | Loss_ST {:.6f}'.format(\n                train_loss.result(),\n                train_loss_short_start_pos.result(),\n                train_loss_short_end_pos.result(),\n                train_loss_short_ans_type.result()\n            ))\n            print(' Acc {:.6f} |  Acc_SS {:.6f} |  Acc_SE {:.6f} |  Acc_ST {:.6f}'.format(\n                train_acc.result(),\n                train_acc_short_start_pos.result(),\n                train_acc_short_end_pos.result(),\n                train_acc_short_ans_type.result()                \n            ))\n            print(\"-\" * 100)\n\n    epoch_end_time = datetime.datetime.now()\n    epoch_elapsed_time = (epoch_end_time - epoch_start_time).total_seconds()\n            \n    if (epoch + 1) % 1 == 0:\n        \n        ckpt_save_path = ckpt_manager.save()\n        print ('\\nSaving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path))\n        \n        print('\\nEpoch {}'.format(epoch + 1))\n        print('Loss {:.6f} | Loss_SS {:.6f} | Loss_SE {:.6f} | Loss_ST {:.6f}'.format(\n                train_loss.result(),\n                train_loss_short_start_pos.result(),\n                train_loss_short_end_pos.result(),\n                train_loss_short_ans_type.result()\n        ))\n        print(' Acc {:.6f} |  Acc_SS {:.6f} |  Acc_SE {:.6f} |  Acc_ST {:.6f}'.format(\n                train_acc.result(),\n                train_acc_short_start_pos.result(),\n                train_acc_short_end_pos.result(),\n                train_acc_short_ans_type.result() \n        ))\n\n    print('\\nTime taken for 1 epoch: {} secs\\n'.format(epoch_elapsed_time))\n    print(\"-\" * 80 + \"\\n\")\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"","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":1}