Upload utils/io.py with huggingface_hub

utils/io.py

@@ -0,0 +1,43 @@
+"""I/O utilities for checkpoint management."""
+
+import os
+import torch
+import torch.nn as nn
+from typing import Dict, Optional
+
+
+def save_checkpoint(model: nn.Module, optimizer, scheduler, epoch: int,
+                    losses: Dict, output_path: str, config: Dict = None):
+    """Save training checkpoint."""
+    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
+
+    state = {
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'scheduler_state_dict': scheduler.state_dict(),
+        'losses': losses,
+    }
+    if config:
+        state['config'] = config
+
+    torch.save(state, output_path)
+    return output_path
+
+
+def load_checkpoint(model: nn.Module, checkpoint_path: str,
+                    optimizer=None, scheduler=None,
+                    strict: bool = True) -> Dict:
+    """Load checkpoint and optionally restore optimizer/scheduler state."""
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+
+    state_dict = checkpoint.get('model_state_dict', checkpoint)
+    model.load_state_dict(state_dict, strict=strict)
+
+    if optimizer and 'optimizer_state_dict' in checkpoint:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+    if scheduler and 'scheduler_state_dict' in checkpoint:
+        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+
+    return checkpoint