diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index d8d95e2f1a..d39815c437 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -31,6 +31,7 @@ __all__ = [ # Config "ModelConfig", "is_vision_model", + "scan_trained_models", "scan_trained_loras", "load_model_defaults", "get_base_model_from_lora", @@ -72,6 +73,7 @@ def __getattr__(name): if name in ( "is_vision_model", "ModelConfig", + "scan_trained_models", "scan_trained_loras", "load_model_defaults", "get_base_model_from_lora", @@ -79,14 +81,15 @@ def __getattr__(name): from utils.models import ( is_vision_model, ModelConfig, - scan_trained_loras, + scan_trained_models, load_model_defaults, get_base_model_from_lora, ) globals()["is_vision_model"] = is_vision_model globals()["ModelConfig"] = ModelConfig - globals()["scan_trained_loras"] = scan_trained_loras + globals()["scan_trained_models"] = scan_trained_models + globals()["scan_trained_loras"] = scan_trained_models globals()["load_model_defaults"] = load_model_defaults globals()["get_base_model_from_lora"] = get_base_model_from_lora return globals()[name] diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3f361ca5eb..cd8606f23f 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -32,8 +32,9 @@ from auth.authentication import get_current_subject # Import backend functions try: from utils.models import ( - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -62,8 +63,9 @@ except ImportError: if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from utils.models import ( - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -791,15 +793,16 @@ async def scan_loras( lora_list = [] # Scan training outputs - trained_loras = scan_trained_loras(outputs_dir = resolved_outputs_dir) - for display_name, adapter_path in trained_loras: - base_model = get_base_model_from_lora(adapter_path) + trained_models = scan_trained_models(outputs_dir = resolved_outputs_dir) + for display_name, model_path, model_type in trained_models: + base_model = get_base_model_from_checkpoint(model_path) lora_list.append( LoRAInfo( display_name = display_name, - adapter_path = adapter_path, + adapter_path = model_path, base_model = base_model, source = "training", + export_type = model_type, ) ) diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py new file mode 100644 index 0000000000..84be681fca --- /dev/null +++ b/studio/backend/tests/test_trained_model_scan.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for Studio trained-model discovery used by Chat.""" + +import json +from pathlib import Path +import sys +import types as _types +import importlib + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from unittest.mock import patch + +from utils.models.model_config import ( + ModelConfig, + get_base_model_from_checkpoint, + get_base_model_from_lora, + scan_trained_models, +) + + +def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: Path): + lora_dir = tmp_path / "unsloth_SmolLM-135M_1775412608" + lora_dir.mkdir() + (lora_dir / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (lora_dir / "adapter_model.safetensors").write_bytes(b"") + + full_dir = tmp_path / "unsloth_SmolLM-135M_full_1775412609" + full_dir.mkdir() + (full_dir / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (full_dir / "model.safetensors").write_bytes(b"") + + found = { + name: (path, model_type) + for name, path, model_type in scan_trained_models(str(tmp_path)) + } + + assert found[lora_dir.name] == (str(lora_dir), "lora") + assert found[full_dir.name] == (str(full_dir), "merged") + + +def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config( + tmp_path: Path, +): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + assert get_base_model_from_checkpoint(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M" + + +def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + assert get_base_model_from_lora(str(tmp_path)) is None + + +@patch("utils.models.model_config.is_audio_input_type", return_value = False) +@patch("utils.models.model_config.detect_audio_type", return_value = None) +@patch("utils.models.model_config.is_vision_model", return_value = False) +def test_model_config_full_finetune_local_path_is_not_lora( + _mock_vision, + _mock_audio_type, + _mock_audio_input, + tmp_path: Path, +): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "unsloth/Qwen3-4B"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + config = ModelConfig.from_identifier(str(tmp_path)) + + assert config is not None + assert config.is_lora is False + assert config.base_model is None + + +def test_scan_trained_loras_aliases_scan_trained_models(): + utils_models = importlib.import_module("utils.models") + core_module = importlib.import_module("core") + + assert utils_models.scan_trained_loras is utils_models.scan_trained_models + assert core_module.scan_trained_loras is core_module.scan_trained_models diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index a81682d8b7..808e2b012e 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -13,8 +13,9 @@ from .model_config import ( detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, load_model_config, @@ -25,6 +26,8 @@ from .model_config import ( ) from .checkpoints import scan_checkpoints +scan_trained_loras = scan_trained_models + __all__ = [ "ModelConfig", "GgufVariantInfo", @@ -33,8 +36,10 @@ __all__ = [ "detect_audio_type", "is_audio_input_type", "VALID_AUDIO_TYPES", + "scan_trained_models", "scan_trained_loras", "scan_exported_models", + "get_base_model_from_checkpoint", "load_model_defaults", "get_base_model_from_lora", "load_model_config", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5be0b183d5..81331c9dde 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1322,46 +1322,89 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def scan_trained_loras(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str]]: +def _has_model_weight_files(model_dir: Path) -> bool: + """Return True when a directory contains loadable model weights.""" + for item in model_dir.iterdir(): + if not item.is_file(): + continue + + suffix = item.suffix.lower() + if suffix == ".safetensors": + return True + if suffix == ".gguf": + return "mmproj" not in item.name.lower() + if suffix == ".bin": + name = item.name.lower() + if ( + name.startswith("pytorch_model") + or name.startswith("model") + or name.startswith("adapter_model") + or name.startswith("consolidated") + ): + return True + return False + + +def _detect_training_output_type(model_dir: Path) -> Optional[str]: + """Classify a Studio training output as LoRA or full finetune.""" + adapter_config = model_dir / "adapter_config.json" + adapter_model = model_dir / "adapter_model.safetensors" + if adapter_config.exists() or adapter_model.exists(): + return "lora" + + config_file = model_dir / "config.json" + if config_file.exists() and _has_model_weight_files(model_dir): + return "merged" + + return None + + +def _looks_like_lora_adapter(model_dir: Path) -> bool: + return model_dir.is_dir() and ( + (model_dir / "adapter_config.json").exists() + or any(model_dir.glob("adapter_model*.safetensors")) + or any(model_dir.glob("adapter_model*.bin")) + ) + + +def scan_trained_models( + outputs_dir: str = str(outputs_root()), +) -> List[Tuple[str, str, str]]: """ - Scan outputs folder for trained LoRA adapters. + Scan outputs folder for trained Studio models. Returns: - List of tuples: [(display_name, adapter_path), ...] - - Example: - [ - ("unsloth_Meta-Llama-3.1_...", "./outputs/unsloth_Meta-Llama-3.1_.../"), - ("my_finetuned_model", "./outputs/my_finetuned_model/"), - ] + List of tuples: [(display_name, model_path, model_type), ...] + model_type is "lora" for adapter runs and "merged" for full finetunes. """ - trained_loras = [] + trained_models = [] outputs_path = resolve_output_dir(outputs_dir) if not outputs_path.exists(): logger.warning(f"Outputs directory not found: {outputs_dir}") - return trained_loras + return trained_models try: for item in outputs_path.iterdir(): if item.is_dir(): - # Check if this directory contains a LoRA adapter - adapter_config = item / "adapter_config.json" - adapter_model = item / "adapter_model.safetensors" + model_type = _detect_training_output_type(item) + if model_type is None: + continue - if adapter_config.exists() or adapter_model.exists(): - display_name = item.name - adapter_path = str(item) - trained_loras.append((display_name, adapter_path)) - logger.debug(f"Found trained LoRA: {display_name}") + display_name = item.name + model_path = str(item) + trained_models.append((display_name, model_path, model_type)) + logger.debug("Found trained model: %s (%s)", display_name, model_type) # Sort by modification time (newest first) - trained_loras.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True) + trained_models.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True) logger.info( - f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}" + "Found %s trained models in %s", + len(trained_models), + outputs_dir, ) - return trained_loras + return trained_models except Exception as e: logger.error(f"Error scanning outputs folder: {e}") @@ -1494,6 +1537,68 @@ def scan_exported_models( return [] +def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]: + """Read the base model name from a local training or checkpoint directory.""" + try: + checkpoint_path_obj = Path(checkpoint_path) + + adapter_config_path = checkpoint_path_obj / "adapter_config.json" + if adapter_config_path.exists(): + with open(adapter_config_path, "r") as f: + config = json.load(f) + base_model = config.get("base_model_name_or_path") + if base_model: + logger.info( + "Detected base model from adapter_config.json: %s", base_model + ) + return base_model + + config_path = checkpoint_path_obj / "config.json" + if config_path.exists(): + with open(config_path, "r") as f: + config = json.load(f) + for key in ("model_name", "_name_or_path"): + base_model = config.get(key) + if base_model and str(base_model) != str(checkpoint_path_obj): + logger.info( + "Detected base model from config.json (%s): %s", + key, + base_model, + ) + return base_model + + training_args_path = checkpoint_path_obj / "training_args.bin" + if training_args_path.exists(): + try: + import torch + + training_args = torch.load(training_args_path) + if hasattr(training_args, "model_name_or_path"): + base_model = training_args.model_name_or_path + logger.info( + "Detected base model from training_args.bin: %s", base_model + ) + return base_model + except Exception as e: + logger.warning(f"Could not load training_args.bin: {e}") + + dir_name = checkpoint_path_obj.name + if dir_name.startswith("unsloth_"): + parts = dir_name.split("_") + if len(parts) >= 2: + model_parts = parts[1:-1] + base_model = "unsloth/" + "_".join(model_parts) + logger.info("Detected base model from directory name: %s", base_model) + return base_model + + logger.warning(f"Could not detect base model for checkpoint: {checkpoint_path}") + return None + + except Exception as e: + logger.error(f"Error reading base model from checkpoint config: {e}") + return None + + def get_base_model_from_lora(lora_path: str) -> Optional[str]: """ Read the base model name from a LoRA adapter's config. @@ -1502,16 +1607,14 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: lora_path: Path to the LoRA adapter directory Returns: - Base model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit") - or None if not found - - Example: - >>> get_base_model_from_lora("./outputs/unsloth_Meta-Llama-3.1_.../") - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit" + Base model identifier or None if not found """ try: lora_path_obj = Path(lora_path) + if not _looks_like_lora_adapter(lora_path_obj): + return None + # Try adapter_config.json first adapter_config_path = lora_path_obj / "adapter_config.json" if adapter_config_path.exists(): @@ -1884,7 +1987,11 @@ class ModelConfig: # Auto-detect LoRA for local paths (check adapter_config.json on disk) if not is_lora and is_local: - detected_base = get_base_model_from_lora(path) + detected_base = ( + get_base_model_from_lora(path) + if _looks_like_lora_adapter(Path(path)) + else None + ) if detected_base: is_lora = True logger.info( diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 441c2b48e4..08e69bbf93 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -203,11 +203,22 @@ export function ModelSelector({ ? lora.name.split("/")[0].trim() : lora.name; // Show type tag instead of base model name + const isLocal = lora.source === "local"; + const isTraining = lora.source === "training"; const isExported = lora.source === "exported"; const isMerged = lora.exportType === "merged"; - const tag = isExported - ? isMerged ? "Merged · Exported" : "LoRA" - : "LoRA"; + const isGguf = lora.exportType === "gguf"; + const tag = isLocal + ? isGguf + ? "GGUF" + : "Local" + : isTraining && isMerged + ? "Full finetune" + : isExported + ? isMerged + ? "Merged · Exported" + : "LoRA · Exported" + : "LoRA"; all.set(lora.id, { ...lora, name: displayName, diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 87d59dfaa3..e410f8cff9 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1330,7 +1330,7 @@ export function LoraModelPicker({ setQuery(event.target.value)} - placeholder="Search local adapters" + placeholder="Search trained models" className="h-9 pl-8" /> @@ -1339,7 +1339,7 @@ export function LoraModelPicker({
{grouped.length === 0 ? (
- No adapters found. + No trained models found.
) : ( grouped.map(([baseModel, adapters], index) => ( @@ -1348,9 +1348,11 @@ export function LoraModelPicker({ {baseModel} {adapters.map((adapter) => { const isLocal = adapter.source === "local"; + const isTraining = adapter.source === "training"; const isExported = adapter.source === "exported"; const isMerged = adapter.exportType === "merged"; const isGguf = adapter.exportType === "gguf"; + const isTrainingFull = isTraining && isMerged; const isLocalGgufDir = isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name)); @@ -1360,6 +1362,8 @@ export function LoraModelPicker({ : "Local" : isGguf ? "GGUF" + : isTrainingFull + ? "Full" : isExported ? isMerged ? "Merged" @@ -1369,6 +1373,8 @@ export function LoraModelPicker({ ? isLocalGgufDir ? "GGUF" : "Local" + : isTrainingFull + ? "Full finetune" : isExported ? `${tag} · Exported` : tag; diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index ebaab75c0f..165eba0f2c 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -71,6 +71,7 @@ type LoraCandidate = { id: string; baseModel: string; updatedAt?: number; + exportType?: "lora" | "merged" | "gguf"; }; function normalizeModelRef(value: string | null | undefined): string { @@ -81,12 +82,13 @@ function pickBestLoraForBase( loras: LoraCandidate[], baseModel: string | null, ): LoraCandidate | null { - if (loras.length === 0) return null; - const sorted = [...loras].sort( + const adapterOnly = loras.filter((lora) => lora.exportType === "lora"); + if (adapterOnly.length === 0) return null; + const sorted = [...adapterOnly].sort( (a, b) => (b.updatedAt ?? -1) - (a.updatedAt ?? -1), ); const normalizedBase = normalizeModelRef(baseModel); - if (!normalizedBase) return sorted[0]; + if (!normalizedBase) return sorted[0] ?? null; const exact = sorted.find( (lora) => normalizeModelRef(lora.baseModel) === normalizedBase, @@ -101,7 +103,7 @@ function pickBestLoraForBase( normalizedBase.includes(normalizedLoraBase) ); }); - return partial ?? sorted[0]; + return partial ?? sorted[0] ?? null; } function messageHasImage(message: MessageRecord): boolean { @@ -154,7 +156,8 @@ type CompareModelSelection = { function useIsLoraCompare(): boolean { return useChatRuntimeStore((s) => { const cp = s.params.checkpoint; - return cp ? s.loras.some((l) => l.id === cp) : false; + const selected = cp ? s.loras.find((l) => l.id === cp) : undefined; + return selected?.exportType === "lora"; }); } @@ -235,7 +238,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
- Fine-tuned (LoRA) + Fine-tuned
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e5c1493872..65b02be965 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -236,6 +236,25 @@ export function useChatModelRuntime() { // Apply inference defaults on reconnect (page refresh with model already loaded) if (statusRes.inference) { const currentParams = useChatRuntimeStore.getState().params; + const reconnectResponse: LoadModelResponse = { + status: "already_loaded", + model: statusRes.active_model, + display_name: statusRes.active_model, + is_vision: statusRes.is_vision, + is_lora: false, + is_gguf: statusRes.is_gguf, + is_audio: statusRes.is_audio, + audio_type: statusRes.audio_type, + has_audio_input: statusRes.has_audio_input, + inference: statusRes.inference, + context_length: statusRes.context_length, + max_context_length: statusRes.max_context_length, + native_context_length: statusRes.native_context_length, + supports_reasoning: statusRes.supports_reasoning, + reasoning_always_on: statusRes.reasoning_always_on, + supports_tools: statusRes.supports_tools, + speculative_type: statusRes.speculative_type, + }; setParams( mergeRecommendedInference(currentParams, statusRes, statusRes.active_model), ); @@ -340,8 +359,9 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.isDownloaded ?? false; const model = models.find((entry) => entry.id === modelId); const lora = loras.find((entry) => entry.id === modelId); + const loraIsAdapter = lora?.exportType === "lora"; const isLora = - explicitIsLora ?? model?.isLora ?? (lora ? true : false); + explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false; const displayName = model?.name || lora?.name || modelId; const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; @@ -355,7 +375,7 @@ export function useChatModelRuntime() { ? loras.find((entry) => entry.id === previousCheckpoint) : undefined; const previousIsLora = - previousModel?.isLora ?? (previousLora ? true : false); + previousModel?.isLora ?? (previousLora?.exportType === "lora"); // Covers Unix absolute (/), relative (./ ../), tilde (~/), Windows drive (C:\), UNC (\\server) const isLocal = /^(\/|\.{1,2}[\\/]|~[\\/]|[A-Za-z]:[\\/]|\\\\)/.test(modelId); const isCachedLora = isLora && isLocal; diff --git a/studio/frontend/src/features/chat/tour/steps.tsx b/studio/frontend/src/features/chat/tour/steps.tsx index 3c43d9fed4..e222fdc8a0 100644 --- a/studio/frontend/src/features/chat/tour/steps.tsx +++ b/studio/frontend/src/features/chat/tour/steps.tsx @@ -30,7 +30,7 @@ export function buildChatTourSteps({ body: ( <> This selects what’s loaded for inference. Hub = base models. Fine-tuned - = your LoRA adapters from Studio. + = trained Studio outputs, including LoRA adapters and full finetunes. ), }, @@ -40,9 +40,9 @@ export function buildChatTourSteps({ title: "Two tabs", body: ( <> - Hub: search Hugging Face models. Fine-tuned: adapters (LoRA) you’ve - trained locally. If results look off, compare base vs LoRA to see what - changed. + Hub: search Hugging Face models. Fine-tuned: local Studio outputs you’ve + trained or exported. If results look off, compare base vs fine-tuned + outputs to see what changed. ), onEnter: openModelSelector,