def evaluate_entity_linking(questions, base_dir, num_hops):
  """Evaluate how often answers can be reached from linked entities."""

  def _check_answers(sp_vec, answers):
    found_ans = np.zeros((len(answers),), dtype=bool)
    for ii, ans in enumerate(answers):
      if sp_vec[0, ans] > 0.:
        found_ans[ii] = True
    return found_ans

  sp_cooccur, sp_coref, entity2id = load_entity_matrices(base_dir)
  num_found_ans = {i: 0. for i in range(num_hops + 1)}
  for ii, question in enumerate(questions):
    if (ii + 1) % 1000 == 0:
      tf.logging.info("Evaluated %d questions...", ii)
    subjects = [
        entity2id.get(ee["kb_id"].lower(), 0) for ee in question["entities"]
    ]
    answers = [
        entity2id.get(ee["kb_id"].lower(), 0)
        for ee in question["supporting_facts"]
    ]
    # Create the initial sparse vector.
    vals = np.ones((len(subjects),), dtype=np.float32)
    vals = vals / vals.sum()
    rows = np.zeros((len(subjects),), dtype=np.int32)
    cols = np.asarray(subjects, dtype=np.int32)
    v_st = sp.csr_matrix((vals, (rows, cols)), shape=(1, len(entity2id)))
    found_ans = _check_answers(v_st, answers)
    if found_ans.all():
      num_found_ans[0] += 1
      continue
    for i in range(num_hops):
      v_m = v_st * sp_cooccur
      v_st = v_m * sp_coref
      # print("then", v_st.getnnz(), v_st[0].nonzero()[:10])
      found_ans = np.logical_or(found_ans, _check_answers(v_st, answers))
      if found_ans.all():
        num_found_ans[i + 1] += 1
        break
  for i in range(num_hops + 1):
    num_found_ans[i] /= len(questions)
  return num_found_ans
Evaluate how often answers can be reached from linked entities.
language/labs/drkit/hotpotqa/preprocessing/link_questions.py