{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!python --version","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport time\nfrom tqdm import tqdm\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":{"trusted":true},"cell_type":"code","source":"%%time\n!pip install /kaggle/input/tf-115-dependencies/tensorboard-1.15.0-py3-none-any.whl\n!pip install /kaggle/input/tf-115-dependencies/tensorflow_estimator-1.15.1-py2.py3-none-any.whl\n!pip install /kaggle/input/tf-115/tensorflow_gpu-1.15.0-cp36-cp36m-manylinux2010_x86_64.whl \n!pip install /kaggle/input/tf-115-dependencies/bert_tensorflow-1.0.1-py2.py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\ntf.__version__","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tf.test.gpu_device_name()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**pred definitions**"},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\nimport json\nimport math\nimport os\nimport random\nfrom bert import modeling\nfrom bert import optimization\nfrom bert import tokenization\nimport six\nimport tensorflow as tf\n\nfrom tensorflow.contrib import cluster_resolver as contrib_cluster_resolver\nfrom tensorflow.contrib import tpu as contrib_tpu\n\nflags = tf.flags\nFLAGS = flags.FLAGS\n\n## Required parameters\nflags.DEFINE_string(\n    \"bert_config_file\", None,\n    \"The config json file corresponding to the pre-trained BERT model. \"\n    \"This specifies the model architecture.\")\n\nflags.DEFINE_string(\"vocab_file\", None,\n                    \"The vocabulary file that the BERT model was trained on.\")\n\nflags.DEFINE_string(\n    \"output_dir\", None,\n    \"The output directory where the model checkpoints will be written.\")\n\n## Other parameters\nflags.DEFINE_string(\"train_file\", None,\n                    \"SQuAD json for training. E.g., train-v1.1.json\")\n\nflags.DEFINE_string(\n    \"predict_file\", None,\n    \"SQuAD json for predictions. E.g., dev-v1.1.json or test-v1.1.json\")\n\nflags.DEFINE_string(\n    \"init_checkpoint\", None,\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\", False, \"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    \"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_bool(\"use_tpu\", False, \"Whether to use TPU or GPU/CPU.\")\n\ntf.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.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.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.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 SQuAD evaluation.\")\n\nflags.DEFINE_bool(\n    \"version_2_with_negative\", False,\n    \"If true, the SQuAD examples contain some that do not have an answer.\")\n\nflags.DEFINE_float(\n    \"null_score_diff_threshold\", 0.0,\n    \"If null_score - best_non_null is greater than the threshold predict null.\")\n\nflags.DEFINE_string('f', '', 'kernel')\n\nflags.DEFINE_string('HistoryManager.hist_file', '', 'kernel')\n\n\nflags.DEFINE_string(\n    \"output_prediction_file\", 'bert_model_output/predictions_tiny_dev.json',\n    \"Where to print predictions in NQ prediction format, to be passed to\"\n    \"natural_questions.nq_eval.\")\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\nTextSpan = collections.namedtuple(\"TextSpan\", \"token_positions text\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimport gzip,re,enum\nimport numpy as np\n\nSpan = collections.namedtuple(\"Span\", [\"start_token_idx\", \"end_token_idx\"])\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\n\ndef get_best_indexes(logits, n_best_size):\n  \"\"\"Get the n-best logits from a list.\"\"\"\n  index_and_score = sorted(\n      enumerate(logits[1:], 1), key=lambda x: x[1], reverse=True)\n  best_indexes = []\n  for i in range(len(index_and_score)):\n    if i >= n_best_size:\n      break\n    best_indexes.append(index_and_score[i][0])\n  return best_indexes\n\n# my version, returning nbest_summary\n\ndef my_compute_predictions(example):\n\n  # example = examples[0]\n  \"\"\"Converts an example into an NQEval object for evaluation.\"\"\"\n  predictions = []\n  n_best_size = 10\n  max_answer_length = 30\n\n  for unique_id, result in example.results.items():\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    start_indexes = get_best_indexes(result[\"start_logits\"], n_best_size)\n    end_indexes = get_best_indexes(result[\"end_logits\"], n_best_size)\n    for start_index in start_indexes:\n      for end_index in end_indexes:\n        if end_index < start_index:\n          continue\n        if token_map[start_index] == -1:\n          continue\n        if token_map[end_index] == -1:\n          continue\n        length = end_index - start_index + 1\n        if length > max_answer_length:\n          continue\n        summary = ScoreSummary()\n        summary.short_span_score = (\n            result[\"start_logits\"][start_index] +\n            result[\"end_logits\"][end_index])\n        summary.cls_token_score = (\n            result[\"start_logits\"][0] + result[\"end_logits\"][0])\n        summary.answer_type_logits = result[\"answer_type_logits\"]\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  # Default empty prediction.\n  score = -10000.0\n  short_span = Span(-1, -1)\n  long_span = Span(-1, -1)\n  summary = ScoreSummary()\n\n  nbest_summary = []\n  seen_labels = set()\n  nbest = 20\n  if predictions:\n    for i in range(min(nbest,len(predictions))):\n      score, summary, start_span, end_span = sorted(predictions, key=lambda x:x[0],reverse=True)[i]\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      summary.predicted_label = {\n          \"example_id\": example.example_id,\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      if (long_span.start_token_idx,long_span.end_token_idx,\\\n          short_span.start_token_idx,short_span.end_token_idx) not in seen_labels:\n        nbest_summary.append(summary)\n        seen_labels.add((long_span.start_token_idx,long_span.end_token_idx,\\\n                        short_span.start_token_idx,short_span.end_token_idx))\n\n  if not predictions:\n    summary.predicted_label = {\n          \"example_id\": example.example_id,\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    nbest_summary.append(summary)\n\n  return nbest_summary\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\n# my version (adapted to python3)\ndef compute_pred_dict(candidates_dict, dev_features, raw_results):\n  \"\"\"Computes official answer key from raw logits.\"\"\"\n  # candidates_dict, dev_features, raw_results = candidates_dict, eval_features, [r._asdict() for r in all_results]\n\n  raw_results_by_id = [(int(res[\"unique_id\"] + 1), res) for res in raw_results]\n\n  # Cast example id to int32 for each example, similarly to the raw results.\n  sess = tf.Session()\n  all_candidates = candidates_dict.items()\n  example_ids = tf.to_int32(np.array([int(k) for k, _ in all_candidates\n                                      ])).eval(session=sess)\n  examples_by_id = list(zip(example_ids, all_candidates))\n\n  # Cast unique_id also to int32 for features.\n  feature_ids = []\n  features = []\n  for f in dev_features:\n    feature_ids.append(f.features.feature[\"unique_ids\"].int64_list.value[0] + 1)\n    features.append(f.features.feature)\n  feature_ids = tf.to_int32(np.array(feature_ids)).eval(session=sess)\n  features_by_id = list(zip(feature_ids, features))\n\n  # Join examplew with features and raw results.\n  examples = []\n  merged = sorted(examples_by_id + raw_results_by_id + features_by_id,key=lambda x:x[0])\n\n  for idx, datum in merged:\n    if isinstance(datum, tuple):\n      examples.append(EvalExample(datum[0], datum[1]))\n    elif \"token_map\" in datum:\n      examples[-1].features[idx] = datum\n    else:\n      examples[-1].results[idx] = datum\n\n  # Construct prediction objects.\n  tf.logging.info(\"Computing predictions...\")\n  nbest_summary_dict = {}\n  summary_dict = {}\n  nq_pred_dict = {}\n  for e in examples:\n    nbest_summary = my_compute_predictions(e)\n    nbest_summary_dict[e.example_id] = nbest_summary\n    summary_dict[e.example_id] = nbest_summary[0]\n    nq_pred_dict[e.example_id] = nbest_summary[0].predicted_label\n    if len(nq_pred_dict) % 100 == 0:\n      tf.logging.info(\"Examples processed: %d\", len(nq_pred_dict))\n  tf.logging.info(\"Done computing predictions.\")\n\n  return nq_pred_dict, nbest_summary_dict\n\ndef read_candidates_from_one_split(input_path):\n  \n  candidates_dict = {}  \n  try:    \n    with gzip.open(input_path) as input_file:\n      tf.logging.info(\"Reading examples from: %s\", input_path)\n      for line in input_file:\n        e = json.loads(line)\n        candidates_dict[e[\"example_id\"]] = e[\"long_answer_candidates\"]\n  except:\n    with tf.io.gfile.GFile(input_path, \"r\") as input_file:\n      tf.logging.info(\"Reading examples from: %s\", input_path)\n      for line in input_file:\n        e = json.loads(line)\n        candidates_dict[e[\"example_id\"]] = e[\"long_answer_candidates\"]   \n#         if 'train' in input_path and len(candidates_dict)==3000:break\n  return candidates_dict  \n\n\ndef read_candidates(input_pattern):\n  \"\"\"Read candidates with real multiple processes.\"\"\"\n  input_paths = tf.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\nRawResult = collections.namedtuple(\n    \"RawResult\",\n    [\"unique_id\", \"start_logits\", \"end_logits\", \"answer_type_logits\"])\n\n\n\ndef create_model(bert_config, is_training, input_ids, input_mask, segment_ids,\n                 use_one_hot_embeddings):\n  \"\"\"Creates a classification model.\"\"\"\n  model = modeling.BertModel(\n      config=bert_config,\n      is_training=is_training,\n      input_ids=input_ids,\n      input_mask=input_mask,\n      token_type_ids=segment_ids,\n      use_one_hot_embeddings=use_one_hot_embeddings)\n\n  # Get the logits for the start and end predictions.\n  final_hidden = model.get_sequence_output()\n\n  final_hidden_shape = modeling.get_shape_list(final_hidden, expected_rank=3)\n  batch_size = final_hidden_shape[0]\n  seq_length = final_hidden_shape[1]\n  hidden_size = final_hidden_shape[2]\n\n  output_weights = tf.get_variable(\n      \"cls/nq/output_weights\", [2, hidden_size],\n      initializer=tf.truncated_normal_initializer(stddev=0.02))\n\n  output_bias = tf.get_variable(\n      \"cls/nq/output_bias\", [2], initializer=tf.zeros_initializer())\n\n  final_hidden_matrix = tf.reshape(final_hidden,\n                                   [batch_size * seq_length, hidden_size])\n  logits = tf.matmul(final_hidden_matrix, output_weights, transpose_b=True)\n  logits = tf.nn.bias_add(logits, output_bias)\n\n  logits = tf.reshape(logits, [batch_size, seq_length, 2])\n  logits = tf.transpose(logits, [2, 0, 1])\n\n  unstacked_logits = tf.unstack(logits, axis=0)\n\n  (start_logits, end_logits) = (unstacked_logits[0], unstacked_logits[1])\n\n  # Get the logits for the answer type prediction.\n  answer_type_output_layer = model.get_pooled_output()\n  answer_type_hidden_size = answer_type_output_layer.shape[-1].value\n\n  num_answer_types = 5  # YES, NO, UNKNOWN, SHORT, LONG\n  answer_type_output_weights = tf.get_variable(\n      \"answer_type_output_weights\", [num_answer_types, answer_type_hidden_size],\n      initializer=tf.truncated_normal_initializer(stddev=0.02))\n\n  answer_type_output_bias = tf.get_variable(\n      \"answer_type_output_bias\", [num_answer_types],\n      initializer=tf.zeros_initializer())\n\n  answer_type_logits = tf.matmul(\n      answer_type_output_layer, answer_type_output_weights, transpose_b=True)\n  answer_type_logits = tf.nn.bias_add(answer_type_logits,\n                                      answer_type_output_bias)\n\n  return (start_logits, end_logits, answer_type_logits)\n\n\ndef validate_flags_or_throw(bert_config):\n  \"\"\"Validate the input FLAGS or throw an exception.\"\"\"\n  if not FLAGS.do_train and not FLAGS.do_predict:\n    raise ValueError(\"At least one of `{do_train,do_predict}` must be True.\")\n\n  if FLAGS.do_train:\n    if not FLAGS.train_precomputed_file:\n      raise ValueError(\"If `do_train` is True, then `train_precomputed_file` \"\n                       \"must be specified.\")\n    if not FLAGS.train_num_precomputed:\n      raise ValueError(\"If `do_train` is True, then `train_num_precomputed` \"\n                       \"must be specified.\")\n\n  if FLAGS.do_predict:\n    if not FLAGS.predict_file:\n      raise ValueError(\n          \"If `do_predict` is True, then `predict_file` must be specified.\")\n\n  if FLAGS.max_seq_length > bert_config.max_position_embeddings:\n    raise ValueError(\n        \"Cannot use sequence length %d because the BERT model \"\n        \"was only trained up to sequence length %d\" %\n        (FLAGS.max_seq_length, bert_config.max_position_embeddings))\n\n  if FLAGS.max_seq_length <= FLAGS.max_query_length + 3:\n    raise ValueError(\n        \"The max_seq_length (%d) must be greater than max_query_length \"\n        \"(%d) + 3\" % (FLAGS.max_seq_length, FLAGS.max_query_length))\n\nclass AnswerType(enum.IntEnum):\n  \"\"\"Type of NQ answer.\"\"\"\n  UNKNOWN = 0\n  YES = 1\n  NO = 2\n  SHORT = 3\n  LONG = 4\n\n\nclass Answer(collections.namedtuple(\"Answer\", [\"type\", \"text\", \"offset\"])):\n  \"\"\"Answer record.\n  An Answer contains the type of the answer and possibly the text (for\n  long) as well as the offset (for extractive).\n  \"\"\"\n\n  def __new__(cls, type_, text=None, offset=None):\n    return super(Answer, cls).__new__(cls, type_, text, offset)\n\n\nclass InputFeatures(object):\n  \"\"\"A single set of features of data.\"\"\"\n\n  def __init__(self,\n               unique_id,\n               example_index,\n               doc_span_index,\n               token_to_orig_map,\n               input_ids,\n               input_mask,\n               segment_ids,\n               start_position=None,\n               end_position=None,\n               answer_text=\"\",\n               answer_type=AnswerType.SHORT):\n    self.unique_id = unique_id\n    self.example_index = example_index\n    self.doc_span_index = doc_span_index\n    self.token_to_orig_map = token_to_orig_map\n    self.input_ids = input_ids\n    self.input_mask = input_mask\n    self.segment_ids = segment_ids\n    self.start_position = start_position\n    self.end_position = end_position\n    self.answer_text = answer_text\n    self.answer_type = answer_type\n\n\n_SPECIAL_TOKENS_RE = re.compile(r\"^\\[[^ ]*\\]$\", re.UNICODE)\n\ndef tokenize(tokenizer, text, apply_basic_tokenization=False):\n  \"\"\"Tokenizes text, optionally looking up special tokens separately.\n  Args:\n    tokenizer: a tokenizer from bert.tokenization.FullTokenizer\n    text: text to tokenize\n    apply_basic_tokenization: If True, apply the basic tokenization. If False,\n      apply the full tokenization (basic + wordpiece).\n  Returns:\n    tokenized text.\n  A special token is any text with no spaces enclosed in square brackets with no\n  space, so we separate those out and look them up in the dictionary before\n  doing actual tokenization.\n  \"\"\"\n  tokenize_fn = tokenizer.tokenize\n  if apply_basic_tokenization:\n    tokenize_fn = tokenizer.basic_tokenizer.tokenize\n  tokens = []\n  for token in text.split(\" \"):\n    if _SPECIAL_TOKENS_RE.match(token):\n      if token in tokenizer.vocab:\n        tokens.append(token)\n      else:\n        tokens.append(tokenizer.wordpiece_tokenizer.unk_token)\n    else:\n      tokens.extend(tokenize_fn(token))\n  return tokens\n\ndef convert_single_example(example, tokenizer, is_training):\n  \"\"\"Converts a single NqExample into a list of InputFeatures.\"\"\"\n  tok_to_orig_index = []\n  orig_to_tok_index = []\n  all_doc_tokens = []\n  features = []\n  for (i, token) in enumerate(example.doc_tokens):\n    orig_to_tok_index.append(len(all_doc_tokens))\n    sub_tokens = tokenize(tokenizer, token)\n    tok_to_orig_index.extend([i] * len(sub_tokens))\n    all_doc_tokens.extend(sub_tokens)\n\n  # `tok_to_orig_index` maps wordpiece indices to indices of whitespace\n  # tokenized word tokens in the contexts. The word tokens might themselves\n  # correspond to word tokens in a larger document, with the mapping given\n  # by `doc_tokens_map`.\n  if example.doc_tokens_map:\n    tok_to_orig_index = [\n        example.doc_tokens_map[index] for index in tok_to_orig_index\n    ]\n\n  # QUERY\n  query_tokens = []\n  query_tokens.append(\"[Q]\")\n  query_tokens.extend(tokenize(tokenizer, example.questions[-1]))\n  if len(query_tokens) > FLAGS.max_query_length:\n    query_tokens = query_tokens[-FLAGS.max_query_length:]\n\n  # ANSWER\n  tok_start_position = 0\n  tok_end_position = 0\n  if is_training:\n    tok_start_position = orig_to_tok_index[example.start_position]\n    if example.end_position < len(example.doc_tokens) - 1:\n      tok_end_position = orig_to_tok_index[example.end_position + 1] - 1\n    else:\n      tok_end_position = len(all_doc_tokens) - 1\n\n  # The -3 accounts for [CLS], [SEP] and [SEP]\n  max_tokens_for_doc = FLAGS.max_seq_length - len(query_tokens) - 3\n\n  # We can have documents that are longer than the maximum sequence length.\n  # To deal with this we do a sliding window approach, where we take chunks\n  # of up to our max length with a stride of `doc_stride`.\n  _DocSpan = collections.namedtuple(  # pylint: disable=invalid-name\n      \"DocSpan\", [\"start\", \"length\"])\n  doc_spans = []\n  start_offset = 0\n  while start_offset < len(all_doc_tokens):\n    length = len(all_doc_tokens) - start_offset\n    length = min(length, max_tokens_for_doc)\n    doc_spans.append(_DocSpan(start=start_offset, length=length))\n    if start_offset + length == len(all_doc_tokens):\n      break\n    start_offset += min(length, FLAGS.doc_stride)\n\n  for (doc_span_index, doc_span) in enumerate(doc_spans):\n    tokens = []\n    token_to_orig_map = {}\n    segment_ids = []\n    tokens.append(\"[CLS]\")\n    segment_ids.append(0)\n    tokens.extend(query_tokens)\n    segment_ids.extend([0] * len(query_tokens))\n    tokens.append(\"[SEP]\")\n    segment_ids.append(0)\n\n    for i in range(doc_span.length):\n      split_token_index = doc_span.start + i\n      token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index]\n      tokens.append(all_doc_tokens[split_token_index])\n      segment_ids.append(1)\n    tokens.append(\"[SEP]\")\n    segment_ids.append(1)\n    assert len(tokens) == len(segment_ids)\n\n    input_ids = tokenizer.convert_tokens_to_ids(tokens)\n\n    # The mask has 1 for real tokens and 0 for padding tokens. Only real\n    # tokens are attended to.\n    input_mask = [1] * len(input_ids)\n\n    # Zero-pad up to the sequence length.\n    padding = [0] * (FLAGS.max_seq_length - len(input_ids))\n    input_ids.extend(padding)\n    input_mask.extend(padding)\n    segment_ids.extend(padding)\n\n    assert len(input_ids) == FLAGS.max_seq_length\n    assert len(input_mask) == FLAGS.max_seq_length\n    assert len(segment_ids) == FLAGS.max_seq_length\n\n    start_position = None\n    end_position = None\n    answer_type = None\n    answer_text = \"\"\n    if is_training:\n      doc_start = doc_span.start\n      doc_end = doc_span.start + doc_span.length - 1\n      # For training, if our document chunk does not contain an annotation\n      # we throw it out, since there is nothing to predict.\n      contains_an_annotation = (\n          tok_start_position >= doc_start and tok_end_position <= doc_end)\n      if ((not contains_an_annotation) or\n          example.answer.type == AnswerType.UNKNOWN):\n        # If an example has unknown answer type or does not contain the answer\n        # span, then we only include it with probability --include_unknowns.\n        # When we include an example with unknown answer type, we set the first\n        # token of the passage to be the annotated short span.\n        if (FLAGS.include_unknowns < 0 or\n            random.random() > FLAGS.include_unknowns):\n          continue\n        start_position = 0\n        end_position = 0\n        answer_type = AnswerType.UNKNOWN\n      else:\n        doc_offset = len(query_tokens) + 2\n        start_position = tok_start_position - doc_start + doc_offset\n        end_position = tok_end_position - doc_start + doc_offset\n        answer_type = example.answer.type\n\n      answer_text = \" \".join(tokens[start_position:(end_position + 1)])\n\n    feature = InputFeatures(\n        unique_id=-1,\n        example_index=-1,\n        doc_span_index=doc_span_index,\n        token_to_orig_map=token_to_orig_map,\n        input_ids=input_ids,\n        input_mask=input_mask,\n        segment_ids=segment_ids,\n        start_position=start_position,\n        end_position=end_position,\n        answer_text=answer_text,\n        answer_type=answer_type)\n\n    features.append(feature)\n\n  return features\n\n\ndef convert_examples_to_features(examples, tokenizer, is_training, output_fn):\n  \"\"\"Converts a list of NqExamples into InputFeatures.\"\"\"\n  num_spans_to_ids = collections.defaultdict(list)\n\n  for example in tqdm(examples):\n    example_index = example.example_id\n    features = convert_single_example(example, tokenizer, is_training)\n    num_spans_to_ids[len(features)].append(example.qas_id)\n\n    for feature in features:\n      feature.example_index = example_index\n      feature.unique_id = feature.example_index + feature.doc_span_index\n      output_fn(feature)\n\n  return num_spans_to_ids\n\ndef read_nq_entry(entry, is_training):\n  \"\"\"Converts a NQ entry into a list of NqExamples.\"\"\"\n\n  def is_whitespace(c):\n    return c in \" \\t\\r\\n\" or ord(c) == 0x202F\n\n  examples = []\n  contexts_id = entry[\"id\"]\n  contexts = entry[\"contexts\"]\n  doc_tokens = []\n  char_to_word_offset = []\n  prev_is_whitespace = True\n  for c in contexts:\n    if is_whitespace(c):\n      prev_is_whitespace = True\n    else:\n      if prev_is_whitespace:\n        doc_tokens.append(c)\n      else:\n        doc_tokens[-1] += c\n      prev_is_whitespace = False\n    char_to_word_offset.append(len(doc_tokens) - 1)\n\n  questions = []\n  for i, question in enumerate(entry[\"questions\"]):\n    qas_id = \"{}\".format(contexts_id)\n    question_text = question[\"input_text\"]\n    start_position = None\n    end_position = None\n    answer = None\n    if is_training:\n      answer_dict = entry[\"answers\"][i]\n      answer = make_nq_answer(contexts, answer_dict)\n\n      # For now, only handle extractive, yes, and no.\n      if answer is None or answer.offset is None:\n        continue\n      start_position = char_to_word_offset[answer.offset]\n      end_position = char_to_word_offset[answer.offset + len(answer.text) - 1]\n\n      # Only add answers where the text can be exactly recovered from the\n      # document. If this CAN'T happen it's likely due to weird Unicode\n      # stuff so we will just skip the example.\n      #\n      # Note that this means for training mode, every example is NOT\n      # guaranteed to be preserved.\n      actual_text = \" \".join(doc_tokens[start_position:(end_position + 1)])\n      cleaned_answer_text = \" \".join(\n          tokenization.whitespace_tokenize(answer.text))\n      if actual_text.find(cleaned_answer_text) == -1:\n        tf.logging.warning(\"Could not find answer: '%s' vs. '%s'\", actual_text,\n                           cleaned_answer_text)\n        continue\n\n    questions.append(question_text)\n    example = NqExample(\n        example_id=int(contexts_id),\n        qas_id=qas_id,\n        questions=questions[:],\n        doc_tokens=doc_tokens,\n        doc_tokens_map=entry.get(\"contexts_map\", None),\n        answer=answer,\n        start_position=start_position,\n        end_position=end_position)\n    examples.append(example)\n  return examples\n\n\n\nclass NqExample(object):\n  \"\"\"A single training/test example.\"\"\"\n\n  def __init__(self,\n               example_id,\n               qas_id,\n               questions,\n               doc_tokens,\n               doc_tokens_map=None,\n               answer=None,\n               start_position=None,\n               end_position=None):\n    self.example_id = example_id\n    self.qas_id = qas_id\n    self.questions = questions\n    self.doc_tokens = doc_tokens\n    self.doc_tokens_map = doc_tokens_map\n    self.answer = answer\n    self.start_position = start_position\n    self.end_position = end_position\n\n\ndef has_long_answer(a):\n  return (a[\"long_answer\"][\"start_token\"] >= 0 and\n          a[\"long_answer\"][\"end_token\"] >= 0)\n\n\ndef should_skip_context(e, idx):\n  if (FLAGS.skip_nested_contexts and\n      not e[\"long_answer_candidates\"][idx][\"top_level\"]):\n    return True\n  elif not get_candidate_text(e, idx).text.strip():\n    # Skip empty contexts.\n    return True\n  else:\n    return False\n\n\ndef get_first_annotation(e):\n  \"\"\"Returns the first short or long answer in the example.\n  Args:\n    e: (dict) annotated example.\n  Returns:\n    annotation: (dict) selected annotation\n    annotated_idx: (int) index of the first annotated candidate.\n    annotated_sa: (tuple) char offset of the start and end token\n        of the short answer. The end token is exclusive.\n  \"\"\"\n  positive_annotations = sorted(\n      [a for a in e[\"annotations\"] if has_long_answer(a)],\n      key=lambda a: a[\"long_answer\"][\"candidate_index\"])\n\n  for a in positive_annotations:\n    if a[\"short_answers\"]:\n      idx = a[\"long_answer\"][\"candidate_index\"]\n      start_token = a[\"short_answers\"][0][\"start_token\"]\n      end_token = a[\"short_answers\"][-1][\"end_token\"]\n      return a, idx, (token_to_char_offset(e, idx, start_token),\n                      token_to_char_offset(e, idx, end_token) - 1)\n\n  for a in positive_annotations:\n    idx = a[\"long_answer\"][\"candidate_index\"]\n    return a, idx, (-1, -1)\n\n  return None, -1, (-1, -1)\n\n\ndef get_text_span(example, span):\n  \"\"\"Returns the text in the example's document in the given token span.\"\"\"\n  token_positions = []\n  tokens = []\n  for i in range(span[\"start_token\"], span[\"end_token\"]):\n    t = example[\"document_tokens\"][i]\n    if not t[\"html_token\"]:\n      token_positions.append(i)\n      token = t[\"token\"].replace(\" \", \"\")\n      tokens.append(token)\n  return TextSpan(token_positions, \" \".join(tokens))\n\n\ndef token_to_char_offset(e, candidate_idx, token_idx):\n  \"\"\"Converts a token index to the char offset within the candidate.\"\"\"\n  c = e[\"long_answer_candidates\"][candidate_idx]\n  char_offset = 0\n  for i in range(c[\"start_token\"], token_idx):\n    t = e[\"document_tokens\"][i]\n    if not t[\"html_token\"]:\n      token = t[\"token\"].replace(\" \", \"\")\n      char_offset += len(token) + 1\n  return char_offset\n\n\ndef get_candidate_type(e, idx):\n  \"\"\"Returns the candidate's type: Table, Paragraph, List or Other.\"\"\"\n  c = e[\"long_answer_candidates\"][idx]\n  first_token = e[\"document_tokens\"][c[\"start_token\"]][\"token\"]\n  if first_token == \"<Table>\":\n    return \"Table\"\n  elif first_token == \"<P>\":\n    return \"Paragraph\"\n  elif first_token in (\"<Ul>\", \"<Dl>\", \"<Ol>\"):\n    return \"List\"\n  elif first_token in (\"<Tr>\", \"<Li>\", \"<Dd>\", \"<Dt>\"):\n    return \"Other\"\n  else:\n    tf.logging.warning(\"Unknoww candidate type found: %s\", first_token)\n    return \"Other\"\n\n\ndef add_candidate_types_and_positions(e):\n  \"\"\"Adds type and position info to each candidate in the document.\"\"\"\n  counts = collections.defaultdict(int)\n  for idx, c in candidates_iter(e):\n    context_type = get_candidate_type(e, idx)\n    if counts[context_type] < FLAGS.max_position:\n      counts[context_type] += 1\n    c[\"type_and_position\"] = \"[%s=%d]\" % (context_type, counts[context_type])\n\n\ndef get_candidate_type_and_position(e, idx):\n  \"\"\"Returns type and position info for the candidate at the given index.\"\"\"\n  if idx == -1:\n    return \"[NoLongAnswer]\"\n  else:\n    return e[\"long_answer_candidates\"][idx][\"type_and_position\"]\n\n\ndef get_candidate_text(e, idx):\n  \"\"\"Returns a text representation of the candidate at the given index.\"\"\"\n  # No candidate at this index.\n  if idx < 0 or idx >= len(e[\"long_answer_candidates\"]):\n    return TextSpan([], \"\")\n\n  # This returns an actual candidate.\n  return get_text_span(e, e[\"long_answer_candidates\"][idx])\n\n\ndef candidates_iter(e):\n  \"\"\"Yield's the candidates that should not be skipped in an example.\"\"\"\n  for idx, c in enumerate(e[\"long_answer_candidates\"]):\n    if should_skip_context(e, idx):\n      continue\n    yield idx, c\n\ndef create_example_from_jsonl(line):\n  \"\"\"Creates an NQ example from a given line of JSON.\"\"\"\n  e = json.loads(line, object_pairs_hook=collections.OrderedDict)\n  add_candidate_types_and_positions(e)\n  annotation, annotated_idx, annotated_sa = get_first_annotation(e)\n\n  # annotated_idx: index of the first annotated context, -1 if null.\n  # annotated_sa: short answer start and end char offsets, (-1, -1) if null.\n  question = {\"input_text\": e[\"question_text\"]}\n  answer = {\n      \"candidate_id\": annotated_idx,\n      \"span_text\": \"\",\n      \"span_start\": -1,\n      \"span_end\": -1,\n      \"input_text\": \"long\",\n  }\n\n  # Yes/no answers are added in the input text.\n  if annotation is not None:\n    assert annotation[\"yes_no_answer\"] in (\"YES\", \"NO\", \"NONE\")\n    if annotation[\"yes_no_answer\"] in (\"YES\", \"NO\"):\n      answer[\"input_text\"] = annotation[\"yes_no_answer\"].lower()\n\n  # Add a short answer if one was found.\n  if annotated_sa != (-1, -1):\n    answer[\"input_text\"] = \"short\"\n    span_text = get_candidate_text(e, annotated_idx).text\n    answer[\"span_text\"] = span_text[annotated_sa[0]:annotated_sa[1]]\n    answer[\"span_start\"] = annotated_sa[0]\n    answer[\"span_end\"] = annotated_sa[1]\n    expected_answer_text = get_text_span(\n        e, {\n            \"start_token\": annotation[\"short_answers\"][0][\"start_token\"],\n            \"end_token\": annotation[\"short_answers\"][-1][\"end_token\"],\n        }).text\n    assert expected_answer_text == answer[\"span_text\"], (expected_answer_text,\n                                                         answer[\"span_text\"])\n\n  # Add a long answer if one was found.\n  elif annotation and annotation[\"long_answer\"][\"candidate_index\"] >= 0:\n    answer[\"span_text\"] = get_candidate_text(e, annotated_idx).text\n    answer[\"span_start\"] = 0\n    answer[\"span_end\"] = len(answer[\"span_text\"])\n\n  context_idxs = [-1]\n  context_list = [{\"id\": -1, \"type\": get_candidate_type_and_position(e, -1)}]\n  context_list[-1][\"text_map\"], context_list[-1][\"text\"] = (\n      get_candidate_text(e, -1))\n  for idx, _ in candidates_iter(e):\n    context = {\"id\": idx, \"type\": get_candidate_type_and_position(e, idx)}\n    context[\"text_map\"], context[\"text\"] = get_candidate_text(e, idx)\n    context_idxs.append(idx)\n    context_list.append(context)\n    if len(context_list) >= FLAGS.max_contexts:\n      break\n\n  # Assemble example.\n  example = {\n      \"name\": e[\"document_title\"],\n      \"id\": str(e[\"example_id\"]),\n      \"questions\": [question],\n      \"answers\": [answer],\n      \"has_correct_context\": annotated_idx in context_idxs\n  }\n\n  single_map = []\n  single_context = []\n  offset = 0\n  for context in context_list:\n    single_map.extend([-1, -1])\n    single_context.append(\"[ContextId=%d] %s\" %\n                          (context[\"id\"], context[\"type\"]))\n    offset += len(single_context[-1]) + 1\n    if context[\"id\"] == annotated_idx:\n      answer[\"span_start\"] += offset\n      answer[\"span_end\"] += offset\n\n    # Many contexts are empty once the HTML tags have been stripped, so we\n    # want to skip those.\n    if context[\"text\"]:\n      single_map.extend(context[\"text_map\"])\n      single_context.append(context[\"text\"])\n      offset += len(single_context[-1]) + 1\n\n  example[\"contexts\"] = \" \".join(single_context)\n  example[\"contexts_map\"] = single_map\n  if annotated_idx in context_idxs:\n    expected = example[\"contexts\"][answer[\"span_start\"]:answer[\"span_end\"]]\n\n    # This is a sanity check to ensure that the calculated start and end\n    # indices match the reported span text. If this assert fails, it is likely\n    # a bug in the data preparation code above.\n    assert expected == answer[\"span_text\"], (expected, answer[\"span_text\"])\n\n  return example\n\n\ndef create_example_from_jsonl_simple(line):\n  \"\"\"Creates an NQ example from a given line of JSON.\"\"\"\n  e = json.loads(line, object_pairs_hook=collections.OrderedDict)\n  document_tokens = e[\"document_text\"].split(\" \")\n  e[\"document_tokens\"] = []\n  for token in document_tokens:\n      e[\"document_tokens\"].append({\"token\":token, \"start_byte\":-1, \"end_byte\":-1, \"html_token\":\"<\" in token})\n\n  add_candidate_types_and_positions(e)\n  try:\n    annotation, annotated_idx, annotated_sa = get_first_annotation(e)\n  except:\n    annotation, annotated_idx, annotated_sa = None, -1, (-1, -1)  \n\n  # annotated_idx: index of the first annotated context, -1 if null.\n  # annotated_sa: short answer start and end char offsets, (-1, -1) if null.\n  question = {\"input_text\": e[\"question_text\"]}\n  answer = {\n      \"candidate_id\": annotated_idx,\n      \"span_text\": \"\",\n      \"span_start\": -1,\n      \"span_end\": -1,\n      \"input_text\": \"long\",\n  }\n\n  # Yes/no answers are added in the input text.\n  if annotation is not None:\n    assert annotation[\"yes_no_answer\"] in (\"YES\", \"NO\", \"NONE\")\n    if annotation[\"yes_no_answer\"] in (\"YES\", \"NO\"):\n      answer[\"input_text\"] = annotation[\"yes_no_answer\"].lower()\n\n  # Add a short answer if one was found.\n  if annotated_sa != (-1, -1):\n    answer[\"input_text\"] = \"short\"\n    span_text = get_candidate_text(e, annotated_idx).text\n    answer[\"span_text\"] = span_text[annotated_sa[0]:annotated_sa[1]]\n    answer[\"span_start\"] = annotated_sa[0]\n    answer[\"span_end\"] = annotated_sa[1]\n    expected_answer_text = get_text_span(\n        e, {\n            \"start_token\": annotation[\"short_answers\"][0][\"start_token\"],\n            \"end_token\": annotation[\"short_answers\"][-1][\"end_token\"],\n        }).text\n    assert expected_answer_text == answer[\"span_text\"], (expected_answer_text,\n                                                         answer[\"span_text\"])\n\n  # Add a long answer if one was found.\n  elif annotation and annotation[\"long_answer\"][\"candidate_index\"] >= 0:\n    answer[\"span_text\"] = get_candidate_text(e, annotated_idx).text\n    answer[\"span_start\"] = 0\n    answer[\"span_end\"] = len(answer[\"span_text\"])\n\n  context_idxs = [-1]\n  context_list = [{\"id\": -1, \"type\": get_candidate_type_and_position(e, -1)}]\n  context_list[-1][\"text_map\"], context_list[-1][\"text\"] = (\n      get_candidate_text(e, -1))\n  for idx, _ in candidates_iter(e):\n    context = {\"id\": idx, \"type\": get_candidate_type_and_position(e, idx)}\n    context[\"text_map\"], context[\"text\"] = get_candidate_text(e, idx)\n    context_idxs.append(idx)\n    context_list.append(context)\n    if len(context_list) >= FLAGS.max_contexts:\n      break\n\n  if \"document_title\" not in e:\n      e[\"document_title\"] = e[\"example_id\"]\n\n  # Assemble example.\n  example = {\n      \"name\": e[\"document_title\"],\n      \"id\": str(e[\"example_id\"]),\n      \"questions\": [question],\n      \"answers\": [answer],\n      \"has_correct_context\": annotated_idx in context_idxs\n  }\n\n  single_map = []\n  single_context = []\n  offset = 0\n  for context in context_list:\n    single_map.extend([-1, -1])\n    single_context.append(\"[ContextId=%d] %s\" %\n                          (context[\"id\"], context[\"type\"]))\n    offset += len(single_context[-1]) + 1\n    if context[\"id\"] == annotated_idx:\n      answer[\"span_start\"] += offset\n      answer[\"span_end\"] += offset\n\n    # Many contexts are empty once the HTML tags have been stripped, so we\n    # want to skip those.\n    if context[\"text\"]:\n      single_map.extend(context[\"text_map\"])\n      single_context.append(context[\"text\"])\n      offset += len(single_context[-1]) + 1\n\n  example[\"contexts\"] = \" \".join(single_context)\n  example[\"contexts_map\"] = single_map\n  if annotated_idx in context_idxs:\n    expected = example[\"contexts\"][answer[\"span_start\"]:answer[\"span_end\"]]\n\n    # This is a sanity check to ensure that the calculated start and end\n    # indices match the reported span text. If this assert fails, it is likely\n    # a bug in the data preparation code above.\n    assert expected == answer[\"span_text\"], (expected, answer[\"span_text\"])\n\n  return example\n\ndef read_nq_examples(input_file, is_training):\n  \"\"\"Read a NQ json file into a list of NqExample.\"\"\"\n  input_paths = tf.gfile.Glob(input_file)\n  input_data = []\n\n  def _open(path):\n    if path.endswith(\".gz\"):\n      return gzip.open(path, \"r\")\n    else:\n      return tf.gfile.Open(path, \"r\")\n\n  for path in input_paths:\n    tf.logging.info(\"Reading: %s\", path)\n    with _open(path) as input_file:\n      for line in input_file:\n        try:\n          input_data.append(create_example_from_jsonl(line))\n        except:\n          input_data.append(create_example_from_jsonl_simple(line))\n#         if 'train' in path and len(input_data)==3000: break\n\n  examples = []\n  for entry in input_data:\n    examples.extend(read_nq_entry(entry, is_training))\n  return examples\n\ndef model_fn_builder(bert_config, init_checkpoint, learning_rate,\n                     num_train_steps, num_warmup_steps, use_tpu,\n                     use_one_hot_embeddings):\n  \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n\n  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n    \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n\n    tf.logging.info(\"*** Features ***\")\n    for name in sorted(features.keys()):\n      tf.logging.info(\"  name = %s, shape = %s\" % (name, features[name].shape))\n\n    unique_ids = features[\"unique_ids\"]\n    input_ids = features[\"input_ids\"]\n    input_mask = features[\"input_mask\"]\n    segment_ids = features[\"segment_ids\"]\n\n    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n\n    (start_logits, end_logits, answer_type_logits) = create_model(\n        bert_config=bert_config,\n        is_training=is_training,\n        input_ids=input_ids,\n        input_mask=input_mask,\n        segment_ids=segment_ids,\n        use_one_hot_embeddings=use_one_hot_embeddings)\n\n    tvars = tf.trainable_variables()\n\n    initialized_variable_names = {}\n    scaffold_fn = None\n    if init_checkpoint:\n      (assignment_map, initialized_variable_names\n      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)\n      if use_tpu:\n\n        def tpu_scaffold():\n          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n          return tf.train.Scaffold()\n\n        scaffold_fn = tpu_scaffold\n      else:\n        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n\n    tf.logging.info(\"**** Trainable Variables ****\")\n    for var in tvars:\n      init_string = \"\"\n      if var.name in initialized_variable_names:\n        init_string = \", *INIT_FROM_CKPT*\"\n      tf.logging.info(\"  name = %s, shape = %s%s\", var.name, var.shape,\n                      init_string)\n\n    output_spec = None\n    if mode == tf.estimator.ModeKeys.TRAIN:\n      seq_length = modeling.get_shape_list(input_ids)[1]\n\n      # Computes the loss for positions.\n      def compute_loss(logits, positions):\n        one_hot_positions = tf.one_hot(\n            positions, depth=seq_length, dtype=tf.float32)\n        log_probs = tf.nn.log_softmax(logits, axis=-1)\n        loss = -tf.reduce_mean(\n            tf.reduce_sum(one_hot_positions * log_probs, axis=-1))\n        return loss\n\n      # Computes the loss for labels.\n      def compute_label_loss(logits, labels):\n        one_hot_labels = tf.one_hot(\n            labels, depth=len(AnswerType), dtype=tf.float32)\n        log_probs = tf.nn.log_softmax(logits, axis=-1)\n        loss = -tf.reduce_mean(\n            tf.reduce_sum(one_hot_labels * log_probs, axis=-1))\n        return loss\n\n      start_positions = features[\"start_positions\"]\n      end_positions = features[\"end_positions\"]\n      answer_types = features[\"answer_types\"]\n\n      start_loss = compute_loss(start_logits, start_positions)\n      end_loss = compute_loss(end_logits, end_positions)\n      answer_type_loss = compute_label_loss(answer_type_logits, answer_types)\n\n      total_loss = (start_loss + end_loss + answer_type_loss) / 3.0\n\n      train_op = optimization.create_optimizer(total_loss, learning_rate,\n                                               num_train_steps,\n                                               num_warmup_steps, use_tpu)\n\n      output_spec = contrib_tpu.TPUEstimatorSpec(\n          mode=mode,\n          loss=total_loss,\n          train_op=train_op,\n          scaffold_fn=scaffold_fn)\n    elif mode == tf.estimator.ModeKeys.PREDICT:\n      predictions = {\n          \"unique_ids\": unique_ids,\n          \"start_logits\": start_logits,\n          \"end_logits\": end_logits,\n          \"answer_type_logits\": answer_type_logits,\n      }\n      output_spec = contrib_tpu.TPUEstimatorSpec(\n          mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)\n    else:\n      raise ValueError(\"Only TRAIN and PREDICT modes are supported: %s\" %\n                       (mode))\n\n    return output_spec\n\n  return model_fn\n\ndef input_fn_builder(input_file, seq_length, is_training, drop_remainder):\n  \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n\n  name_to_features = {\n      \"unique_ids\": tf.FixedLenFeature([], tf.int64),\n      \"input_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n      \"input_mask\": tf.FixedLenFeature([seq_length], tf.int64),\n      \"segment_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n  }\n\n  if is_training:\n    name_to_features[\"start_positions\"] = tf.FixedLenFeature([], tf.int64)\n    name_to_features[\"end_positions\"] = tf.FixedLenFeature([], tf.int64)\n    name_to_features[\"answer_types\"] = tf.FixedLenFeature([], tf.int64)\n\n  def _decode_record(record, name_to_features):\n    \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n    example = tf.parse_single_example(record, name_to_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.to_int32(t)\n      example[name] = t\n\n    return example\n\n  def input_fn(params):\n    \"\"\"The actual input function.\"\"\"\n    batch_size = params[\"batch_size\"]\n\n    # For training, we want a lot of parallel reading and shuffling.\n    # For eval, we want no shuffling and parallel reading doesn't matter.\n    d = tf.data.TFRecordDataset(input_file)\n    if is_training:\n      d = d.repeat()\n      d = d.shuffle(buffer_size=100)\n\n    d = d.apply(\n        tf.data.experimental.map_and_batch(\n            lambda record: _decode_record(record, name_to_features),\n            batch_size=batch_size,\n            drop_remainder=drop_remainder))\n\n    return d\n\n  return input_fn\n\n\nclass FeatureWriter(object):\n  \"\"\"Writes InputFeature to TF example file.\"\"\"\n\n  def __init__(self, filename, is_training):\n    self.filename = filename\n    self.is_training = is_training\n    self.num_features = 0\n    self._writer = tf.python_io.TFRecordWriter(filename)\n\n  def process_feature(self, feature):\n    \"\"\"Write a InputFeature to the TFRecordWriter as a tf.train.Example.\"\"\"\n    self.num_features += 1\n\n    def create_int_feature(values):\n      feature = tf.train.Feature(\n          int64_list=tf.train.Int64List(value=list(values)))\n      return feature\n\n    features = collections.OrderedDict()\n    features[\"unique_ids\"] = create_int_feature([feature.unique_id])\n    features[\"input_ids\"] = create_int_feature(feature.input_ids)\n    features[\"input_mask\"] = create_int_feature(feature.input_mask)\n    features[\"segment_ids\"] = create_int_feature(feature.segment_ids)\n\n    if self.is_training:\n      features[\"start_positions\"] = create_int_feature([feature.start_position])\n      features[\"end_positions\"] = create_int_feature([feature.end_position])\n      features[\"answer_types\"] = create_int_feature([feature.answer_type])\n    else:\n      token_map = [-1] * len(feature.input_ids)\n      for k, v in feature.token_to_orig_map.items():\n        token_map[k] = v\n      features[\"token_map\"] = create_int_feature(token_map)\n\n    tf_example = tf.train.Example(features=tf.train.Features(feature=features))\n    self._writer.write(tf_example.SerializeToString())\n\n  def close(self):\n    self._writer.close()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"main_pred()"},{"metadata":{"trusted":true},"cell_type":"code","source":"FLAGS.vocab_file = '/kaggle/input/bert-config-vocab/vocab-nq.txt' \nFLAGS.bert_config_file = '/kaggle/input/bert-config-vocab/bert_config.json'\n\nFLAGS.do_train=False\nFLAGS.do_predict=True\n\nFLAGS.predict_batch_size=32\nFLAGS.max_seq_length=512\nFLAGS.doc_stride=256\nFLAGS.output_dir='tfrecords'\nFLAGS.use_tpu=False\n\nFLAGS.max_contexts = 80\n\neval_set = 'test'\nFLAGS.predict_file = '/kaggle/input/tensorflow2-question-answering/simplified-nq-test.jsonl'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\neval_examples = read_nq_examples(input_file=FLAGS.predict_file, is_training=False) \nprint(len(eval_examples))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ncandidates_dict = read_candidates(FLAGS.predict_file) \nprint(len(candidates_dict))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from tensorflow.contrib import cluster_resolver as contrib_cluster_resolver\nfrom tensorflow.contrib import tpu as contrib_tpu\n\ndef main_pred():\n    tf.logging.set_verbosity(tf.logging.INFO)\n    bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)\n    validate_flags_or_throw(bert_config)\n    tf.gfile.MakeDirs(FLAGS.output_dir)\n\n    tokenizer = tokenization.FullTokenizer(\n    vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)\n\n    FLAGS.use_tpu = False\n    tpu_cluster_resolver = None\n\n    is_per_host = contrib_tpu.InputPipelineConfig.PER_HOST_V2\n    run_config = contrib_tpu.RunConfig(\n    cluster=tpu_cluster_resolver,\n    master=FLAGS.master,\n    model_dir=FLAGS.output_dir,\n    keep_checkpoint_max=100,\n    save_checkpoints_steps=FLAGS.save_checkpoints_steps,\n    tpu_config=contrib_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\n    num_train_steps = None\n    num_warmup_steps = None\n\n    model_fn = 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 will fall back to normal Estimator on CPU\n    # or GPU.\n    estimator = tf.contrib.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    \n    if FLAGS.do_lower_case:\n        tfrecord_file = os.path.join(f\"{eval_set}.tf_record\")\n    else:\n        tfrecord_file = os.path.join(f\"{eval_set}_cased.tf_record\")\n        \n    if not os.path.exists(tfrecord_file):\n        eval_writer = FeatureWriter(\n          filename=tfrecord_file,\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        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\n    tf.logging.info(\"***** Running predictions *****\")\n    tf.logging.info(\"  Num orig examples = %d\", len(eval_examples))\n#     tf.logging.info(\"  Num split examples = %d\", len(eval_features))\n    tf.logging.info(\"  Batch size = %d\", FLAGS.predict_batch_size)\n\n    predict_input_fn = input_fn_builder(\n      input_file=tfrecord_file,\n      seq_length=FLAGS.max_seq_length,\n      is_training=False,\n      drop_remainder=False)\n\n    eval_features = [\n      tf.train.Example.FromString(r)\n      for r in tf.python_io.tf_record_iterator(tfrecord_file)\n    ]\n\n    # If running eval on the TPU, you will need to specify the number of\n    # steps.\n    all_results = []\n    from tqdm import tqdm\n    for result in tqdm(estimator.predict(\n      predict_input_fn, yield_single_examples=True),total=len(eval_features)):\n        if len(all_results) % 1000 == 0:\n            tf.logging.info(\"Processing example: %d\" % (len(all_results)))\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        all_results.append(\n            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    nq_pred_dict,nbest_summary_dict = 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    with open(FLAGS.output_prediction_file, \"w\") as f:\n        json.dump(predictions_json, f, indent=4)      \n\n    import pickle\n    pickle.dump(nbest_summary_dict,open(\n        FLAGS.output_prediction_file.replace('pred_jsons/pred','nbest_pkl/nbest_dict').replace('.json','.pkl'),\n        'wb'))      ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        if 'model' in dirname and 'index' in filename:\n            print(os.path.join(dirname, filename))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model L: _wwm_stride_192_neg_0.01_0.04-64-2.00E-05 (model d in writeup)\n# model O: _wwm_fix_top_level_bug_max_contexts_200_0.01_0.04-64-4.00E-05 (model c in writeup)\n# model R2: _wwm_cased_fix_top_level_bug_0.01_0.04-64-4.50E-05 (model e in writeup)\n\nckpts = ['/kaggle/input/model-l/_wwm_stride_192_neg_0.01_0.04-64-2.00E-05/model.ckpt-7000',\n         '/kaggle/input/model-o/model.ckpt-9500',\n         '/kaggle/input/model-r2/model.ckpt-8500'\n        ]\nensemble_wts = [1, 1.1, 0.8]\n\nlong_thr =   2.3972\nshort_thr = 2.895\n\nyes_thr = 5.5\nno_thr = 5.5\n\n\n# scores on dev\n# --------------------\n# LONG ANSWER R@P TABLE:\n# Optimal threshold: 2.3972\n#  F1     /  P      /  R\n#  71.62% /  71.77% /  71.46%\n# --------------------\n# SHORT ANSWER R@P TABLE:\n# Optimal threshold: 2.895\n#  F1     /  P      /  R\n#  58.50% /  70.07% /  50.20%\n# --------------------\n#  F1     /  P      /  R\n#  66.47% /  71.18% /  62.35%\n\npred_files = [x.split('/')[-2] for x in ckpts]\n\npred_files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nFLAGS.do_lower_case = True\nFLAGS.vocab_file = '/kaggle/input/bert-config-vocab/vocab-nq.txt' \nFLAGS.bert_config_file = '/kaggle/input/bert-config-vocab/bert_config.json'\nfor i in range(2):\n    FLAGS.init_checkpoint = ckpts[i]\n    FLAGS.output_prediction_file = pred_files[i]+'.json'\n    main_pred()\n    \nFLAGS.do_lower_case = False\nFLAGS.vocab_file = '/kaggle/input/bert-config-vocab/vocab_cased-nq.txt' \nFLAGS.bert_config_file = '/kaggle/input/bert-config-vocab/bert_config_cased.json'\ni=2\nFLAGS.init_checkpoint = ckpts[i]\nFLAGS.output_prediction_file = pred_files[i]+'.json'\nmain_pred()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### make sub"},{"metadata":{"trusted":true},"cell_type":"code","source":"NQLabel = collections.namedtuple(\n    'NQLabel',\n    [\n        'example_id',  # the unique id for each NQ example.\n        'long_answer_span',  # A Span object for long answer.\n        'short_answer_span_list',  # A list of Spans for short answer.\n        #   Note that In NQ, the short answers\n        #   do not need to be in a single span.\n        'yes_no_answer',  # Indicate if the short answer is an yes/no answer\n        #   The possible values are \"yes\", \"no\", \"none\".\n        #   (case insensitive)\n        #   If the field is \"yes\", short_answer_span_list\n        #   should be empty or only contain null spans.\n        'long_score',  # The prediction score for the long answer prediction.\n        'short_score'  # The prediction score for the short answer prediction.\n    ])\n\n\nclass uitl_Span(object):\n  \"\"\"A class for handling token and byte spans.\n\n    The logic is:\n\n    1) if both start_byte !=  -1 and end_byte != -1 then the span is defined\n       by byte offsets\n    2) else, if start_token != -1 and end_token != -1 then the span is define\n       by token offsets\n    3) else, this is a null span.\n\n    Null spans means that there is no (long or short) answers.\n    If your systems only care about token spans rather than byte spans, set all\n    byte spans to -1.\n\n  \"\"\"\n\n  def __init__(self, start_byte, end_byte, start_token_idx, end_token_idx):\n\n    if ((start_byte < 0 and end_byte >= 0) or\n        (start_byte >= 0 and end_byte < 0)):\n      raise ValueError('Inconsistent Null Spans (Byte).')\n\n    if ((start_token_idx < 0 and end_token_idx >= 0) or\n        (start_token_idx >= 0 and end_token_idx < 0)):\n      raise ValueError('Inconsistent Null Spans (Token).')\n\n    if start_byte >= 0 and end_byte >= 0 and start_byte >= end_byte:\n      raise ValueError('Invalid byte spans (start_byte >= end_byte).')\n\n    if ((start_token_idx >= 0 and end_token_idx >= 0) and\n        (start_token_idx >= end_token_idx)):\n      raise ValueError('Invalid token spans (start_token_idx >= end_token_idx)')\n\n    self.start_byte = start_byte\n    self.end_byte = end_byte\n    self.start_token_idx = start_token_idx\n    self.end_token_idx = end_token_idx\n\n  def is_null_span(self):\n    \"\"\"A span is a null span if the start and end are both -1.\"\"\"\n\n    if (self.start_byte < 0 and self.end_byte < 0 and\n        self.start_token_idx < 0 and self.end_token_idx < 0):\n      return True\n    return False\n\n  def __str__(self):\n    byte_str = 'byte: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'\n    tok_str = ('tok: [' + str(self.start_token_idx) + ',' + str(\n        self.end_token_idx) + ')')\n\n    return byte_str + ' ' + tok_str\n\n  def __repr__(self):\n    return self.__str__()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def ensemble_nbest_to_nq_pred_dict(nbest_lst, wts, yes_thr=None, no_thr=None):\n  def sigmoid(x): return 1 / (1 + np.exp(-x))\n    \n  assert len(nbest_lst)==len(wts) \n  nq_pred_dict = {}\n  for i in range(1,len(nbest_lst)): assert set(nbest_lst[0].keys()) == set(nbest_lst[i].keys())\n  \n  for example_id in nbest_lst[0].keys():\n    d_long = {} # key: (start_tok,end_tok), value: prob\n    d_short = {}\n    yes_logits,no_logits = 0,0\n    yes_exception, no_exception = False,False\n    for i,nbest in enumerate(nbest_lst):        \n        if yes_thr and (not yes_exception): \n            try:\n                yes_logits += nbest[example_id][0].answer_type_logits[1] * wts[i]\n            except:\n                print(\"exception yes_logits += nbest[example_id][0].answer_type_logits[1] * wts[i]\")\n                yes_exception = True\n        if no_thr and (not no_exception): \n            try:\n                no_logits += nbest[example_id][0].answer_type_logits[2] * wts[i]\n            except:\n                print(\"exception no_logits += nbest[example_id][0].answer_type_logits[2] * wts[i]\")\n                no_exception = True                \n        long_seen = set()\n        short_seen = set()\n        for score_summary in nbest[example_id]:\n          pred_lbl = score_summary.predicted_label\n          long_ans =  pred_lbl['long_answer'] \n          short_ans = pred_lbl['short_answers'][0]\n          long_key = long_ans['start_token'], long_ans['end_token']\n          short_key = short_ans['start_token'], short_ans['end_token']          \n          if long_key not in long_seen:\n              long_seen.add(long_key)\n              if long_key not in d_long: d_long[long_key] = sigmoid(pred_lbl['long_answer_score']) * wts[i]\n              else:                      d_long[long_key]+= sigmoid(pred_lbl['long_answer_score']) * wts[i]\n          if short_key not in short_seen:\n              short_seen.add(short_key)\n              if short_key not in d_short: d_short[short_key] = sigmoid(pred_lbl['short_answers_score']) * wts[i]\n              else:                        d_short[short_key]+= sigmoid(pred_lbl['short_answers_score']) * wts[i]                  \n\n    (start_tok_long,end_tok_long),prob_long = sorted(list(d_long.items()),key=lambda x:x[1],reverse=True)[0]\n    (start_tok_short,end_tok_short),prob_short = sorted(list(d_short.items()),key=lambda x:x[1],reverse=True)[0]\n\n    short_answer_span_list = [uitl_Span(-1,-1,start_tok_short,end_tok_short)]\n    yes_no_answer='none'      \n    try:\n        if yes_thr and yes_logits > yes_thr:\n            yes_no_answer,prob_short = 'yes', 999999\n            short_answer_span_list = [uitl_Span(-1,-1,-1,-1)]\n        elif no_thr and no_logits > no_thr:\n            yes_no_answer,prob_short = 'no', 999999\n            short_answer_span_list = [uitl_Span(-1,-1,-1,-1)]\n    except:\n        print('exception if yes_thr and yes_logits > yes_thr ...')\n        \n    nq_pred_dict[int(example_id)] = NQLabel(\n                        example_id=example_id,\n                        long_answer_span = uitl_Span(-1, -1, start_tok_long, end_tok_long),\n                        short_answer_span_list = short_answer_span_list,\n                        yes_no_answer=yes_no_answer,\n                        long_score=prob_long,\n                        short_score=prob_short)   \n  return nq_pred_dict  ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls -lrt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pkl_lst = [x+'.pkl' for x in pred_files]\npkl_lst","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pickle\nnbest_test_lst = [pickle.load(open(nbest_file,'rb')) for nbest_file in pkl_lst]\n\nnq_pred_dict_en = ensemble_nbest_to_nq_pred_dict(nbest_test_lst, ensemble_wts, yes_thr, no_thr)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\n\ndef nqlabel_to_dict(nqlabel):\n    nqlabel.example_id\n    nqlabel.long_answer_span\n    nqlabel.short_answer_span_list\n    nqlabel.yes_no_answer\n    nqlabel.long_score\n    nqlabel.short_score\n    return {'example_id': int(nqlabel.example_id),\n             'long_answer': {'start_token': nqlabel.long_answer_span.start_token_idx,\n              'end_token': nqlabel.long_answer_span.end_token_idx,\n              'start_byte': nqlabel.long_answer_span.start_byte,\n              'end_byte': nqlabel.long_answer_span.end_byte},\n             'long_answer_score': nqlabel.long_score,\n             'short_answers': [{'start_token': x.start_token_idx,\n               'end_token': x.end_token_idx,\n               'start_byte': x.start_byte,\n               'end_byte': x.end_byte} for x in nqlabel.short_answer_span_list],\n             'short_answers_score': nqlabel.short_score,\n             'yes_no_answer': nqlabel.yes_no_answer.upper()\n                }    \n\ntest_answers_df = pd.Series([nqlabel_to_dict(v) for (k,v) in nq_pred_dict_en.items()]).to_frame()\ntest_answers_df = test_answers_df.rename(columns={0:'predictions'})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_short_answer(entry):\n    if entry[\"short_answers_score\"] < short_thr: \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\"] < long_thr: \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)\n\ntest_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\"])\n\ntest_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\"]))\n\nsample_submission = pd.read_csv('/kaggle/input/tensorflow2-question-answering/sample_submission.csv')\n\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\n\nsample_submission.to_csv('submission.csv',index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.read_csv('submission.csv').count()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"(pd.read_csv('submission.csv').PredictionString=='YES').sum(), (pd.read_csv('submission.csv').PredictionString=='NO').sum()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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}