{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Using Tensorflow 2.0 with Bert on Natural Questions\n\nIn this notebook, we'll be using the Bert baseline for Tensorflow to create predictions for the Natural Questions test set. Note that this uses a model that has already been pre-trained - we're only doing inference here. A GPU is required, and this should take between 1-2 hours to run.\n\nThe original script can be found [here](https://github.com/google-research/language/blob/master/language/question_answering/bert_joint/run_nq.py).\nThe supporting modules were drawn from the [official Tensorflow model repository](https://github.com/tensorflow/models/tree/master/official/nlp). The bert-joint-baseline data is described [here](https://github.com/google-research/language/tree/master/language/question_answering/bert_joint).\n\n**Note:** This baseline uses code that was migrated from TF1.x. Be aware that it contains use of tf.compat.v1, which is not permitted to be eligible for [TF2.0 prizes in this competition](https://www.kaggle.com/c/tensorflow2-question-answering/overview/prizes). It is intended to be used as a starting point, but we're excited to see how much better you can do using TF2.0!"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport tensorflow as tf\nimport tf2_0_baseline_w_bert as tf2baseline\nimport bert_modeling as modeling\nimport bert_optimization as optimization\nimport bert_tokenization as tokenization\nimport json\n\ntf.compat.v1.disable_eager_execution()\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n# In this case, we've got some extra BERT model files under `/kaggle/input/bertjointbaseline`\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\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.","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### Tensorflow flags are variables that can be passed around within the TF system. Every flag below has some context provided regarding what the flag is and how it's used.\n\n#### Most of these can be changed as desired, with the exception of the Special Flags at the bottom, which _must_ stay as-is to work with the Kaggle back end."},{"metadata":{"trusted":true},"cell_type":"code","source":"def del_all_flags(FLAGS):\n    flags_dict = FLAGS._flags()\n    keys_list = [keys for keys in flags_dict]\n    for keys in keys_list:\n        FLAGS.__delattr__(keys)\n\ndel_all_flags(tf.compat.v1.app.flags.FLAGS)\n\nflags = tf.compat.v1.app.flags\n\nflags.DEFINE_string(\n    \"bert_config_file\", \"/kaggle/input/bertjointbaseline/bert_config.json\",\n    \"The config json file corresponding to the pre-trained BERT model. \"\n    \"This specifies the model architecture.\")\n\nflags.DEFINE_string(\"vocab_file\", \"/kaggle/input/bertjointbaseline/vocab-nq.txt\",\n                    \"The vocabulary file that the BERT model was trained on.\")\n\nflags.DEFINE_string(\n    \"output_dir\", \"outdir\",\n    \"The output directory where the model checkpoints will be written.\")\n\nflags.DEFINE_string(\"train_precomputed_file\", None,\n                    \"Precomputed tf records for training.\")\n\nflags.DEFINE_integer(\"train_num_precomputed\", None,\n                     \"Number of precomputed tf records for training.\")\n\nflags.DEFINE_string(\n    \"output_prediction_file\", \"predictions.json\",\n    \"Where to print predictions in NQ prediction format, to be passed to\"\n    \"natural_questions.nq_eval.\")\n\nflags.DEFINE_string(\n    \"init_checkpoint\", \"/kaggle/input/bertjointbaseline/bert_joint.ckpt\",\n    \"Initial checkpoint (usually from a pre-trained BERT model).\")\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\nflags.DEFINE_integer(\n    \"max_seq_length\", 384,\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_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\nflags.DEFINE_bool(\"do_train\", False, \"Whether to run training.\")\n\nflags.DEFINE_bool(\"do_predict\", True, \"Whether to run eval on the dev set.\")\n\nflags.DEFINE_integer(\"train_batch_size\", 32, \"Total batch size for training.\")\n\nflags.DEFINE_integer(\"predict_batch_size\", 8,\n                     \"Total batch size for predictions.\")\n\nflags.DEFINE_float(\"learning_rate\", 5e-5, \"The initial learning rate for Adam.\")\n\nflags.DEFINE_float(\"num_train_epochs\", 3.0,\n                   \"Total number of training epochs to perform.\")\n\nflags.DEFINE_float(\n    \"warmup_proportion\", 0.1,\n    \"Proportion of training to perform linear learning rate warmup for. \"\n    \"E.g., 0.1 = 10% of training.\")\n\nflags.DEFINE_integer(\"save_checkpoints_steps\", 1000,\n                     \"How often to save the model checkpoint.\")\n\nflags.DEFINE_integer(\"iterations_per_loop\", 1000,\n                     \"How many steps to make in each estimator call.\")\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    \"verbosity\", 1, \"How verbose our error messages should be\")\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_float(\n    \"include_unknowns\", -1.0,\n    \"If positive, probability of including answers of type `UNKNOWN`.\")\n\nflags.DEFINE_bool(\"use_tpu\", False, \"Whether to use TPU or GPU/CPU.\")\n\ntf.compat.v1.flags.DEFINE_string(\n    \"tpu_name\", None,\n    \"The Cloud TPU to use for training. This should be either the name \"\n    \"used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 \"\n    \"url.\")\n\ntf.compat.v1.flags.DEFINE_string(\n    \"tpu_zone\", None,\n    \"[Optional] GCE zone where the Cloud TPU is located in. If not \"\n    \"specified, we will attempt to automatically detect the GCE project from \"\n    \"metadata.\")\n\ntf.compat.v1.flags.DEFINE_string(\n    \"gcp_project\", None,\n    \"[Optional] Project name for the Cloud TPU-enabled project. If not \"\n    \"specified, we will attempt to automatically detect the GCE project from \"\n    \"metadata.\")\n\ntf.compat.v1.flags.DEFINE_string(\"master\", None, \"[Optional] TensorFlow master URL.\")\n\nflags.DEFINE_integer(\n    \"num_tpu_cores\", 8,\n    \"Only used if `use_tpu` is True. Total number of TPU cores to use.\")\n\nflags.DEFINE_bool(\n    \"verbose_logging\", False,\n    \"If true, all of the warnings related to data processing will be printed. \"\n    \"A number of warnings are expected for a normal NQ evaluation.\")\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(\"task_id\", 0,\n                     \"Train and dev shard to read from and write to.\")\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\n\n## Special flags - do not change\n\nflags.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\")\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\nFLAGS = flags.FLAGS","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### Here, we:\n1. Set up Bert\n2. Read in the test set\n3. Run it past the pre-built Bert model to create embeddings\n4. Use those embeddings to make predictions\n5. Write those predictions to `predictions.json`\n\nFeel free to change the code below. Code for the `tf2baseline.*` functions is included in the `tf2_0_baseline_w_bert` utility script, and can be customized, whether by forking the utility script and updating it, or by creating your own non-`tf2baseline` versions in this kernel.\n\nNote: the `tf2_0_baseline_w_bert` utility script contains code for training your own embeddings. Here that code is removed."},{"metadata":{"trusted":true},"cell_type":"code","source":"bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)\n\ntf2baseline.validate_flags_or_throw(bert_config)\ntf.io.gfile.makedirs(FLAGS.output_dir)\n\ntokenizer = tokenization.FullTokenizer(\n    vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)\n\ntpu_cluster_resolver = None\nif FLAGS.use_tpu and FLAGS.tpu_name:\n  tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(\n      FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)\n\nis_per_host = tf.compat.v1.estimator.tpu.InputPipelineConfig.PER_HOST_V2\nrun_config = tf.compat.v1.estimator.tpu.RunConfig(\n    cluster=tpu_cluster_resolver,\n    master=FLAGS.master,\n    model_dir=FLAGS.output_dir,\n    save_checkpoints_steps=FLAGS.save_checkpoints_steps,\n    tpu_config=tf.compat.v1.estimator.tpu.TPUConfig(\n        iterations_per_loop=FLAGS.iterations_per_loop,\n        num_shards=FLAGS.num_tpu_cores,\n        per_host_input_for_training=is_per_host))\n\nnum_train_steps = None\nnum_warmup_steps = None\n\nmodel_fn = tf2baseline.model_fn_builder(\n    bert_config=bert_config,\n    init_checkpoint=FLAGS.init_checkpoint,\n    learning_rate=FLAGS.learning_rate,\n    num_train_steps=num_train_steps,\n    num_warmup_steps=num_warmup_steps,\n    use_tpu=FLAGS.use_tpu,\n    use_one_hot_embeddings=FLAGS.use_tpu)\n\n# If TPU is not available, this falls back to normal Estimator on CPU or GPU.\nestimator = tf.compat.v1.estimator.tpu.TPUEstimator(\n    use_tpu=FLAGS.use_tpu,\n    model_fn=model_fn,\n    config=run_config,\n    train_batch_size=FLAGS.train_batch_size,\n    predict_batch_size=FLAGS.predict_batch_size)\n\nif FLAGS.do_predict:\n  if not FLAGS.output_prediction_file:\n    raise ValueError(\n        \"--output_prediction_file must be defined in predict mode.\")\n    \n  eval_examples = tf2baseline.read_nq_examples(\n      input_file=FLAGS.predict_file, is_training=False)\n\n  print(\"FLAGS.predict_file\", FLAGS.predict_file)\n\n  eval_writer = tf2baseline.FeatureWriter(\n      filename=os.path.join(FLAGS.output_dir, \"eval.tf_record\"),\n      is_training=False)\n  eval_features = []\n\n  def append_feature(feature):\n    eval_features.append(feature)\n    eval_writer.process_feature(feature)\n\n  num_spans_to_ids = tf2baseline.convert_examples_to_features(\n      examples=eval_examples,\n      tokenizer=tokenizer,\n      is_training=False,\n      output_fn=append_feature)\n  eval_writer.close()\n  eval_filename = eval_writer.filename\n\n  print(\"***** Running predictions *****\")\n  print(f\"  Num orig examples = %d\" % len(eval_examples))\n  print(f\"  Num split examples = %d\" % len(eval_features))\n  print(f\"  Batch size = %d\" % FLAGS.predict_batch_size)\n  for spans, ids in num_spans_to_ids.items():\n    print(f\"  Num split into %d = %d\" % (spans, len(ids)))\n\n  predict_input_fn = tf2baseline.input_fn_builder(\n      input_file=eval_filename,\n      seq_length=FLAGS.max_seq_length,\n      is_training=False,\n      drop_remainder=False)\n\n  print(eval_filename)\n\n  # If running eval on the TPU, you will need to specify the number of steps.\n  all_results = []\n\n  for result in estimator.predict(\n      predict_input_fn, yield_single_examples=True):\n    if len(all_results) % 1000 == 0:\n      print(\"Processing example: %d\" % (len(all_results)))\n\n    unique_id = int(result[\"unique_ids\"])\n    start_logits = [float(x) for x in result[\"start_logits\"].flat]\n    end_logits = [float(x) for x in result[\"end_logits\"].flat]\n    answer_type_logits = [float(x) for x in result[\"answer_type_logits\"].flat]\n\n    all_results.append(\n        tf2baseline.RawResult(\n            unique_id=unique_id,\n            start_logits=start_logits,\n            end_logits=end_logits,\n            answer_type_logits=answer_type_logits))\n\n  print (\"Going to candidates file\")\n\n  candidates_dict = tf2baseline.read_candidates(FLAGS.predict_file)\n\n  print (\"setting up eval features\")\n\n  eval_features = [\n      tf.train.Example.FromString(r)\n      for r in tf.compat.v1.python_io.tf_record_iterator(eval_filename)\n  ]\n\n  print (\"compute_pred_dict\")\n\n  nq_pred_dict = tf2baseline.compute_pred_dict(candidates_dict, eval_features,\n                                   [r._asdict() for r in all_results])\n  predictions_json = {\"predictions\": list(nq_pred_dict.values())}\n\n  print (\"writing json\")\n\n  with tf.io.gfile.GFile(FLAGS.output_prediction_file, \"w\") as f:\n    json.dump(predictions_json, f, indent=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### Now, we turn `predictions.json` into a `submission.csv` file."},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df = pd.read_json(\"/kaggle/working/predictions.json\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The 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_short_answer(entry):\n    # if entry[\"short_answers_score\"] < 1.5:\n    #     return \"\"\n    \n    answer = []    \n    for short_answer in entry[\"short_answers\"]:\n        if short_answer[\"start_token\"] > -1:\n            answer.append(str(short_answer[\"start_token\"]) + \":\" + str(short_answer[\"end_token\"]))\n    if entry[\"yes_no_answer\"] != \"NONE\":\n        answer.append(entry[\"yes_no_answer\"])\n    return \" \".join(answer)\n\ndef create_long_answer(entry):\n   # if entry[\"long_answer_score\"] < 1.5:\n   # return \"\"\n\n    answer = []\n    if entry[\"long_answer\"][\"start_token\"] > -1:\n        answer.append(str(entry[\"long_answer\"][\"start_token\"]) + \":\" + str(entry[\"long_answer\"][\"end_token\"]))\n    return \" \".join(answer)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df[\"long_answer_score\"] = test_answers_df[\"predictions\"].apply(lambda q: q[\"long_answer_score\"])\ntest_answers_df[\"short_answer_score\"] = test_answers_df[\"predictions\"].apply(lambda q: q[\"short_answers_score\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df[\"long_answer_score\"].describe()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"An example of what each sample's answers look like in `prediction.json`:"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df.predictions.values[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We re-format the JSON answers to match the requirements for submission."},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df[\"long_answer\"] = test_answers_df[\"predictions\"].apply(create_long_answer)\ntest_answers_df[\"short_answer\"] = test_answers_df[\"predictions\"].apply(create_short_answer)\ntest_answers_df[\"example_id\"] = test_answers_df[\"predictions\"].apply(lambda q: str(q[\"example_id\"]))\n\nlong_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"long_answer\"]))\nshort_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"short_answer\"]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"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."},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission = pd.read_csv(\"/kaggle/input/tensorflow2-question-answering/sample_submission.csv\")\n\nlong_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_long\")].apply(lambda q: long_answers[q[\"example_id\"].replace(\"_long\", \"\")], axis=1)\nshort_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_short\")].apply(lambda q: short_answers[q[\"example_id\"].replace(\"_short\", \"\")], axis=1)\n\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_long\"), \"PredictionString\"] = long_prediction_strings\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_short\"), \"PredictionString\"] = short_prediction_strings","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And finally, we write out our submission!"},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission.to_csv(\"submission.csv\", index=False)","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}