diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 01f161516b..ef883013ad 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -10,6 +10,7 @@ class CheckpointInfo(BaseModel): display_name: str = Field(..., description="User-friendly checkpoint name (folder name)") path: str = Field(..., description="Full path to the checkpoint directory") + loss: Optional[float] = Field(None, description="Training loss at this checkpoint") class ModelCheckpoints(BaseModel): diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index bba6f7ae34..a5de630fee 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -291,8 +291,8 @@ async def list_checkpoints( ModelCheckpoints( name=model_name, checkpoints=[ - CheckpointInfo(display_name=display_name, path=path) - for display_name, path in checkpoints + CheckpointInfo(display_name=display_name, path=path, loss=loss) + for display_name, path, loss in checkpoints ], ) for model_name, checkpoints in raw_models diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index 3e7bd9cd6c..cbc6c69d95 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -1,19 +1,44 @@ """ Checkpoint scanning utilities for discovering training runs and their checkpoints. """ +import json import logging from pathlib import Path -from typing import List, Tuple +from typing import List, Optional, Tuple logger = logging.getLogger(__name__) -def scan_checkpoints(outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tuple[str, str]]]]: +def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]: + """ + Read the training loss from a checkpoint's trainer_state.json. + + Returns the loss from the last log_history entry, or None if unavailable. + """ + trainer_state = checkpoint_path / "trainer_state.json" + if not trainer_state.exists(): + return None + try: + with open(trainer_state) as f: + state = json.load(f) + log_history = state.get("log_history", []) + if log_history: + return log_history[-1].get("loss") + except Exception as e: + logger.debug(f"Could not read loss from {trainer_state}: {e}") + return None + + +def scan_checkpoints( + outputs_dir: str = "./outputs", +) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]]]]: """ Scan outputs folder for training runs and their checkpoints. Returns: - List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...] + List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...]), ...] + The first entry in each checkpoint list is the main adapter; its loss is + set to the loss of the last (highest-step) intermediate checkpoint. """ models = [] outputs_path = Path(outputs_dir) @@ -36,8 +61,8 @@ def scan_checkpoints(outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tup # This is a valid training run checkpoints = [] - # Add the final model checkpoint - checkpoints.append((item.name, str(item))) + # Placeholder for the main adapter — loss filled from last checkpoint below + checkpoints.append((item.name, str(item), None)) # Scan for intermediate checkpoints (checkpoint-N subdirs) for sub in sorted(item.iterdir()): @@ -46,7 +71,13 @@ def scan_checkpoints(outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tup sub_config = sub / "config.json" sub_adapter = sub / "adapter_config.json" if sub_config.exists() or sub_adapter.exists(): - checkpoints.append((sub.name, str(sub))) + loss = _read_checkpoint_loss(sub) + checkpoints.append((sub.name, str(sub), loss)) + + # Assign the last checkpoint's loss to the main adapter entry + if len(checkpoints) > 1: + last_checkpoint_loss = checkpoints[-1][2] + checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss) models.append((item.name, checkpoints)) logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")