Merge pull request #114 from unslothai/feature/checkpoint-loss-in-api

feat: include training loss per checkpoint in API response

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):

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

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)")