def wikimovie_eval_fn(dataset, results, name_map, output_prediction_file,
                      **kwargs):
  """Compute evaluation metrics for OneHopDataset or TwoHopDataset.

  Args:
    dataset: An object of type OneHopDataset.
    results: A list of result dicts from running estimator.predict.
    name_map: A mapping from prediction indices to text strings.
    output_prediction_file: File to store predictions to.
    **kwargs: Variable keyword arguments.

  Returns:
    metrics: A dict mapping metric names to values.
  """
  del kwargs

  # Collect ground truth answers.
  gt_answer = {ex.qas_id: ex.answer_entity for ex in dataset.examples}
  gt_ques = {ex.qas_id: ex.question_text for ex in dataset.examples}
  gt_entity = {ex.qas_id: ex.subject_entity[0] for ex in dataset.examples}
  inf_chain = {ex.qas_id: ex.inference_chain for ex in dataset.examples}

  # Compute basic metrics.
  num_correct = 0.
  all_predictions = {}
  chain2stats = {ch: [0., 0.] for ch in inf_chain.values()}
  incorrect_results, correct_results = [], []
  for result in results:
    qas_id = result["qas_ids"]
    prediction = result["predictions"]
    if prediction in gt_answer[qas_id]:
      num_correct += 1
      chain2stats[inf_chain[qas_id]][0] += 1
      correct_results.append({
          "qas_id": result["qas_ids"],
          "question": gt_ques[qas_id],
          "answers": gt_answer[qas_id],
          "subject": gt_entity[qas_id],
          "inf-chain": inf_chain[qas_id],
          "predictions": result["predictions"],
      })
      for hop in range(3):
        if "sparse_%d" % hop in result:
          correct_results[-1].update({
              "sparse_%d" % hop: result["sparse_%d" % hop],
              "dense_%d" % hop: result["dense_%d" % hop],
              "mention_%d" % hop: result["mention_%d" % hop],
              "entity_%d" % hop: result["entity_%d" % hop],
              "sparse_scores_%d" % hop: result["sparse_scores_%d" % hop],
              "dense_scores_%d" % hop: result["dense_scores_%d" % hop],
              "mention_scores_%d" % hop: result["mention_scores_%d" % hop],
              "entity_scores_%d" % hop: result["entity_scores_%d" % hop],
          })
    else:
      incorrect_results.append({
          "qas_id": result["qas_ids"],
          "question": gt_ques[qas_id],
          "answers": gt_answer[qas_id],
          "subject": gt_entity[qas_id],
          "inf-chain": inf_chain[qas_id],
          "predictions": result["predictions"],
      })
      for hop in range(3):
        if "sparse_%d" % hop in result:
          incorrect_results[-1].update({
              "sparse_%d" % hop: result["sparse_%d" % hop],
              "dense_%d" % hop: result["dense_%d" % hop],
              "mention_%d" % hop: result["mention_%d" % hop],
              "entity_%d" % hop: result["entity_%d" % hop],
              "sparse_scores_%d" % hop: result["sparse_scores_%d" % hop],
              "dense_scores_%d" % hop: result["dense_scores_%d" % hop],
              "mention_scores_%d" % hop: result["mention_scores_%d" % hop],
              "entity_scores_%d" % hop: result["entity_scores_%d" % hop],
          })
    chain2stats[inf_chain[qas_id]][1] += 1
    all_predictions[qas_id] = name_map[str(prediction)]
  accuracy = num_correct / len(all_predictions)
  json.dump(all_predictions, tf.gfile.Open(output_prediction_file, "w"))
  json.dump(
      random.sample(incorrect_results, 100),
      tf.gfile.Open(output_prediction_file + ".incorrect", "w"),
      cls=NumpyEncoder)
  json.dump(
      random.sample(correct_results, 100),
      tf.gfile.Open(output_prediction_file + ".correct", "w"),
      cls=NumpyEncoder)

  # Return metrics.
  metrics = {
      "accuracy": accuracy,
  }
  for ch, stats in chain2stats.items():
    metrics["inference-chains-acc/" + ch] = stats[0] / stats[1]
  return metrics
Compute evaluation metrics for OneHopDataset or TwoHopDataset.

  Args:
    dataset: An object of type OneHopDataset.
    results: A list of result dicts from running estimator.predict.
    name_map: A mapping from prediction indices to text strings.
    output_prediction_file: File to store predictions to.
    **kwargs: Variable keyword arguments.

  Returns:
    metrics: A dict mapping metric names to values.
  
language/labs/drkit/evaluate.py