Function evaluate

official/projects/triviaqa/train.py:232–275

Run evaluation.

(sp_processor, features_map_fn, labels_map_fn, logits_fn,
             decode_logits_fn, split_and_pad_fn, distribute_strategy,
             validation_dataset, ground_truth)

def evaluate(sp_processor, features_map_fn, labels_map_fn, logits_fn,
             decode_logits_fn, split_and_pad_fn, distribute_strategy,
             validation_dataset, ground_truth):
  """Run evaluation."""
  loss_metric = tf_keras.metrics.Mean()

  @tf.function
  def update_loss(y, logits):
    loss_fn = modeling.SpanOrCrossEntropyLoss(
        reduction=tf_keras.losses.Reduction.NONE)
    return loss_metric(loss_fn(y, logits))

  predictions = collections.defaultdict(list)
  for _, (features, labels) in validation_dataset.enumerate():
    token_ids = features['token_ids']
    y = labels_map_fn(token_ids, labels)
    x = split_and_pad_fn(features_map_fn(features))
    logits = tf.concat(
        distribute_strategy.experimental_local_results(logits_fn(x)), 0)
    logits = logits[:features['token_ids'].shape[0]]
    update_loss(y, logits)
    end_limit = token_ids.row_lengths() - 1  # inclusive
    begin, end, scores = decode_logits_fn(logits, end_limit)
    answers = prediction.decode_answer(features['context'], begin, end,
                                       features['token_offsets'],
                                       end_limit).numpy()
    for _, (qid, token_id, offset, score, answer) in enumerate(
        zip(features['qid'].numpy(),
            tf.gather(features['token_ids'], begin, batch_dims=1).numpy(),
            tf.gather(features['token_offsets'], begin, batch_dims=1).numpy(),
            scores, answers)):
      if not answer:
        continue
      if sp_processor.IdToPiece(int(token_id)).startswith('▁') and offset > 0:
        answer = answer[1:]
      predictions[qid.decode('utf-8')].append((score, answer.decode('utf-8')))
  predictions = {
      qid: evaluation.normalize_answer(
          sorted(answers, key=operator.itemgetter(0), reverse=True)[0][1])
      for qid, answers in predictions.items()
  }
  metrics = evaluation.evaluate_triviaqa(ground_truth, predictions, mute=True)
  metrics['loss'] = loss_metric.result().numpy()
  return metrics


def main(argv):