File size: 1,438 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | import os
import re
from typing import Union
def get_version_number(fname: str) -> Union[int, None]:
"""
Gets version numnber of a last checkpoint, or None if not a last checkpoint.
Args:
fname: name of the file
Returns:
version number if in the format last-v<X>.ckpt, with X and integer > 0,
or last.ckpt yields 0. If not the right naming returns None.
"""
match = re.search(r"last-v(\d+).ckpt", fname)
if match:
return int(match.group(1))
elif fname == "last.ckpt":
return 0 # last.ckpt is base version
return None # ignore if does not match
def fetch_last_ckpt(ckpt_dir: str) -> Union[str, None]:
"""
Returns the name of the latest last-v<X>.ckpt where X is an integer. Defaults to last.ckpt if just that's the only one.
If no last ckpt then returns None.
Args:
ckpt_dir: directory where checkpoints are stored.
Returns:
Name of the latest checkpoint, None if no such checkpoint present.
"""
if not os.path.exists(ckpt_dir):
return None
last_ckpts = [
f
for f in os.listdir(ckpt_dir)
if "last" in f and f.endswith(".ckpt") and get_version_number(f) is not None
]
if len(last_ckpts) == 0:
return None
sorted_files = sorted(
last_ckpts, key=get_version_number, reverse=True
) # fort by version #, highest first
return sorted_files[0]
|