{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Use Hugging Face's Tensorflow 2 transformer models for NQ - Inference"},{"metadata":{},"cell_type":"markdown","source":"# Introduction\nThis inference kernel is the continuation of my TF2 training kerenl [Use Hugging Face's Tensorflow 2 transformer models for NQ](https://www.kaggle.com/yihdarshieh/use-hugging-face-s-tensorflow-2-transformer-models). These demonstrate how to use Hugging Face's [transformers](https://github.com/huggingface/transformers) package, more precisely, theier `Tensorflow 2` models, for this competition."},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","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 # For using flags without tf.compat.v1.flags.Flag\nimport datetime\n\nprint(f'tf={tf.__version__}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.show_versions()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Input data files are available in the \"../input/\" directory.\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# If you want to use your own .tfrecord or new trained checkpoints, you can put them under you own nq dir (`MY_OWN_NQ_DIR`)\n# Default to NQ_DIR. You have to change it to the dir containing your own working files.\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# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\nfor dirname, _, filenames in os.walk(INPUT_DIR):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# Any results you write to the current directory are saved as output.\n\nINPUT_DIR","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# HuggingFace pretrained Bert model names"},{"metadata":{"trusted":true},"cell_type":"code","source":"PRETRAINED_MODELS = {\n    \"BERT\": [\n        'bert-base-uncased',\n        'bert-large-uncased',\n        'bert-base-cased',\n        'bert-large-cased',\n        'bert-base-multilingual-uncased',\n        'bert-base-multilingual-cased',\n        'bert-base-chinese',\n        'bert-base-german-cased',\n        'bert-large-uncased-whole-word-masking',\n        'bert-large-cased-whole-word-masking',\n        'bert-large-uncased-whole-word-masking-finetuned-squad',\n        'bert-large-cased-whole-word-masking-finetuned-squad',\n        'bert-base-cased-finetuned-mrpc'\n    ],\n    \"DISTILBERT\": [\n        'distilbert-base-uncased',\n        'distilbert-base-uncased-distilled-squad'\n    ]\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Abseil Flags - Datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"flags = 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.\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Abseil Flags - Hyperparameters and file paths"},{"metadata":{"trusted":true},"cell_type":"code","source":"if os.path.isfile(os.path.join(MY_OWN_NQ_DIR, \"nq_train.tfrecord\")):\n    TRAIN_TF_RECORD = os.path.join(MY_OWN_NQ_DIR, \"nq_train.tfrecord\")\nelif os.path.isfile(os.path.join(MY_OWN_NQ_DIR, \"nq-train.tfrecords-00000-of-00001\")):\n    TRAIN_TF_RECORD = os.path.join(MY_OWN_NQ_DIR, \"nq-train.tfrecords-00000-of-00001\")\nelse:\n    TRAIN_TF_RECORD = os.path.join(BERT_JOINT_BASE_DIR, \"nq-train.tfrecords-00000-of-00001\")\n    \nflags.DEFINE_string(\"train_tf_record\", TRAIN_TF_RECORD,\n                    \"Precomputed tf records for training dataset.\")\n\nflags.DEFINE_string(\"valid_tf_record\", os.path.join(NQ_DIR, \"nq_valid.tfrecord\"),\n                    \"Precomputed tf records for validation dataset.\")\n\nflags.DEFINE_string(\"valid_small_tf_record\", os.path.join(NQ_DIR, \"nq_valid_small.tfrecord\"),\n                    \"Precomputed tf records for a smaller validation dataset.\")\n\nflags.DEFINE_string(\"valid_tf_record_with_labels\", \"nq_valid_with_labels.tfrecord\",\n                    \"Precomputed tf records for validation dataset with labels.\")\n\nflags.DEFINE_string(\"valid_small_tf_record_with_labels\", \"nq_valid_small_with_labels.tfrecord\",\n                    \"Precomputed tf records for a smaller validation dataset with labels.\")\n\n# This file should be generated when the kernel is running using the provided test dataset!\nflags.DEFINE_string(\"test_tf_record\", \"nq_test.tfrecord\",\n                    \"Precomputed tf records for test dataset.\")\n\nflags.DEFINE_bool(\"do_train\", False, \"Whether to run training dataset.\")\n\nflags.DEFINE_bool(\"do_valid\", True, \"Whether to run validation dataset.\")\n\nflags.DEFINE_bool(\"smaller_valid_dataset\", True, \"Whether to use the smaller validation dataset\")\n\nflags.DEFINE_bool(\"do_predict\", True, \"Whether to run test dataset.\")\n\nflags.DEFINE_string(\n    \"validation_prediction_output_file\", \"validatioin_predictions.json\",\n    \"Where to print predictions for validation dataset in NQ prediction format, to be passed to natural_questions.nq_eval.\")\n\nflags.DEFINE_string(\n    \"validation_small_prediction_output_file\", \"validatioin_small_predictions.json\",\n    \"Where to print predictions for validation dataset in NQ prediction format, to be passed to natural_questions.nq_eval.\")\n\nflags.DEFINE_string(\n    \"prediction_output_file\", \"predictions.json\",\n    \"Where to print predictions for test dataset in NQ prediction format, to be passed to natural_questions.nq_eval.\")\n\nflags.DEFINE_string(\n    \"input_checkpoint_dir\", os.path.join(MY_OWN_NQ_DIR, \"checkpoints\"),\n    \"The root directory that contains checkpoints to be loaded of all trained models.\")\n\nflags.DEFINE_string(\n    \"output_checkpoint_dir\", \"checkpoints\",\n    \"The output directory where the model checkpoints will be written to.\")\n\n# If you want to use other Hugging Face's models, change this to `MY_OWN_NQ_DIR` and put the downloaded models at the right place.\nflags.DEFINE_string(\"model_dir\", NQ_DIR, \"Root dir of all Hugging Face's models\")\n\n#flags.DEFINE_string(\"model_name\", \"distilbert-base-uncased-distilled-squad\", \"Name of Hugging Face's model to use.\")\n# flags.DEFINE_string(\"model_name\", \"bert-base-uncased\", \"Name of Hugging Face's model to use.\")\nflags.DEFINE_string(\"model_name\", \"bert-large-uncased-whole-word-masking-finetuned-squad\", \"Name of Hugging Face's model to use.\")\n\nflags.DEFINE_integer(\"epochs\", 40, \"Total epochs for training.\")\n\nflags.DEFINE_integer(\"train_batch_size\", 5, \"Batch size for training.\")\n\nflags.DEFINE_integer(\"shuffle_buffer_size\", 10000, \"Shuffle buffer size for training.\")\n\nflags.DEFINE_integer(\"batch_accumulation_size\", 100, \"Number of batches to accumulate gradient before applying optimization.\")\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\nflags.DEFINE_integer(\"predict_batch_size\", 25, \"Batch size for predictions.\")\n\n# ----------------------------------------------------------------------------------------\nflags.DEFINE_integer(\n    \"n_best_size\", 20,\n    \"The total number of n-best predictions to generate in the \"\n    \"nbest_predictions.json output file.\")\n\nflags.DEFINE_integer(\n    \"max_answer_length\", 30,\n    \"The maximum length of an answer that can be generated. This is needed \"\n    \"because the start and end predictions are not conditioned on one another.\")\n\nflags.DEFINE_string(\n    \"validation_predict_file\", os.path.join(NQ_DIR, \"simplified-nq-dev.jsonl\"),\n    \"\")\n\nflags.DEFINE_string(\n    \"validation_predict_file_small\", os.path.join(NQ_DIR, \"simplified-nq-dev-small.jsonl\"),\n    \"\")\n\n# ----------------------------------------------------------------------------------------\n## Special flags - do not change\n\nif IS_KAGGLE: \n    flags.DEFINE_string(\n        \"predict_file\", \"/kaggle/input/tensorflow2-question-answering/simplified-nq-test.jsonl\",\n        \"NQ json for predictions. E.g., dev-v1.1.jsonl.gz or test-v1.1.jsonl.gz\")\nelse:\n    flags.DEFINE_string(\n        \"predict_file\", os.path.join(NQ_DIR, \"simplified-nq-test.jsonl\"),\n        \"NQ json for predictions. E.g., dev-v1.1.jsonl.gz or test-v1.1.jsonl.gz\")\n    \nif IS_KAGGLE: \n    flags.DEFINE_string(\n        \"sample_submission_csv\", \"/kaggle/input/tensorflow2-question-answering/sample_submission.csv\",\n        \"path to sample submission csv file.\")\nelse:\n    flags.DEFINE_string(\n        \"sample_submission_csv\", os.path.join(NQ_DIR, \"sample_submission.csv\"),\n        \"path to sample submission csv file.\")    \n    \nflags.DEFINE_boolean(\"logtostderr\", True, \"Logs to stderr\")\nflags.DEFINE_boolean(\"undefok\", True, \"it's okay to be undefined\")\nflags.DEFINE_string('f', '', 'kernel')\nflags.DEFINE_string('HistoryManager.hist_file', '', 'kernel')\n\n# Make the default flags as parsed flags\nFLAGS.mark_as_parsed()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IS_SUBMITTING = False\n\ntest_answers_df = pd.read_csv(FLAGS.sample_submission_csv)\n\nif IS_KAGGLE and len(test_answers_df) != 692:\n    IS_SUBMITTING = True\n    FLAGS.do_train = False\n    FLAGS.do_valid = False\n    FLAGS.do_predict = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(test_answers_df))\nprint(IS_SUBMITTING)\nprint(FLAGS.do_train)\nprint(FLAGS.do_valid)\nprint(FLAGS.do_predict)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"NB_ANSWER_TYPES = 5","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Make TF record file for test dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def jsonl_iterator(jsonl_files, to_json=False):\n\n    for file_path in jsonl_files:\n        with open(file_path, \"r\", encoding=\"UTF-8\") as fp:\n            for jsonl in fp:\n                raw_example = jsonl\n                if to_json:\n                    raw_example = json.loads(jsonl)\n                yield raw_example\n\n                \n# Convert test examples to tf records.\ncreator = TFExampleCreator(is_training=False)\nnq_lines = jsonl_iterator([FLAGS.predict_file])\ncreator.process_nq_lines(nq_lines=nq_lines, output_tfrecord=FLAGS.test_tf_record, max_examples=0, collect_nq_features=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Get Datasets from TF Record files"},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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            \"token_map\": tf.io.FixedLenFeature([seq_length], tf.int64)\n        }        \n\n    # Taken from the TensorFlow models repository: https://github.com/tensorflow/models/blob/befbe0f9fe02d6bc1efb1c462689d069dae23af1/official/nlp/bert/input_pipeline.py#L24\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                        \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 not is_training:\n            x['token_map'] = record['token_map']\n\n        if is_training:\n            y = {\n                'start_positions': record['start_positions'],\n                'end_positions': record['end_positions'],\n                '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)\n    \n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Check a small batch in validation / test datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_tf_record = FLAGS.valid_tf_record\nif FLAGS.smaller_valid_dataset:\n    valid_tf_record = FLAGS.valid_small_tf_record\n    \nvalid_tf_record_with_labels = FLAGS.valid_tf_record_with_labels\nif FLAGS.smaller_valid_dataset:\n    valid_tf_record_with_labels = FLAGS.valid_small_tf_record_with_labels\n\ntrain_dataset = get_dataset(FLAGS.train_tf_record,\n                    seq_length=FLAGS.max_seq_length_for_training,\n                    batch_size=2,\n                    shuffle_buffer_size=FLAGS.shuffle_buffer_size,\n                    is_training=True\n                )    \n\nvalidation_dataset = get_dataset(os.path.join(NQ_DIR, valid_tf_record),\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=2,\n                         is_training=False\n                     )\n\nvalidation_dataset_with_labels = get_dataset(os.path.join(NQ_DIR, valid_tf_record_with_labels),\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=2,\n                         is_training=True\n                     )\n\ntest_dataset = get_dataset(FLAGS.test_tf_record,\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=2,\n                         is_training=False\n                     )\n\n# Can't use next(train_dataset)!\nfeatures, targets = next(iter(train_dataset))\nprint(features)\nprint(targets)\n\nfeatures = next(iter(validation_dataset))\nprint(features)\n\nfeatures, labels = next(iter(validation_dataset_with_labels))\nprint(features)\n\nfeatures = next(iter(test_dataset))\nprint(features)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Get the actual validation / test datasets from TF Records"},{"metadata":{"trusted":true},"cell_type":"code","source":"validation_dataset = get_dataset(os.path.join(NQ_DIR, valid_tf_record),\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=FLAGS.predict_batch_size,\n                         is_training=False\n                     )\n\nvalidation_dataset_with_labels = get_dataset(os.path.join(NQ_DIR, valid_tf_record_with_labels),\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=FLAGS.predict_batch_size,\n                         is_training=True\n                     )\n\ntest_dataset = get_dataset(FLAGS.test_tf_record,\n                         seq_length=FLAGS.max_seq_length,\n                         batch_size=FLAGS.predict_batch_size,\n                         is_training=False\n                     )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Get validation / test features"},{"metadata":{"trusted":true},"cell_type":"code","source":"# r is a string Tensor, we use r.numpy() to get underlying byte string\nvalidation_features = (tf.train.Example.FromString(r.numpy()) for r in tf.data.TFRecordDataset(valid_tf_record))\ntest_features = (tf.train.Example.FromString(r.numpy()) for r in tf.data.TFRecordDataset(FLAGS.test_tf_record))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# An interface  for NQ models"},{"metadata":{"trusted":true},"cell_type":"code","source":"from transformers import BertTokenizer\nfrom transformers import TFBertModel, TFDistilBertModel\nfrom transformers import TFBertMainLayer, TFDistilBertMainLayer, TFBertPreTrainedModel, TFDistilBertPreTrainedModel\nfrom transformers.modeling_tf_utils import get_initializer\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.pos_classifier = tf.keras.layers.Dense(2,\n                                        kernel_initializer=get_initializer(config.initializer_range),\n                                        name='pos_classifier')       \n\n        self.answer_type_classifier = tf.keras.layers.Dense(NB_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        pos_logits = self.pos_classifier(sequence_output)  # shape = (batch_size, seq_len, 2)\n        start_pos_logits = pos_logits[:, :, 0]  # shape = (batch_size, seq_len)\n        end_pos_logits = pos_logits[:, :, 1]  # shape = (batch_size, seq_len)\n        \n        answer_type_logits = self.answer_type_classifier(pooled_output)  # shape = (batch_size, NB_ANSWER_TYPES)\n\n        outputs = (start_pos_logits, end_pos_logits, 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        \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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Try Bert / DistillBert models for NQ"},{"metadata":{"trusted":true},"cell_type":"code","source":"bert_tokenizer, bert_for_nq = get_pretrained_model('bert-base-uncased')\n_, distil_bert_for_nq = get_pretrained_model('distilbert-base-uncased-distilled-squad')\n\ninput_ids = tf.constant(bert_tokenizer.encode(\"Hello, my dog is cute\"))[None, :]  # Batch size 1\ninput_masks = tf.constant(0, shape=input_ids.shape)\nsegment_ids = tf.constant(0, shape=input_ids.shape)\n\n# Actual inputs to model\ninputs = (input_ids, input_masks, segment_ids)\n\n# Outputs from bert_for_nq using backend_call()\noutputs = bert_for_nq(inputs)\n(start_pos_logits, end_pos_logits, answer_type_logits) = outputs\nprint(start_pos_logits.shape)\nprint(end_pos_logits.shape)\nprint(answer_type_logits.shape)\n\nlen(bert_for_nq.trainable_variables)\n\n# Outputs from distil_bert_for_nq using backend_call()\noutputs = distil_bert_for_nq(inputs)\n(start_pos_logits, end_pos_logits, answer_type_logits) = outputs\nprint(start_pos_logits.shape)\nprint(end_pos_logits.shape)\nprint(answer_type_logits.shape)\n\nlen(distil_bert_for_nq.trainable_variables)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Choose the model to use"},{"metadata":{"trusted":true},"cell_type":"code","source":"bert_tokenizer, bert_nq = get_pretrained_model(FLAGS.model_name)\n\nif not IS_KAGGLE:\n    bert_nq.trainable_variables","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load the latest checkpoint"},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint_path = os.path.join(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=5)\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    \n# Reset saving path, because the FLAGS.input_checkpoint_dir is not writable on Kaggle\nprint(ckpt_manager._directory)\nckpt_manager._directory = os.path.join(FLAGS.output_checkpoint_dir, FLAGS.model_name)\nckpt_manager._checkpoint_prefix = os.path.join(ckpt_manager._directory, \"ckpt\")\nprint(ckpt_manager._directory)\n\nfrom tensorflow.python.lib.io.file_io import recursive_create_dir\nrecursive_create_dir(ckpt_manager._directory)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Evaluation"},{"metadata":{"trusted":true},"cell_type":"code","source":"Span = collections.namedtuple(\"Span\", [\"start_token_idx\", \"end_token_idx\"])\n\n\nclass EvalExample(object):\n    \"\"\"Eval data available for a single example.\"\"\"\n\n    def __init__(self, example_id, candidates):\n        self.example_id = example_id\n        self.candidates = candidates\n        self.results = {}\n        self.features = {}\n\n\nclass ScoreSummary(object):\n\n    def __init__(self):\n        self.predicted_label = None\n        self.short_span_score = None\n        self.cls_token_score = None\n        self.answer_type_logits = None\n        self.start_prob = None\n        self.end_prob = None\n        self.answer_type_prob_dist = None\n\n        \ndef read_candidates_from_one_split(input_path):\n    \"\"\"Read candidates from a single jsonl file.\"\"\"\n    candidates_dict = {}\n    if input_path.endswith(\".gz\"):\n        with gzip.GzipFile(fileobj=tf.io.gfile.GFile(input_path, \"rb\")) as input_file:\n            print(\"Reading examples from: {}\".format(input_path))\n            for index, line in enumerate(input_file):\n                e = json.loads(line)\n                candidates_dict[e[\"example_id\"]] = e[\"long_answer_candidates\"]\n                # if index > 100:\n                #     break\n    else:\n        with tf.io.gfile.GFile(input_path, \"r\") as input_file:\n            print(\"Reading examples from: {}\".format(input_path))\n            for index, line in enumerate(input_file):\n                e = json.loads(line)\n                candidates_dict[e[\"example_id\"]] = e[\"long_answer_candidates\"]\n                # if index > 100:\n                #     break\n\n\n    return candidates_dict\n\n\ndef read_candidates(input_pattern):\n    \"\"\"Read candidates with real multiple processes.\"\"\"\n    input_paths = tf.io.gfile.glob(input_pattern)\n    final_dict = {}\n    for input_path in input_paths:\n        final_dict.update(read_candidates_from_one_split(input_path))\n    return final_dict\n\n\ndef get_best_indexes(logits, n_best_size, token_map=None):\n    # Return a sorted list of (idx, logit)\n    index_and_score = sorted(enumerate(logits[1:], 1), key=lambda x: x[1], reverse=True)\n        \n    best_indexes = []\n    for i in range(len(index_and_score)):\n        \n        idx = index_and_score[i][0]\n\n        if token_map is not None and token_map[idx] == -1:\n            continue\n\n        best_indexes.append(idx)\n\n        if len(best_indexes) >= n_best_size:\n            break    \n    \n    return best_indexes\n\n\ndef compute_predictions(example):\n    \"\"\"Converts an example into an NQEval object for evaluation.\n    \n       Unlike the starter kernel, this returns a list of `ScoreSummary`, sorted by score.\n    \"\"\"\n    \n    predictions = []\n    max_answer_length = FLAGS.max_answer_length\n\n    for unique_id, result in example.results.items():\n        \n        if unique_id not in example.features:\n            raise ValueError(\"No feature found with unique_id:\", unique_id)\n        token_map = example.features[unique_id][\"token_map\"].int64_list.value\n        \n        for start_index, start_logit, start_prob in zip(result[\"start_indexes\"], result[\"start_logits\"], result[\"start_pos_prob_dist\"]):\n\n            if token_map[start_index] == -1:\n                continue            \n            \n            for end_index, end_logit, end_prob in zip(result[\"end_indexes\"], result[\"end_logits\"], result[\"end_pos_prob_dist\"]):\n\n                if token_map[end_index] == -1:\n                    continue\n\n                if end_index < start_index:\n                    continue                    \n                    \n                length = end_index - start_index + 1\n                if length > max_answer_length:\n                    continue\n                    \n                summary = ScoreSummary()\n                \n                summary.instance_id = unique_id\n                \n                summary.short_span_score = start_logit + end_logit\n                summary.cls_token_score = result[\"cls_start_logit\"] + result[\"cls_end_logit\"]\n                summary.answer_type_logits = result[\"answer_type_logits\"]\n                \n                summary.start_indexes = result[\"start_indexes\"]\n                summary.end_indexes = result[\"end_indexes\"]\n\n                summary.start_logits = result[\"start_logits\"]\n                summary.end_logits = result[\"end_logits\"]                \n                \n                summary.start_pos_prob_dist = result[\"start_pos_prob_dist\"]\n                summary.end_pos_prob_dist = result[\"end_pos_prob_dist\"]                \n                           \n                summary.start_index = start_index\n                summary.end_index = end_index\n                \n                summary.start_logit = start_logit\n                summary.end_logit = end_logit\n                \n                answer_type_prob_dist = result[\"answer_type_prob_dist\"]\n                summary.start_prob = start_prob\n                summary.end_prob = end_prob\n                summary.answer_type_prob_dist = {\n                    \"unknown\": answer_type_prob_dist[0],\n                    \"yes\": answer_type_prob_dist[1],\n                    \"no\": answer_type_prob_dist[2],\n                    \"short\": answer_type_prob_dist[3],\n                    \"long\": answer_type_prob_dist[4]\n                }\n                \n                start_span = token_map[start_index]\n                end_span = token_map[end_index] + 1\n\n                # Span logits minus the cls logits seems to be close to the best.\n                score = summary.short_span_score - summary.cls_token_score\n                predictions.append((score, summary, start_span, end_span))\n                \n    all_summaries = []            \n                    \n    if predictions:\n        \n        predictions = sorted(predictions, key=lambda x: (x[0], x[2], x[3]), reverse=True)\n        \n        for prediction in predictions:\n            \n            long_span = Span(-1, -1)\n          \n            score, summary, start_span, end_span = prediction\n            short_span = Span(start_span, end_span)\n            for c in example.candidates:\n                start = short_span.start_token_idx\n                end = short_span.end_token_idx\n                if c[\"top_level\"] and c[\"start_token\"] <= start and c[\"end_token\"] >= end:\n                    long_span = Span(c[\"start_token\"], c[\"end_token\"])\n                    break\n\n            summary.predicted_label = {\n                    \"example_id\": example.example_id,\n                    \"instance_id\": summary.instance_id,\n                    \"long_answer\": {\n                            \"start_token\": long_span.start_token_idx,\n                            \"end_token\": long_span.end_token_idx\n                    },\n                    \"short_answers\": [{\n                            \"start_token\": short_span.start_token_idx,\n                            \"end_token\": short_span.end_token_idx\n                    }],\n                    \"yes_no_answer\": \"NONE\",\n                    \"long_answer_score\": score,                \n                    \"short_answers_score\": score,                \n                    \"answer_type_prob_dist\": summary.answer_type_prob_dist,\n                    \"start_index\": summary.start_index,\n                    \"end_index\": summary.end_index,\n                    \"start_logit\": summary.start_logit,\n                    \"end_logit\": summary.end_logit,\n                    \"start_prob\": summary.start_prob,\n                    \"end_prob\": summary.end_prob,\n                    \"start_indexes\": summary.start_indexes,\n                    \"end_indexes\": summary.end_indexes,\n                    \"start_logits\": summary.start_logits,\n                    \"end_logits\": summary.end_logits,\n                    \"start_pos_prob_dist\": summary.start_pos_prob_dist,\n                    \"end_pos_prob_dist\": summary.end_pos_prob_dist\n            }\n            \n            all_summaries.append(summary)\n\n    if len(all_summaries) == 0:\n\n        short_span = Span(-1, -1)\n        long_span = Span(-1, -1)\n        score = 0\n        summary = ScoreSummary()        \n        \n        summary.predicted_label = {\n                \"example_id\": example.example_id,\n                \"instance_id\": None,\n                \"long_answer\": {\n                        \"start_token\": long_span.start_token_idx,\n                        \"end_token\": long_span.end_token_idx,\n                        \"start_byte\": -1,\n                        \"end_byte\": -1\n                },\n                \"long_answer_score\": score,\n                \"short_answers\": [{\n                        \"start_token\": short_span.start_token_idx,\n                        \"end_token\": short_span.end_token_idx,\n                        \"start_byte\": -1,\n                        \"end_byte\": -1\n                }],\n                \"short_answers_score\": score,\n                \"yes_no_answer\": \"NONE\"\n        }        \n        \n        all_summaries.append(summary)\n            \n    all_summaries = all_summaries[:min(FLAGS.n_best_size, len(all_summaries))]        \n    \n    return all_summaries\n\n\ndef compute_pred_dict(candidates_dict, dev_features, raw_results):\n    \"\"\"Computes official answer key from raw logits.\n    \n       Unlike the starter kernel, each nq_pred_dict[example_id] is a list of `predicted_label`\n       that is defined in `compute_predictions`.\n    \"\"\"\n\n    raw_results_by_id = [(int(res[\"unique_id\"]), 1, res, None) for res in raw_results]\n\n    examples_by_id = [(int(tf.cast(int(k), dtype=tf.int32)), 0, v, k) for k, v in candidates_dict.items()]\n    \n    features_by_id = [(int(tf.cast(f.features.feature[\"unique_ids\"].int64_list.value[0], dtype=tf.int32)), 2, f.features.feature, None) for f in dev_features]\n    \n    print('merging examples...')\n    merged = sorted(examples_by_id + raw_results_by_id + features_by_id)\n    print('done.')\n    \n    examples = []\n    for idx, type_, datum, orig_example_id in merged:\n        if type_ == 0: # Here, datum the list `long_answer_candidates`\n            examples.append(EvalExample(orig_example_id, datum))\n        elif type_ == 2: # Here, datum is a feature with `token_map`\n            examples[-1].features[idx] = datum\n        else: # Here, datum is a raw_result given by the model\n            examples[-1].results[idx] = datum    \n    \n    # Construct prediction objects.\n    summary_dict = {}\n    nq_pred_dict = {}\n    for e in examples:\n        \n        all_summaries = compute_predictions(e)\n        summary_dict[e.example_id] = all_summaries\n        nq_pred_dict[e.example_id] = [summary.predicted_label for summary in all_summaries]\n        if len(nq_pred_dict) % 100 == 0:\n            print(\"Examples processed: %d\" % len(nq_pred_dict))\n\n    return nq_pred_dict","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Get logits from model - Save to a json file"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_prediction_json(mode, max_nb_pos_logits=-1):\n    \n    if mode == 'valid':\n        dataset = validation_dataset\n        if FLAGS.smaller_valid_dataset:\n            predict_file = FLAGS.validation_predict_file_small\n            prediction_output_file = FLAGS.validation_small_prediction_output_file\n        else:\n            predict_file = FLAGS.validation_predict_file\n            prediction_output_file = FLAGS.validation_prediction_output_file\n        eval_features = validation_features\n    else:\n        dataset = test_dataset\n        predict_file = FLAGS.predict_file\n        eval_features = test_features\n        prediction_output_file = FLAGS.prediction_output_file\n    \n    print(predict_file)\n    print(prediction_output_file)\n    \n    all_results = []\n\n    prediction_start_time = datetime.datetime.now()\n\n    for (batch_idx, features) in enumerate(dataset):\n\n        batch_start_time = datetime.datetime.now()\n\n        unique_ids = features['unique_ids']\n        token_maps = features['token_map']       \n        \n        (input_ids, input_masks, segment_ids) = (features['input_ids'], features['input_mask'], features['segment_ids'])\n        \n        nq_inputs = (input_ids, input_masks, segment_ids)\n        nq_logits = bert_nq(nq_inputs, training=False)\n\n        (start_pos_logits, end_pos_logits, answer_type_logits) = nq_logits\n\n        unique_ids = unique_ids.numpy().tolist()\n        \n        token_maps = token_maps.numpy().tolist()\n        \n        start_pos_prob_dist = tf.nn.softmax(start_pos_logits, axis=-1).numpy().tolist()\n        end_pos_prob_dist = tf.nn.softmax(end_pos_logits, axis=-1).numpy().tolist()\n        answer_type_prob_dist = tf.nn.softmax(answer_type_logits, axis=-1).numpy().tolist()\n        \n        start_pos_logits = start_pos_logits.numpy().tolist()\n        end_pos_logits = end_pos_logits.numpy().tolist()\n        answer_type_logits = answer_type_logits.numpy().tolist()\n\n        for uid, token_map, s, e, a, sp, ep, ap in zip(unique_ids, token_maps, start_pos_logits, end_pos_logits, answer_type_logits, start_pos_prob_dist, end_pos_prob_dist, answer_type_prob_dist):\n\n            if max_nb_pos_logits < 0:\n                max_nb_pos_logits = len(start_pos_logits)\n            \n            # full_start_logits = s\n            # full_end_logits = e\n            \n            cls_start_logit = s[0]\n            cls_end_logit = e[0]\n            \n            start_indexes = get_best_indexes(s, max_nb_pos_logits, token_map)\n            end_indexes = get_best_indexes(e, max_nb_pos_logits, token_map)            \n            \n            s = [s[idx] for idx in start_indexes]\n            e = [e[idx] for idx in end_indexes]\n            sp = [sp[idx] for idx in start_indexes]\n            ep = [ep[idx] for idx in end_indexes]            \n            \n            raw_result = {\n                \"unique_id\": uid,\n                \"start_indexes\": start_indexes,\n                \"end_indexes\": end_indexes,\n                \"start_logits\": s,\n                \"end_logits\": e,\n                \"answer_type_logits\": a,\n                \"start_pos_prob_dist\": sp,\n                \"end_pos_prob_dist\": ep,\n                \"answer_type_prob_dist\": ap,\n                \"cls_start_logit\": cls_start_logit,\n                \"cls_end_logit\": cls_end_logit\n                # \"full_start_logits\": full_start_logits,\n                # \"full_end_logits\": full_end_logits\n            }\n            all_results.append(raw_result)\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('Batch {} | Elapsed Time {}'.format(\n                batch_idx + 1,\n                batch_elapsed_time\n            ))\n      \n    prediction_end_time = datetime.datetime.now()\n    prediction_elapsed_time = (prediction_end_time - prediction_start_time).total_seconds()\n\n    print('\\nTime taken for prediction: {} secs\\n'.format(prediction_elapsed_time))\n    print(\"-\" * 80 + \"\\n\")\n\n    print(\"Going to candidates file\")\n    candidates_dict = read_candidates(predict_file)\n\n    print (\"setting up eval features\")\n    # eval_features = ...\n\n    print (\"compute_pred_dict\")\n    nq_pred_dict = compute_pred_dict(candidates_dict, eval_features, all_results)\n    \n    predictions_json = {\"predictions\": list(nq_pred_dict.values())}\n\n    print (\"writing json\")\n    with tf.io.gfile.GFile(prediction_output_file, \"w\") as f:\n        json.dump(predictions_json, f, indent=4)\n        \n    return predictions_json","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Run on validation dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"if FLAGS.do_valid:\n    \n    validation_features = (tf.train.Example.FromString(r.numpy()) for r in tf.data.TFRecordDataset(valid_tf_record))\n    valid_predictions_json = get_prediction_json(mode='valid', max_nb_pos_logits=FLAGS.n_best_size)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Code for metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import f1_score\n\ndef create_answer_from_token_indices(answer):\n    \n    if answer[\"start_token\"] == -1 or answer[\"end_token\"] == -1:\n        return \"\"\n    else:\n        return str(answer[\"start_token\"]) + \":\" + str(answer[\"end_token\"])\n    \ndef create_long_answer_from_1_pred(pred):\n    \"\"\"\n    Args:\n        pred: A `predicted_label` as defined in `compute_predictions`.\n    \n    Returns:\n        A string. It's either an empty string \"\" or a string of the form \"start_token:end_token\",\n        where start_token and end_token are string forms of integers.\n    \"\"\"\n    \n    long_answer = create_answer_from_token_indices(pred[\"long_answer\"])\n    \n    return long_answer\n    \n    \ndef create_short_answers_from_1_pred(pred):\n    \"\"\"\n    Args:\n        pred: A `predicted_label` as defined in `compute_predictions`.\n    \n    Returns:\n        A list of strings. Each element can be [\"\"], [\"YES\"], [\"NO\"] or an list of strings with\n        the form \"start_token:end_token\" as describe in `create_long_answer_from_1_pred`.\n    \"\"\"\n    \n    short_answers = []\n    \n    for predicted_short_answer in pred[\"short_answers\"]:\n        \n        short_answer = create_answer_from_token_indices(predicted_short_answer)\n\n        # Custom\n        if \"answer_type_prob_dist\" in pred:\n            if pred[\"answer_type_prob_dist\"][\"yes\"] > 0.5:\n                short_answer = \"YES\"\n            elif pred[\"answer_type_prob_dist\"][\"no\"] > 0.5:\n                short_answer = \"NO\"\n\n            if pred[\"answer_type_prob_dist\"][\"short\"] < 0.5 or pred[\"answer_type_prob_dist\"][\"unknown\"] > 0.5:\n                if short_answer not in [\"YES\", \"NO\"]:\n                    short_answer = \"\"\n    \n        short_answers.append(short_answer)\n    \n    return short_answers\n\n\ndef is_pred_ok(pred, annotations):\n    \"\"\"\n    Args:\n        pred: A `predicted_label` as defined in `compute_predictions`.\n        annotations: A list of annotations. See `simplified-nq-dev.jsonl` for the format.\n        \n    Returns:\n        has_long_label: bool\n        has_short_label: bool\n        has_long_pred: bool\n        has_short_pred: bool\n        is_long_pred_correct: bool\n        is_short_pred_correct: bool\n    \"\"\"    \n        \n    long_labels = []\n    \n    for annotation in annotations:\n        \n        long_label = create_answer_from_token_indices(annotation[\"long_answer\"])\n        long_labels.append(long_label)\n        \n    non_null_long_labels = [x for x in long_labels if x != \"\"]\n    has_long_label = len(non_null_long_labels) > 1\n    \n    long_pred = create_long_answer_from_1_pred(pred)\n    has_long_pred = (long_pred != \"\")\n    \n    short_label_lists = []\n    \n    for annotation in annotations:\n        \n        if len(annotation[\"short_answers\"]) == 0:\n            if annotation[\"yes_no_answer\"] == \"YES\":\n                short_label_lists.append([\"YES\"])\n            elif annotation[\"yes_no_answer\"] == \"NO\":\n                short_label_lists.append([\"NO\"])\n            else:\n                short_label_lists.append([\"\"])\n        else:\n            \n            short_labels = []\n            for anno_short_answer in annotation[\"short_answers\"]:\n                short_label = create_answer_from_token_indices(anno_short_answer)\n                short_labels.append(short_label)\n            \n            short_label_lists.append(short_labels)        \n    \n    non_null_short_label_lists = [x for x in short_label_lists if x != [\"\"]]\n    has_short_label = len(non_null_short_label_lists) > 1\n    \n    # It can be [\"\"], [\"YES\"], [\"NO\"] or [\"start_token:end_token\"].\n    short_preds = create_short_answers_from_1_pred(pred)\n    has_short_pred = (short_preds != [\"\"])\n    \n    is_long_pred_correct = False\n    is_short_pred_correct = False\n    \n    for long_label in long_labels:\n        \n        if has_long_label and long_label == \"\":\n            continue\n        if not has_long_label and long_label != \"\":\n            continue\n            \n        if long_pred == long_label:\n            is_long_pred_correct = True\n            break\n\n    for short_labels in short_label_lists:\n\n        if has_short_label and short_labels == [\"\"]:\n            continue\n        if not has_short_label and short_labels != [\"\"]:\n            continue        \n            \n        if has_short_label:\n            \n            if short_labels == [\"YES\"] or short_labels == [\"NO\"]:\n\n                if short_preds == short_labels:\n                    is_short_pred_correct = True\n                    break\n\n            else:\n\n                if short_preds[0] in short_labels:\n\n                    is_short_pred_correct = True\n                    break\n                        \n        else:\n            \n            if short_preds == short_labels:\n                is_short_pred_correct = True\n                break\n\n     \n    return has_long_label, has_short_label, has_long_pred, has_short_pred, is_long_pred_correct, is_short_pred_correct\n\n\ndef compute_f1_scores(predictions_json, gold_jsonl_file):\n    \n    predictions = predictions_json[\"predictions\"]\n    \n    golden_nq_lines = jsonl_iterator([gold_jsonl_file])\n    golden_dict = dict()\n    for nq_line in golden_nq_lines:\n        nq_data = json.loads(nq_line)\n        golden = dict()\n        golden[\"example_id\"] = nq_data[\"example_id\"]\n        golden[\"annotations\"] = nq_data[\"annotations\"]\n        golden_dict[golden[\"example_id\"]] = golden\n        \n    long_labels = []\n    long_preds = []\n    short_labels = []\n    short_preds = []\n\n    for preds in predictions:\n        \n        # Let's take only the 1st pred for now. We can play with multiple preds if we want.\n        pred = preds[0]\n        \n        example_id = pred[\"example_id\"]\n        assert example_id in golden_dict\n        golden = golden_dict[example_id]\n        assert example_id == golden[\"example_id\"]\n        \n        has_long_label, has_short_label, has_long_pred, has_short_pred, is_long_correct, is_short_correct = is_pred_ok(pred, golden[\"annotations\"])\n        \n        if has_long_label or has_long_pred:\n            if is_long_correct:\n                long_labels.append(1)\n                long_preds.append(1)\n            else:\n                long_labels.append(1)\n                long_preds.append(0)            \n        \n        if has_short_label or has_short_pred:        \n            if is_short_correct:\n                short_labels.append(1)\n                short_preds.append(1)\n            else:\n                short_labels.append(1)\n                short_preds.append(0)\n            \n    f1 = f1_score(long_labels + short_labels, long_preds + short_preds)\n    long_f1 = f1_score(long_labels, long_preds)\n    short_f1 = f1_score(short_labels, short_preds)\n\n    return f1, long_f1, short_f1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Get metrics for validation dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"if FLAGS.do_valid:\n\n    if FLAGS.smaller_valid_dataset:\n        predict_file = FLAGS.validation_predict_file_small\n    else:\n        predict_file = FLAGS.validation_predict_file\n\n    f1, long_f1, short_f1 = compute_f1_scores(valid_predictions_json, predict_file)\n\n    print(f\"      valid f1: {f1}\\n valid long_f1: {long_f1}\\nvalid short_f1: {short_f1}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Run on test dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"if FLAGS.do_predict:\n    test_features = (tf.train.Example.FromString(r.numpy()) for r in tf.data.TFRecordDataset(FLAGS.test_tf_record))\n    predictions_json = get_prediction_json(mode='test', max_nb_pos_logits=FLAGS.n_best_size)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submission files\n\n\nThe Bert model produces a confidence score, which the Kaggle metric does not use. You, however, can use that score to determine which answers get submitted. See the limits commented out in create_short_answer and create_long_answer below for an example.\n\nValues for confidence will range between 1.0 and 2.0."},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_long_answer(preds):\n    \"\"\"\n    Args:\n        preds: A list of `predicted_label` as defined in `compute_predictions`.\n    \n    Returns:\n        A string represented a long answer.\n    \"\"\"\n    \n    # Currently, return the long answer from the 1st pred in preds.\n    return create_long_answer_from_1_pred(preds[0])  \n    \n    \ndef create_short_answer(preds):\n    \"\"\"\n    Args:\n        pred: A list of `predicted_label` as defined in `compute_predictions`.\n    \n    Returns:\n        A string represented a short answer.\n    \"\"\"\n    \n    # Currently, return the short answer from the 1st pred in preds.\n    return create_short_answers_from_1_pred(preds[0])[0]\n\n\nif FLAGS.do_predict:\n\n    test_answers_df = pd.read_json(FLAGS.prediction_output_file)\n\n    test_answers_df[\"long_answer_score\"] = test_answers_df[\"predictions\"].apply(lambda q: q[0][\"long_answer_score\"])\n    test_answers_df[\"short_answer_score\"] = test_answers_df[\"predictions\"].apply(lambda q: q[0][\"short_answers_score\"])\n\n    test_answers_df[\"long_answer_score\"].describe()\n\n\n    # An example of what each sample's answers look like in prediction.json:\n    test_answers_df.predictions.values[0]\n\n    # We re-format the JSON answers to match the requirements for submission.\n    test_answers_df[\"long_answer\"] = test_answers_df[\"predictions\"].apply(create_long_answer)\n    test_answers_df[\"short_answer\"] = test_answers_df[\"predictions\"].apply(create_short_answer)\n    test_answers_df[\"example_id\"] = test_answers_df[\"predictions\"].apply(lambda q: str(q[0][\"example_id\"]))\n\n    long_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"long_answer\"]))\n    short_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"short_answer\"]))\n\n    # Then we add them to our sample submission. Recall that each sample has both a _long and _short entry in the sample submission, one for each type of answer.\n    sample_submission = pd.read_csv(FLAGS.sample_submission_csv)\n\n    long_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_long\")].apply(lambda q: long_answers[q[\"example_id\"].replace(\"_long\", \"\")], axis=1)\n    short_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_short\")].apply(lambda q: short_answers[q[\"example_id\"].replace(\"_short\", \"\")], axis=1)\n\n    sample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_long\"), \"PredictionString\"] = long_prediction_strings\n    sample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_short\"), \"PredictionString\"] = short_prediction_strings\n\n    # And finally, we write out our submission!\n    sample_submission.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for dirname, _, filenames in os.walk('.'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","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}