def embed_wiki(*, data_dir, model_checkpoint, shard,
               num_shards):
  """Embed wikipedia using the given model checkpoint.

  Args:
    data_dir: Data directory for Wiki TFDS dataset
    model_checkpoint: Checkpoint of model to use
    shard: The shard to embed, intention is to run this in parallel
    num_shards: Number of shards to write.
  """
  trainer = training.Trainer.load(model_checkpoint)
  wiki_builder = wiki_tfds.WikipediaText(data_dir=data_dir)
  wiki_sents = wiki_builder.as_dataset(split='validation').shard(
      num_shards=num_shards, index=shard)
  wikipedia_urls, sentence_ids, encodings = trainer.embed_wiki_dataset(
      wiki_sents)
  enc_path = os.path.join(model_checkpoint, 'wikipedia',
                          f'embeddings_{shard}_{num_shards}.npy')
  with util.safe_open(enc_path, 'wb') as f:
    np.save(f, encodings, allow_pickle=False)
  wiki_url_path = os.path.join(model_checkpoint, 'wikipedia',
                               f'urls_{shard}_{num_shards}.npy')
  with util.safe_open(wiki_url_path, 'wb') as f:
    np.save(f, wikipedia_urls, allow_pickle=False)
  sentence_id_path = os.path.join(model_checkpoint, 'wikipedia',
                                  f'sentence_ids_{shard}_{num_shards}.npy')
  with util.safe_open(sentence_id_path, 'wb') as f:
    np.save(f, sentence_ids, allow_pickle=False)
Embed wikipedia using the given model checkpoint.

  Args:
    data_dir: Data directory for Wiki TFDS dataset
    model_checkpoint: Checkpoint of model to use
    shard: The shard to embed, intention is to run this in parallel
    num_shards: Number of shards to write.
  
language/serene/fever_cli.py