def get_checkpoint(wrapper, model_dir, checkpoint):
  """Return checkpoint path and step, or (None, None)."""
  if checkpoint:
    checkpoint = os.path.join(model_dir, checkpoint)
  else:
    checkpoint = tf.train.latest_checkpoint(model_dir)
  step = None
  if checkpoint is not None:
    wrapper.restore_checkpoint(checkpoint)
    step = wrapper.checkpoint.current_step.numpy().item()
  print("Using checkpoint %s at step %s" % (checkpoint, step))
  return checkpoint, step
Return checkpoint path and step, or (None, None).
language/compgen/csl/model/inference/inference_utils.py