def load_entity_matrices(base_dir):
  """Load entity co-occurrence and co-reference matrices."""
  cooccur_ckpt = os.path.join(base_dir, "ent2ment.npz")
  coref_ckpt = os.path.join(base_dir, "coref.npz")

  tf.reset_default_graph()
  co_data, co_indices, co_rowsplits = search_utils.load_ragged_matrix(
      "ent2ment", cooccur_ckpt)
  coref_map = search_utils.load_database(
      "coref", None, coref_ckpt, dtype=tf.int32)

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    tf.logging.info("Loading ragged matrix...")
    np_data, np_indices, np_indptr = sess.run(
        [co_data, co_indices, co_rowsplits])
    tf.logging.info("Loading coref map...")
    np_coref = sess.run(coref_map)
    num_entities = np_indptr.shape[0] - 1
    num_mentions = np_coref.shape[0]
    tf.logging.info("Creating sparse matrix %d x %d...", num_entities,
                    num_mentions)
    sp_cooccur = sp.csr_matrix((np_data, np_indices, np_indptr),
                               shape=(num_entities, num_mentions))
    tf.logging.info("Creating sparse matrix %d x %d...", num_mentions,
                    num_entities)
    sp_coref = sp.csr_matrix((np.ones_like(np_coref, dtype=np.int32),
                              (np.arange(np_coref.shape[0]), np_coref)),
                             shape=(num_mentions, num_entities))

  metadata_file = os.path.join(base_dir, "entities.json")
  entity2id, _ = json.load(tf.gfile.Open(metadata_file))

  return sp_cooccur, sp_coref, entity2id
Load entity co-occurrence and co-reference matrices.
language/labs/drkit/hotpotqa/preprocessing/link_questions.py