def tfidf_linking(questions, base_dir, tokenizer, top_k, batch_size=100):
  """Match questions to entities via Tf-IDF."""
  # Load entity ids and masks.
  tf.reset_default_graph()
  id_ckpt = os.path.join(base_dir, "entity_ids")
  entity_ids = search_utils.load_database(
      "entity_ids", None, id_ckpt, dtype=tf.int32)
  mask_ckpt = os.path.join(base_dir, "entity_mask")
  entity_mask = search_utils.load_database("entity_mask", None, mask_ckpt)
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    tf.logging.info("Loading entity ids and masks...")
    np_ent_ids, np_ent_mask = sess.run([entity_ids, entity_mask])
  tf.logging.info("Building entity count matrix...")
  entity_count_matrix = search_utils.build_count_matrix(np_ent_ids, np_ent_mask)

  # Tokenize questions and build count matrix.
  tf.logging.info("Tokenizing questions...")
  ques_toks, ques_masks = [], []
  for question in questions:
    toks = tokenizer.tokenize(question["question"])
    tok_ids = tokenizer.convert_tokens_to_ids(toks)
    ques_toks.append(tok_ids)
    ques_masks.append([1 for _ in tok_ids])
  tf.logging.info("Building question count matrix...")
  question_count_matrix = search_utils.build_count_matrix(ques_toks, ques_masks)

  # Tf-IDF.
  tf.logging.info("Computing IDFs...")
  idfs = search_utils.counts_to_idfs(entity_count_matrix, cutoff=1e-5)
  tf.logging.info("Computing entity Tf-IDFs...")
  ent_tfidfs = search_utils.counts_to_tfidf(entity_count_matrix, idfs)
  ent_tfidfs = normalize(ent_tfidfs, norm="l2", axis=0)
  tf.logging.info("Computing question TF-IDFs...")
  qry_tfidfs = search_utils.counts_to_tfidf(question_count_matrix, idfs)
  qry_tfidfs = normalize(qry_tfidfs, norm="l2", axis=0)
  tf.logging.info("Searching...")
  top_doc_indices = np.empty((len(questions), top_k), dtype=np.int32)
  top_doc_distances = np.empty((len(questions), top_k), dtype=np.float32)
  # distances = qry_tfidfs.transpose().dot(ent_tfidfs)
  num_batches = len(questions) // batch_size
  tf.logging.info("Computing distances in %d batches of size %d",
                  num_batches + 1, batch_size)
  for nb in tqdm(range(num_batches + 1)):
    min_ = nb * batch_size
    max_ = (nb + 1) * batch_size
    if min_ >= len(questions):
      break
    if max_ > len(questions):
      max_ = len(questions)
    distances = qry_tfidfs[:, min_:max_].transpose().dot(ent_tfidfs).tocsr()
    for ii in range(min_, max_):
      my_distances = distances[ii - min_, :].tocsr()
      if len(my_distances.data) <= top_k:
        o_sort = np.argsort(-my_distances.data)
        top_doc_indices[ii, :len(o_sort)] = my_distances.indices[o_sort]
        top_doc_distances[ii, :len(o_sort)] = my_distances.data[o_sort]
        top_doc_indices[ii, len(o_sort):] = 0
        top_doc_distances[ii, len(o_sort):] = 0
      else:
        o_sort = np.argpartition(-my_distances.data, top_k)[:top_k]
        top_doc_indices[ii, :] = my_distances.indices[o_sort]
        top_doc_distances[ii, :] = my_distances.data[o_sort]

  # Load entity metadata and conver to kb_id.
  metadata_file = os.path.join(base_dir, "entities.json")
  entity2id, entity2name = json.load(tf.gfile.Open(metadata_file))
  id2entity = {i: e for e, i in entity2id.items()}
  id2name = {i: entity2name[e] for e, i in entity2id.items()}
  mentions = []
  for ii in range(len(questions)):
    my_mentions = []
    for m in range(top_k):
      my_mentions.append({
          "kb_id": id2entity[top_doc_indices[ii, m]],
          "score": str(top_doc_distances[ii, m]),
          "name": id2name[top_doc_indices[ii, m]],
      })
    mentions.append(my_mentions)

  return mentions
Match questions to entities via Tf-IDF.
language/labs/drkit/hotpotqa/preprocessing/link_questions.py