import re


def get_checkpoint_step(ckpt_path):
    """Helper function to get step number for a given checkpoint
    Based on t5.models.mtf_model_get_latest_checkpoint_from_dir

    Args:
        ckpt_path: str, checkpoint path

    Returns:
        int, checkpoint step number
    """
    return int(re.sub(".*ckpt-", "", ckpt_path))