diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 66df893f6c..2ee05090f3 100644 --- a/studio/backend/core/export/export.py +++ b/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) diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 4a3dd664ee..a0526e80db 100644 --- a/studio/backend/models/export.py +++ b/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): diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index ef883013ad..4e23c69e6d 100644 --- a/studio/backend/models/models.py +++ b/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): diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 162a71b2c4..6616c9fbd8 100644 --- a/studio/backend/routes/export.py +++ b/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: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index a5de630fee..5ed1de1ada 100644 --- a/studio/backend/routes/models.py +++ b/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( diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index cbc6c69d95..7a50d0900d 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/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)