from typing import Union
from pathlib import Path


def get_latest_save(directory: Union[str, Path]) -> Path:
    """Gets the path to the latest model checkpoint, assuming it was saved with
    the format 'model-XXXX.pt'

    Parameters
    ----------
    directory: Path or str
        The path to a directory containing the model save files.

    Returns:
        The path to the latest model save file (i.e. the one with the highest
        epoch number).
    """
    directory = Path(directory)
    saves = list(directory.glob("model-*.pt"))
    if len(saves) == 0:
        raise FileNotFoundError(f"No save files found in directory {directory}")
    return max(saves, key=lambda p: int(p.stem.split("-")[1]))