def get_num_total_memories(config):
  """Computes the total number of mentions in the corpus."""
  logging.info('Estimating the total number of memories to be generated.')
  data_iter = get_data_iterator(config)
  start_time = time.time()
  num_total_memories = 0
  for batch in data_iter:
    batch_jax = jax.tree.map(np.asarray, batch)
    num_total_memories += batch_jax['mention_target_weights'].sum()
  logging.info('Found %d memories in %.2f seconds', num_total_memories,
               time.time() - start_time)
  return num_total_memories
Computes the total number of mentions in the corpus.
language/mentionmemory/experiments/memory_generation_main.py