def embed_claims(*, claim_tfds_data, train_claim_ids_npy_filename,
                 train_embeddings_npy_filename,
                 val_claim_ids_npy_filename,
                 val_embeddings_npy_filename, model_checkpoint):
  """Embed the claims using the given model checkpoint.

  Args:
    claim_tfds_data: path to claim tfds data
    train_claim_ids_npy_filename: Path to write train claim ids to
    train_embeddings_npy_filename: Path to write train embeddings to
    val_claim_ids_npy_filename: Path to write validation ids to
    val_embeddings_npy_filename: Path to write validation embeddings to
    model_checkpoint: The checkpoint of the model to use for embedding
  """
  logging.info('Loading model')
  trainer = training.Trainer.load(model_checkpoint)
  logging.info('Building claim datasets')
  claim_builder = claim_tfds.ClaimDataset(data_dir=claim_tfds_data)
  train_claims = claim_builder.as_dataset(split='train')
  val_claims = claim_builder.as_dataset(split='validation')
  logging.info('Embedding claims')
  train_claim_ids, train_embeddings = trainer.embed_claim_dataset(train_claims)
  val_claim_ids, val_embeddings = trainer.embed_claim_dataset(val_claims)

  out_dir = pathlib.Path(model_checkpoint) / 'claims'
  train_claim_id_path = out_dir / train_claim_ids_npy_filename
  train_emb_path = out_dir / train_embeddings_npy_filename
  val_claim_id_path = out_dir / val_claim_ids_npy_filename
  val_emb_path = out_dir / val_embeddings_npy_filename

  with util.safe_open(train_claim_id_path, 'wb') as f:
    np.save(f, train_claim_ids, allow_pickle=False)

  with util.safe_open(train_emb_path, 'wb') as f:
    np.save(f, train_embeddings, allow_pickle=False)

  with util.safe_open(val_claim_id_path, 'wb') as f:
    np.save(f, val_claim_ids, allow_pickle=False)

  with util.safe_open(val_emb_path, 'wb') as f:
    np.save(f, val_embeddings, allow_pickle=False)
Embed the claims using the given model checkpoint.

  Args:
    claim_tfds_data: path to claim tfds data
    train_claim_ids_npy_filename: Path to write train claim ids to
    train_embeddings_npy_filename: Path to write train embeddings to
    val_claim_ids_npy_filename: Path to write validation ids to
    val_embeddings_npy_filename: Path to write validation embeddings to
    model_checkpoint: The checkpoint of the model to use for embedding
  
language/serene/fever_cli.py