def read_examples(
    *, embedding_dir, n_shards,
    copy_to_tmp):
  """Read and yield examples from embeddings in directory.

  Args:
    embedding_dir: The directory of .npy embedding files
    n_shards: Number of shards used to create the data
    copy_to_tmp: Whether to copy embeddings to /tmp before reading, this can
      significantly improve throughput compared to remote filesystem reads

  Yields:
    Tuples of an integer identifier, wikipedia_url, sentence_id, and embedding
  """
  idx = 0
  sharded_paths = ShardedEmbPaths(
      embedding_dir=embedding_dir, n_shards=n_shards)
  logging.info('Copying files to tmp')
  if copy_to_tmp:
    sharded_paths.to_tmp()
  logging.info('Starting example read')
  for shard in tqdm.trange(n_shards):
    emb_path = sharded_paths.shard_emb(shard, tmp=copy_to_tmp)
    urls_path = sharded_paths.shard_urls(shard, tmp=copy_to_tmp)
    sentence_ids_path = sharded_paths.shard_sentence_ids(shard, tmp=copy_to_tmp)
    logging.info('Emb path: %s', emb_path)
    logging.info('Urls path: %s', urls_path)
    logging.info('Sent path: %s', sentence_ids_path)

    with \
        util.safe_open(emb_path, 'rb') as emb_f,\
        util.safe_open(urls_path, 'rb') as url_f,\
        util.safe_open(sentence_ids_path, 'rb') as sid_f:
      load_start = time.time()
      embeddings = np.load(emb_f)
      wikipedia_urls = np.load(url_f)
      sentence_ids = np.load(sid_f)
      load_end = time.time()
      logging.info('Reading shard %s, Seconds: %s', shard,
                   load_end - load_start)
      for wiki_url, sid, emb in zip(wikipedia_urls, sentence_ids, embeddings):
        yield idx, wiki_url, sid, emb
        idx += 1
Read and yield examples from embeddings in directory.

  Args:
    embedding_dir: The directory of .npy embedding files
    n_shards: Number of shards used to create the data
    copy_to_tmp: Whether to copy embeddings to /tmp before reading, this can
      significantly improve throughput compared to remote filesystem reads

  Yields:
    Tuples of an integer identifier, wikipedia_url, sentence_id, and embedding
  
language/serene/wiki_index.py