Merge pull request #129 from unslothai/fix/adding-meta-data-for-checkpointing-api

Adding metadata for checkpoints

studio/backend/core/export/export.py

@@ -236,7 +236,8 @@ class ExportBackend:
                          push_to_hub: bool = False,
                          repo_id: Optional[str] = None,
                          hf_token: Optional[str] = None,
-                         private: bool = False) -> Tuple[bool, str]:
+                         private: bool = False,
+                         base_model_id: Optional[str] = None) -> Tuple[bool, str]:
         """
         Export base model (for non-PEFT models).
 
@@ -266,8 +267,8 @@ class ExportBackend:
 
                 logger.info(f"Pushing base model to Hub: {repo_id}")
 
-                # Get base model name
-                base_model = self.current_model.config._name_or_path
+                # Get base model name from request or model config
+                base_model = base_model_id or self.current_model.config._name_or_path or "unknown"
 
                 # Create repo
                 hf_api = HfApi(token=hf_token)

studio/backend/models/export.py

@@ -72,6 +72,10 @@ class ExportCommonOptions(BaseModel):
         False,
         description="If True, create a private repository on the Hub (where applicable)",
     )
+    base_model_id: Optional[str] = Field(
+        None,
+        description="HuggingFace model ID of the base model (for model card metadata)",
+    )
 
 
 class ExportMergedModelRequest(ExportCommonOptions):

studio/backend/models/models.py

@@ -21,6 +21,18 @@ class ModelCheckpoints(BaseModel):
         default_factory=list,
         description="List of checkpoints for this training run (final + intermediate)",
     )
+    base_model: Optional[str] = Field(
+        None,
+        description="Base model name from adapter_config.json or config.json",
+    )
+    peft_type: Optional[str] = Field(
+        None,
+        description="PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
+    )
+    lora_rank: Optional[int] = Field(
+        None,
+        description="LoRA rank (r) if applicable",
+    )
 
 
 class CheckpointListResponse(BaseModel):

studio/backend/routes/export.py

@@ -190,6 +190,7 @@ async def export_base_model(
             repo_id=request.repo_id,
             hf_token=request.hf_token,
             private=request.private,
+            base_model_id=request.base_model_id,
         )
 
         if not success:

studio/backend/routes/models.py

@@ -294,8 +294,11 @@ async def list_checkpoints(
                     CheckpointInfo(display_name=display_name, path=path, loss=loss)
                     for display_name, path, loss in checkpoints
                 ],
+                base_model=metadata.get("base_model"),
+                peft_type=metadata.get("peft_type"),
+                lora_rank=metadata.get("lora_rank"),
             )
-            for model_name, checkpoints in raw_models
+            for model_name, checkpoints, metadata in raw_models
         ]
 
         return CheckpointListResponse(

studio/backend/utils/models/checkpoints.py

@@ -31,12 +31,13 @@ def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
 
 def scan_checkpoints(
     outputs_dir: str = "./outputs",
-) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]]]]:
+) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]], dict]]:
     """
     Scan outputs folder for training runs and their checkpoints.
 
     Returns:
-        List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...]), ...]
+        List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...]
+        metadata keys: base_model, peft_type, lora_rank (all optional)
         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.
     """
@@ -58,6 +59,32 @@ def scan_checkpoints(
             if not (config_file.exists() or adapter_config.exists()):
                 continue
 
+            # Extract training metadata from adapter_config.json / config.json
+            metadata: dict = {}
+            try:
+                if adapter_config.exists():
+                    cfg = json.loads(adapter_config.read_text())
+                    metadata["base_model"] = cfg.get("base_model_name_or_path")
+                    metadata["peft_type"] = cfg.get("peft_type")
+                    metadata["lora_rank"] = cfg.get("r")
+                elif config_file.exists():
+                    cfg = json.loads(config_file.read_text())
+                    metadata["base_model"] = cfg.get("_name_or_path")
+            except Exception:
+                pass
+
+            # Fallback: extract base model name from folder name
+            # e.g. "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct"
+            if not metadata.get("base_model"):
+                parts = item.name.rsplit("_", 1)
+                if len(parts) == 2 and parts[1].isdigit():
+                    name_part = parts[0]
+                    idx = name_part.find("_")
+                    if idx > 0:
+                        metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1:]
+                    else:
+                        metadata["base_model"] = name_part
+
             # This is a valid training run
             checkpoints = []
 
@@ -79,7 +106,7 @@ def scan_checkpoints(
                 last_checkpoint_loss = checkpoints[-1][2]
                 checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss)
 
-            models.append((item.name, checkpoints))
+            models.append((item.name, checkpoints, metadata))
             logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
 
         # Sort by modification time (newest first)