text: # -*- coding: utf-8 -*-
  #!/usr/bin/env python3
  #
  # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  import helpers.tokenization as tokenization
  import collections
  import numpy as np
  import six
  import math
  import json
  
  
  def convert_doc_tokens(paragraph_text):
      """ Return the list of tokens from the doc text """
      def is_whitespace(c):
          if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
              return True
          return False
  
      doc_tokens = []
      prev_is_whitespace = True
      for c in paragraph_text:
          if is_whitespace(c):
              prev_is_whitespace = True
          else:
              if prev_is_whitespace:
                  doc_tokens.append(c)
              else:
                  doc_tokens[-1] += c
              prev_is_whitespace = False
  
      return doc_tokens
  
  
  def _check_is_max_context(doc_spans, cur_span_index, position):
      """Check if this is the 'max context' doc span for the token."""
  
      # Because of the sliding window approach taken to scoring documents, a single
      # token can appear in multiple documents. E.g.
      #  Doc: the man went to the store and bought a gallon of milk
      #  Span A: the man went to the
      #  Span B: to the store and bought
      #  Span C: and bought a gallon of
      #  ...
      #
      # Now the word 'bought' will have two scores from spans B and C. We only
      # want to consider the score with "maximum context", which we define as
      # the *minimum* of its left and right context (the *sum* of left and
      # right context will always be the same, of course).
      #
      # In the example the maximum context for 'bought' would be span C since
      # it has 1 left context and 3 right context, while span B has 4 left context
      # and 0 right context.
      best_score = None
      best_span_index = None
      for (span_index, doc_span) in enumerate(doc_spans):
          end = doc_span.start + doc_span.length - 1
          if position < doc_span.start:
              continue
          if position > end:
              continue
          num_left_context = position - doc_span.start
          num_right_context = end - position
          score = min(num_left_context,
                      num_right_context) + 0.01 * doc_span.length
          if best_score is None or score > best_score:
              best_score = score
              best_span_index = span_index
  
      return cur_span_index == best_span_index
  
  
  def convert_example_to_features(doc_tokens, question_text, tokenizer,
                                  max_seq_length, doc_stride, max_query_length):
      """Loads a data file into a list of `InputBatch`s."""
  
      query_tokens = tokenizer.tokenize(question_text)
  
      if len(query_tokens) > max_query_length:
          query_tokens = query_tokens[0:max_query_length]
  
      tok_to_orig_index = []
      orig_to_tok_index = []
      all_doc_tokens = []
      for (i, token) in enumerate(doc_tokens):
          orig_to_tok_index.append(len(all_doc_tokens))
          sub_tokens = tokenizer.tokenize(token)
          for sub_token in sub_tokens:
              tok_to_orig_index.append(i)
              all_doc_tokens.append(sub_token)
  
      # The -3 accounts for [CLS], [SEP] and [SEP]
      max_tokens_for_doc = max_seq_length - len(query_tokens) - 3
  
      # We can have documents that are longer than the maximum sequence length.
      # To deal with this we do a sliding window approach, where we take chunks
      # of the up to our max length with a stride of `doc_stride`.
      _DocSpan = collections.namedtuple(  # pylint: disable=invalid-name
          "DocSpan", ["start", "length"])
      doc_spans = []
      start_offset = 0
      while start_offset < len(all_doc_tokens):
          length = len(all_doc_tokens) - start_offset
          if length > max_tokens_for_doc:
              length = max_tokens_for_doc
          doc_spans.append(_DocSpan(start=start_offset, length=length))
          if start_offset + length == len(all_doc_tokens):
              break
          start_offset += min(length, doc_stride)
  
      _Feature = collections.namedtuple(  # pylint: disable=invalid-name
          "Feature", [
              "input_ids", "input_mask", "segment_ids", "tokens",
              "token_to_orig_map", "token_is_max_context"
          ])
  
      features = []
      for (doc_span_index, doc_span) in enumerate(doc_spans):
          tokens = []
          token_to_orig_map = {}
          token_is_max_context = {}
          segment_ids = []
          tokens.append("[CLS]")
          segment_ids.append(0)
          for token in query_tokens:
              tokens.append(token)
              segment_ids.append(0)
          tokens.append("[SEP]")
          segment_ids.append(0)
  
          for i in range(doc_span.length):
              split_token_index = doc_span.start + i
              token_to_orig_map[len(
                  tokens)] = tok_to_orig_index[split_token_index]
  
              is_max_context = _check_is_max_context(doc_spans, doc_span_index,
                                                     split_token_index)
              token_is_max_context[len(tokens)] = is_max_context
              tokens.append(all_doc_tokens[split_token_index])
              segment_ids.append(1)
          tokens.append("[SEP]")
          segment_ids.append(1)
  
          input_ids = tokenizer.convert_tokens_to_ids(tokens)
  
          # The mask has 1 for real tokens and 0 for padding tokens. Only real
          # tokens are attended to.
          input_mask = [1] * len(input_ids)
  
          # Zero-pad up to the sequence length.
          while len(input_ids) < max_seq_length:
              input_ids.append(0)
              input_mask.append(0)
              segment_ids.append(0)
  
          assert len(input_ids) == max_seq_length
          assert len(input_mask) == max_seq_length
          assert len(segment_ids) == max_seq_length
  
          def create_int_feature(values):
              feature = np.asarray(values, dtype=np.int32, order=None)
              return feature
  
          features.append(
              _Feature(input_ids=create_int_feature(input_ids),
                       input_mask=create_int_feature(input_mask),
                       segment_ids=create_int_feature(segment_ids),
                       tokens=tokens,
                       token_to_orig_map=token_to_orig_map,
                       token_is_max_context=token_is_max_context))
      return features
  
  
  def read_squad_json(input_file):
      """read from squad json into a list of examples"""
      with open(input_file, "r", encoding='utf-8') as reader:
          input_data = json.load(reader)["data"]
  
      _Example = collections.namedtuple(  # pylint: disable=invalid-name
          "Example", ["id", "question_text", "doc_tokens"])
  
      examples = []
      for entry in input_data:
          for paragraph in entry["paragraphs"]:
              paragraph_text = paragraph["context"]
              doc_tokens = convert_doc_tokens(paragraph_text)
  
              for qa in paragraph["qas"]:
                  examples.append(
                      _Example(id=qa["id"],
                               question_text=qa["question"],
                               doc_tokens=doc_tokens))
  
      return examples
  
  
  def _get_best_indexes(logits, n_best_size):
      """Get the n-best logits from a list."""
  
      index_and_score = sorted(enumerate(logits),
                               key=lambda x: x[1],
                               reverse=True)
  
      best_indexes = []
      for i in range(len(index_and_score)):
          if i >= n_best_size:
              break
          best_indexes.append(index_and_score[i][0])
      return best_indexes
  
  
  def get_final_text(pred_text, orig_text, do_lower_case):
      """Project the tokenized prediction back to the original text."""
  
      # When we created the data, we kept track of the alignment between original
      # (whitespace tokenized) tokens and our WordPiece tokenized tokens. So
      # now `orig_text` contains the span of our original text corresponding to the
      # span that we predicted.
      #
      # However, `orig_text` may contain extra characters that we don't want in
      # our prediction.
      #
      # For example, let's say:
      #   pred_text = steve smith
      #   orig_text = Steve Smith's
      #
      # We don't want to return `orig_text` because it contains the extra "'s".
      #
      # We don't want to return `pred_text` because it's already been normalized
      # (the SQuAD eval script also does punctuation stripping/lower casing but
      # our tokenizer does additional normalization like stripping accent
      # characters).
      #
      # What we really want to return is "Steve Smith".
      #
      # Therefore, we have to apply a semi-complicated alignment heruistic between
      # `pred_text` and `orig_text` to get a character-to-charcter alignment. This
      # can fail in certain cases in which case we just return `orig_text`.
  
      def _strip_spaces(text):
          ns_chars = []
          ns_to_s_map = collections.OrderedDict()
          for (i, c) in enumerate(text):
              if c == " ":
                  continue
              ns_to_s_map[len(ns_chars)] = i
              ns_chars.append(c)
          ns_text = "".join(ns_chars)
          return (ns_text, ns_to_s_map)
  
      # We first tokenize `orig_text`, strip whitespace from the result
      # and `pred_text`, and check if they are the same length. If they are
      # NOT the same length, the heuristic has failed. If they are the same
      # length, we assume the characters are one-to-one aligned.
      tokenizer = tokenization.BasicTokenizer(do_lower_case=do_lower_case)
  
      tok_text = " ".join(tokenizer.tokenize(orig_text))
  
      start_position = tok_text.find(pred_text)
      if start_position == -1:
          return orig_text
      end_position = start_position + len(pred_text) - 1
  
      (orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text)
      (tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text)
  
      if len(orig_ns_text) != len(tok_ns_text):
          return orig_text
  
      # We then project the characters in `pred_text` back to `orig_text` using
      # the character-to-character alignment.
      tok_s_to_ns_map = {}
      for (i, tok_index) in six.iteritems(tok_ns_to_s_map):
          tok_s_to_ns_map[tok_index] = i
  
      orig_start_position = None
      if start_position in tok_s_to_ns_map:
          ns_start_position = tok_s_to_ns_map[start_position]
          if ns_start_position in orig_ns_to_s_map:
              orig_start_position = orig_ns_to_s_map[ns_start_position]
  
      if orig_start_position is None:
          return orig_text
  
      orig_end_position = None
      if end_position in tok_s_to_ns_map:
          ns_end_position = tok_s_to_ns_map[end_position]
          if ns_end_position in orig_ns_to_s_map:
              orig_end_position = orig_ns_to_s_map[ns_end_position]
  
      if orig_end_position is None:
          return orig_text
  
      output_text = orig_text[orig_start_position:(orig_end_position + 1)]
      return output_text
  
  
  def _compute_softmax(scores):
      """Compute softmax probability over raw logits."""
      if not scores:
          return []
  
      max_score = None
      for score in scores:
          if max_score is None or score > max_score:
              max_score = score
  
      exp_scores = []
      total_sum = 0.0
      for score in scores:
          x = math.exp(score - max_score)
          exp_scores.append(x)
          total_sum += x
  
      probs = []
      for score in exp_scores:
          probs.append(score / total_sum)
      return probs
  
  
  def get_predictions(doc_tokens, features, results, n_best_size,
                      max_answer_length):
      _PrelimPrediction = collections.namedtuple(  # pylint: disable=invalid-name
          "PrelimPrediction", [
              "feature_index", "start_index", "end_index", "start_logit",
              "end_logit"
          ])
  
      prediction = ""
      scores_diff_json = 0.0
  
      prelim_predictions = []
      # keep track of the minimum score of null start+end of position 0
      score_null = 1000000  # large and positive
      min_null_feature_index = 0  # the paragraph slice with min mull score
      null_start_logit = 0  # the start logit at the slice with min null score
      null_end_logit = 0  # the end logit at the slice with min null score
      version_2_with_negative = False
  
      for result in results:
          start_indexes = _get_best_indexes(result.start_logits, n_best_size)
          end_indexes = _get_best_indexes(result.end_logits, n_best_size)
          feature = features[result.feature_index]
  
          # if we could have irrelevant answers, get the min score of irrelevant
          if version_2_with_negative:
              feature_null_score = result.start_logits[0] + result.end_logits[0]
              if feature_null_score < score_null:
                  score_null = feature_null_score
                  min_null_feature_index = 0
                  null_start_logit = result.start_logits[0]
                  null_end_logit = result.end_logits[0]
  
          for start_index in start_indexes:
              for end_index in end_indexes:
                  # We could hypothetically create invalid predictions, e.g., predict
                  # that the start of the span is in the question. We throw out all
                  # invalid predictions.
                  if start_index >= len(feature.tokens):
                      continue
                  if end_index >= len(feature.tokens):
                      continue
                  if start_index not in feature.token_to_orig_map:
                      continue
                  if end_index not in feature.token_to_orig_map:
                      continue
                  if not feature.token_is_max_context.get(start_index, False):
                      continue
                  if end_index < start_index:
                      continue
                  length = end_index - start_index + 1
                  if length > max_answer_length:
                      continue
                  prelim_predictions.append(
                      _PrelimPrediction(
                          feature_index=result.feature_index,
                          start_index=start_index,
                          end_index=end_index,
                          start_logit=result.start_logits[start_index],
                          end_logit=result.end_logits[end_index]))
  
      if version_2_with_negative:
          prelim_predictions.append(
              _PrelimPrediction(feature_index=result.feature_index,
                                start_index=0,
                                end_index=0,
                                start_logit=null_start_logit,
                                end_logit=null_end_logit))
  
      prelim_predictions = sorted(prelim_predictions,
                                  key=lambda x: (x.start_logit + x.end_logit),
                                  reverse=True)
  
      _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
          "NbestPrediction", ["text", "start_logit", "end_logit"])
  
      seen_predictions = {}
      nbest = []
      for pred in prelim_predictions:
          if len(nbest) >= n_best_size:
              break
  
          if pred.start_index > 0:  # this is a non-null prediction
              feature = features[pred.feature_index]
              tok_tokens = feature.tokens[pred.start_index:(pred.end_index + 1)]
              orig_doc_start = feature.token_to_orig_map[pred.start_index]
              orig_doc_end = feature.token_to_orig_map[pred.end_index]
              orig_tokens = doc_tokens[orig_doc_start:(orig_doc_end + 1)]
              tok_text = " ".join(tok_tokens)
  
              # De-tokenize WordPieces that have been split off.
              tok_text = tok_text.replace(" ##", "")
              tok_text = tok_text.replace("##", "")
  
              # Clean whitespace
              tok_text = tok_text.strip()
              tok_text = " ".join(tok_text.split())
              orig_text = " ".join(orig_tokens)
  
              final_text = get_final_text(tok_text, orig_text, True)
              if final_text in seen_predictions:
                  continue
  
              seen_predictions[final_text] = True
          else:
              final_text = ""
              seen_predictions[final_text] = True
  
          if len(final_text):
              nbest.append(
                  _NbestPrediction(text=final_text,
                                   start_logit=pred.start_logit,
                                   end_logit=pred.end_logit))
  
      # if we didn't inlude the empty option in the n-best, inlcude it
      if version_2_with_negative:
          if "" not in seen_predictions:
              nbest.append(
                  _NbestPrediction(text="",
                                   start_logit=null_start_logit,
                                   end_logit=null_end_logit))
      # In very rare edge cases we could have no valid predictions. So we
      # just create a nonce prediction in this case to avoid failure.
      if not nbest:
          nbest.append(
              _NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0))
  
      assert len(nbest) >= 1
  
      total_scores = []
      best_non_null_entry = None
      for entry in nbest:
          total_scores.append(entry.start_logit + entry.end_logit)
          if not best_non_null_entry:
              if entry.text:
                  best_non_null_entry = entry
  
      probs = _compute_softmax(total_scores)
  
      nbest_json = []
      for (i, entry) in enumerate(nbest):
          output = collections.OrderedDict()
          output["text"] = entry.text
          output["probability"] = probs[i]
          output["start_logit"] = entry.start_logit
          output["end_logit"] = entry.end_logit
          nbest_json.append(output)
  
      assert len(nbest_json) >= 1
  
      null_score_diff_threshold = 0.0
      if not version_2_with_negative:
          prediction = nbest_json[0]["text"]
      else:
          # predict "" iff the null score - the score of best non-null > threshold
          score_diff = score_null - best_non_null_entry.start_logit - (
              best_non_null_entry.end_logit)
          scores_diff_json = score_diff
          if score_diff > null_score_diff_threshold:
              prediction = ""
          else:
              prediction = best_non_null_entry.text
  
      return prediction, nbest_json, scores_diff_json