diff options
Diffstat (limited to 'src/training/checkpointing.py')
| -rw-r--r-- | src/training/checkpointing.py | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/src/training/checkpointing.py b/src/training/checkpointing.py index 9ff02df..b53ce4f 100644 --- a/src/training/checkpointing.py +++ b/src/training/checkpointing.py @@ -6,7 +6,9 @@ Frozen models (OLMo, Qwen) are not checkpointed — they load from HuggingFace. from __future__ import annotations +import glob import os +import re from typing import Any, Optional import torch @@ -55,6 +57,35 @@ def save_checkpoint( return path +def find_latest_checkpoint(save_dir: str) -> Optional[str]: + """Find the latest checkpoint in save_dir by step number. + + Returns: + Path to latest checkpoint, or None if no checkpoints found. + """ + if not os.path.isdir(save_dir): + return None + + pattern = os.path.join(save_dir, "checkpoint_step*.pt") + files = glob.glob(pattern) + if not files: + return None + + # Extract step numbers and find max + step_re = re.compile(r"checkpoint_step(\d+)\.pt$") + best_step = -1 + best_path = None + for f in files: + m = step_re.search(f) + if m: + step = int(m.group(1)) + if step > best_step: + best_step = step + best_path = f + + return best_path + + def load_checkpoint( path: str, predictor: nn.Module, |