def get_latest_checkpoint(file_names, return_epoch=False):
  """Get latest checkpoint from list of file names.

  Only necessary if manually saving/loading Keras models, i.e., not using the
  tf.train.Checkpoint API.

  Args:
    file_names: List[str], file names located with the `parse_keras_models`
      method.

  Returns:
    str, the file name with the most recent checkpoint or
    Tuple[int, str], the epoch and file name of the most recent checkpoint
  """
  if not file_names:
    return None

  checkpoint_epoch_and_file_name = []
  for file_name in file_names:
    try:
      checkpoint_epoch = file_name.split('/')[-2].split('_')[-1]
    except ValueError:
      raise Exception('Expected Keras checkpoint directory path of format '
                      'gs://path_to_checkpoint/keras_model_{checkpoint_epoch}/')
    checkpoint_epoch = int(checkpoint_epoch)
    checkpoint_epoch_and_file_name.append((checkpoint_epoch, file_name))

  checkpoint_epoch_and_file_name = sorted(
      checkpoint_epoch_and_file_name, reverse=True)

  most_recent_checkpoint_epoch_and_file_name = checkpoint_epoch_and_file_name[0]
  if return_epoch:
    return most_recent_checkpoint_epoch_and_file_name
  else:
    return most_recent_checkpoint_epoch_and_file_name[1]