def rank_metrics(raw_results_list,
                 output_dir,
                 global_step,
                 eval_name,
                 save_raw=True):
  """Computes recall at N and other metrics for the ranking model."""
  to_return = {}
  predictions_per_query = collections.OrderedDict()
  n_queries = 0
  max_examples = FLAGS.max_examples
  total_reciprocal_rank = 0.0
  total_average_precision = 0.0
  for elem in raw_results_list:
    if elem.q_id not in predictions_per_query:
      predictions_per_query[elem.q_id] = []
    predictions_per_query[elem.q_id].append(elem)

  correct_top = np.zeros(max_examples)

  mrr_s = {}
  map_s = {}
  ks_of_interest = [10, 20, 100, 200, 1000]
  for k in ks_of_interest:
    mrr_s[k] = 0
    map_s[k] = 0

  for _, pred_q in predictions_per_query.items():
    sorted_pred = sorted(pred_q, key=lambda x: (x.prob), reverse=True)
    unique_sorted_pred = []
    set_cands = set()
    for cand in sorted_pred:
      if cand.cand_num in set_cands:
        continue
      set_cands.add(cand.cand_num)
      unique_sorted_pred.append(cand)

    sorted_pred = unique_sorted_pred
    cands = [x.cand_num for x in pred_q]
    s_cands = set(cands)
    if len(unique_sorted_pred) < len(s_cands):
      print("Multiple passages from document.")

    if len(s_cands) > max_examples:
      print("Too many examples!")

    correct_example_index = max_examples

    for i in range(len(sorted_pred)):
      if sorted_pred[i].label_id == 1:
        correct_example_index = i
        break
    # collect labels and scores for average precision
    y_true = []
    y_scores = []
    for elem in sorted_pred:
      y_true.append(elem.label_id)
      y_scores.append(elem.prob)

    mrr_increment = rr(y_true)
    precision_increment = 0.0
    if mrr_increment > 0:
      precision_increment = average_precision_score(y_true, y_scores)
    total_reciprocal_rank += mrr_increment
    total_average_precision += precision_increment

    # get all mrr and map quantities of interest
    for k in ks_of_interest:
      mrr_k = bounded_rr(y_true, k)
      mrr_s[k] += mrr_k
      if mrr_k > 0:
        map_s[k] += average_precision(y_true, y_scores, k)

    j = correct_example_index
    n_queries += 1
    while j < max_examples:
      correct_top[j] += 1
      j += 1

  for i in range(max_examples):
    correct_top[i] = correct_top[i] / n_queries

  print("Queries: %d", n_queries)
  print(correct_top)

  to_return["recall_at_1"] = correct_top[0]
  to_return["recall_at_3"] = correct_top[2]
  to_return["recall_at_5"] = correct_top[4]
  to_return["recall_at_10"] = correct_top[9]
  to_return["recall_at_100"] = correct_top[99]
  to_return["recall_at_200"] = correct_top[199]
  to_return["recall_at_" + str(max_examples)] = correct_top[max_examples - 1]

  to_return["map"] = total_average_precision / n_queries
  to_return["mrr"] = total_reciprocal_rank / n_queries
  for k in ks_of_interest:
    to_return["mrr_" + str(k)] = mrr_s[k] / n_queries
    to_return["map_" + str(k)] = map_s[k] / n_queries

  out_predictions = results_file_outputs(output_dir, global_step, eval_name)
  out_metrics = results_file_metrics(output_dir, global_step, eval_name)
  if save_raw:
    write_predictions(raw_results_list, out_predictions)
  write_metrics(to_return, out_metrics)
  print(to_return)
  print("max examples " + str(max_examples))
  return to_return
Computes recall at N and other metrics for the ranking model.
language/multivec/models/metrics.py