diff --git a/install.ps1 b/install.ps1 index ee28a14e5f..3fc9ac4690 100644 --- a/install.ps1 +++ b/install.ps1 @@ -100,22 +100,115 @@ function Install-UnslothStudio { Write-Host "" # ── Helper: refresh PATH from registry (deduplicating entries) ── + # Merge order: venv Scripts (if active) > Machine > User > current $env:Path. + # Dedup compares both raw and expanded forms (%VAR% vs literal). function Refresh-SessionPath { $machine = [System.Environment]::GetEnvironmentVariable("Path", "Machine") $user = [System.Environment]::GetEnvironmentVariable("Path", "User") - $merged = "$machine;$user;$env:Path" + $venvScripts = if ($env:VIRTUAL_ENV) { Join-Path $env:VIRTUAL_ENV "Scripts" } else { $null } + $sources = @() + if ($venvScripts) { $sources += $venvScripts } + $sources += @($machine, $user, $env:Path) + $merged = ($sources | Where-Object { $_ }) -join ";" $seen = @{} - $unique = @() + $unique = New-Object System.Collections.Generic.List[string] foreach ($p in $merged -split ";") { - $key = $p.TrimEnd("\").ToLowerInvariant() - if ($key -and -not $seen.ContainsKey($key)) { - $seen[$key] = $true - $unique += $p + $rawKey = $p.Trim().Trim('"').TrimEnd("\").ToLowerInvariant() + $expKey = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd("\").ToLowerInvariant() + if ($rawKey -and -not $seen.ContainsKey($rawKey) -and -not $seen.ContainsKey($expKey)) { + $seen[$rawKey] = $true + if ($expKey -and $expKey -ne $rawKey) { $seen[$expKey] = $true } + $unique.Add($p) } } $env:Path = $unique -join ";" } + # ── Helper: safely add a directory to the persistent User PATH ── + # Direct registry access preserves REG_EXPAND_SZ (avoids dotnet/runtime#1442). + # Append (default) keeps existing tools first; Prepend for must-win entries. + function Add-ToUserPath { + param( + [Parameter(Mandatory = $true)][string]$Directory, + [ValidateSet('Append','Prepend')] + [string]$Position = 'Append' + ) + try { + $regKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Environment') + try { + $rawPath = $regKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames) + [string[]]$entries = if ($rawPath) { $rawPath -split ';' } else { @() } # string[] prevents scalar collapse + $normalDir = $Directory.Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $expNormalDir = [Environment]::ExpandEnvironmentVariables($Directory).Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $kept = New-Object System.Collections.Generic.List[string] + $matchIndices = New-Object System.Collections.Generic.List[int] + for ($i = 0; $i -lt $entries.Count; $i++) { + $stripped = $entries[$i].Trim().Trim('"') + $rawNorm = $stripped.TrimEnd('\').ToLowerInvariant() + $expNorm = [Environment]::ExpandEnvironmentVariables($stripped).TrimEnd('\').ToLowerInvariant() + $isMatch = ($rawNorm -and ($rawNorm -eq $normalDir -or $rawNorm -eq $expNormalDir)) -or + ($expNorm -and ($expNorm -eq $normalDir -or $expNorm -eq $expNormalDir)) + if ($isMatch) { + $matchIndices.Add($i) + continue + } + $kept.Add($entries[$i]) + } + $alreadyPresent = $matchIndices.Count -gt 0 + if ($alreadyPresent -and $Position -eq 'Append') { # Append: idempotent no-op + return $false + } + if ($alreadyPresent -and $Position -eq 'Prepend' -and # Prepend: no-op if already at front + $matchIndices.Count -eq 1 -and $matchIndices[0] -eq 0) { + return $false + } + # One-time backup under HKCU\Software\Unsloth\PathBackup + if ($rawPath) { + try { + $backupKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Software\Unsloth') + try { + $existingBackup = $backupKey.GetValue('PathBackup', $null) + if (-not $existingBackup) { + $backupKey.SetValue('PathBackup', $rawPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + } + } finally { + $backupKey.Close() + } + } catch { } + } + if (-not $rawPath) { + Write-Host "[WARN] User PATH is empty - initializing with $Directory" -ForegroundColor Yellow + } + $newPath = if ($rawPath) { + if ($Position -eq 'Prepend') { + (@($Directory) + $kept) -join ';' + } else { + ($kept + @($Directory)) -join ';' + } + } else { + $Directory + } + if ($newPath -ceq $rawPath) { # no actual change + return $false + } + $regKey.SetValue('Path', $newPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + # Broadcast WM_SETTINGCHANGE via dummy env-var roundtrip. + # [NullString]::Value avoids PS 7.5+/.NET 9 $null-to-"" coercion. + try { + $d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))" + [Environment]::SetEnvironmentVariable($d, '1', 'User') + [Environment]::SetEnvironmentVariable($d, [NullString]::Value, 'User') + } catch { } + return $true + } finally { + $regKey.Close() + } + } catch { + Write-Host "[WARN] Could not update User PATH: $($_.Exception.Message)" -ForegroundColor Yellow + return $false + } + } + function step { param( [Parameter(Mandatory = $true)][string]$Label, @@ -945,18 +1038,76 @@ shell.Run cmd, 0, False New-StudioShortcuts -UnslothExePath $UnslothExe - # ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ── - $ScriptsDir = Join-Path $VenvDir "Scripts" - $UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User") - if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") { - if ($UserPath) { - [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User") - } else { - [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User") + # ── Expose `unsloth` via a shim dir containing only unsloth.exe ── + # We do NOT add the venv Scripts dir to PATH (it also holds python.exe + # and pip.exe, which would hijack the user's system interpreter). + # Hardlink preferred; falls back to copy if cross-volume or non-NTFS. + # + # Remove the legacy venv Scripts PATH entry that older installers wrote. + $LegacyScriptsDir = Join-Path $VenvDir "Scripts" + try { + $legacyKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Environment') + try { + $rawPath = $legacyKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames) + if ($rawPath) { + [string[]]$pathEntries = $rawPath -split ';' + $normalLegacy = $LegacyScriptsDir.Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $expNormalLegacy = [Environment]::ExpandEnvironmentVariables($LegacyScriptsDir).Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $filtered = @($pathEntries | Where-Object { + $stripped = $_.Trim().Trim('"') + $rawNorm = $stripped.TrimEnd('\').ToLowerInvariant() + $expNorm = [Environment]::ExpandEnvironmentVariables($stripped).TrimEnd('\').ToLowerInvariant() + ($rawNorm -ne $normalLegacy -and $rawNorm -ne $expNormalLegacy) -and + ($expNorm -ne $normalLegacy -and $expNorm -ne $expNormalLegacy) + }) + $cleanedPath = $filtered -join ';' + if ($cleanedPath -ne $rawPath) { + $legacyKey.SetValue('Path', $cleanedPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + try { + $d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))" + [Environment]::SetEnvironmentVariable($d, '1', 'User') + [Environment]::SetEnvironmentVariable($d, [NullString]::Value, 'User') + } catch { } + } + } + } finally { + $legacyKey.Close() + } + } catch { } + $ShimDir = Join-Path $StudioHome "bin" + New-Item -ItemType Directory -Force -Path $ShimDir | Out-Null + $ShimExe = Join-Path $ShimDir "unsloth.exe" + # try/catch: if unsloth.exe is locked (Studio running), keep the old shim. + $shimUpdated = $false + try { + if (Test-Path $ShimExe) { Remove-Item $ShimExe -Force -ErrorAction Stop } + try { + New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null + } catch { + Copy-Item -Path $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy + } + $shimUpdated = $true + } catch { + if (Test-Path $ShimExe) { + Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow + Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow + Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow + Write-Host " Continuing with the existing launcher." -ForegroundColor Yellow + } else { + Write-Host "[WARN] Could not create unsloth launcher at $ShimExe" -ForegroundColor Yellow + Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow + Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow } - Refresh-SessionPath - step "path" "added unsloth to PATH" } + # Only add to PATH when the launcher actually exists on disk. + $pathAdded = $false + if (Test-Path $ShimExe) { + $pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend' + } + if ($shimUpdated -and $pathAdded) { + step "path" "added unsloth launcher to PATH" + } + Refresh-SessionPath # sync current session with registry # Launch studio automatically in interactive terminals; # in non-interactive environments (CI, Docker) just print instructions. diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b53fc513de..77b58e22fb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1703,6 +1703,28 @@ class LlamaCppBackend: # Wait for llama-server to become healthy if not self._wait_for_health(timeout = 600.0): self._kill_process() + _gguf = gguf_path or "" + _is_ollama = ( + ".studio_links" in _gguf + or os.sep + "ollama_links" + os.sep in _gguf + or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf + or (self._model_identifier or "").startswith("ollama/") + ) + # Only show the Ollama-specific message when the server + # output indicates a GGUF compatibility issue, not for + # unrelated failures like OOM or missing binaries. + if _is_ollama: + _output = "\n".join(self._stdout_lines[-50:]).lower() + _gguf_compat_hints = ( + "key not found", + "unknown model architecture", + "failed to load model", + ) + if any(h in _output for h in _gguf_compat_hints): + raise RuntimeError( + "Some Ollama models do not work with llama.cpp. " + "Try a different model, or use this model directly through Ollama instead." + ) raise RuntimeError( "llama-server failed to start. " "Check that the GGUF file is valid and you have enough memory." diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 9e7168eed6..db27ce1907 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -5,8 +5,11 @@ Model Management API routes """ +import hashlib +import json import os import sys +import uuid from pathlib import Path from fastapi import APIRouter, Body, Depends, HTTPException, Query from typing import List, Optional @@ -411,6 +414,267 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: return found +def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: + """Return a writable directory for Ollama ``.gguf`` symlinks. + + Prefers ``/.studio_links/`` so the links sit next to the + blobs they point at. Falls back to a per-ollama-dir namespace under + Studio's own cache when the models directory is read-only (common + for system installs under ``/usr/share/ollama`` or ``/var/lib/ollama``) + so we still surface Ollama models in those environments. + """ + from utils.paths.storage_roots import cache_root + + primary = ollama_dir / ".studio_links" + try: + primary.mkdir(exist_ok = True) + return primary + except OSError as e: + logger.debug( + "Ollama dir %s not writable for .studio_links (%s); " + "falling back to Studio cache", + ollama_dir, + e, + ) + + # Fallback: namespace by a hash of the ollama_dir so two different + # Ollama roots don't collide. This is a cache path, not a security + # boundary. + try: + digest = hashlib.sha256(str(ollama_dir.resolve()).encode()).hexdigest()[:12] + except OSError: + digest = "default" + fallback = cache_root() / "ollama_links" / digest + try: + fallback.mkdir(parents = True, exist_ok = True) + return fallback + except OSError as e: + logger.warning( + "Could not create Ollama symlink cache at %s: %s", + fallback, + e, + ) + return None + + +def _scan_ollama_dir( + ollama_dir: Path, limit: Optional[int] = None +) -> List[LocalModelInfo]: + """Scan an Ollama models directory for downloaded models. + + Ollama stores models in a content-addressable layout:: + + /manifests//// + /blobs/sha256-... + + The default host is ``registry.ollama.ai`` with namespace + ``library`` (official models), but users can pull from custom + namespaces (``mradermacher/llama3``) or entirely different hosts + (``hf.co/org/repo:tag``). We iterate all manifest files via + ``rglob`` so every layout depth is discovered. + + Each manifest is JSON with a ``layers`` array. The layer with + ``mediaType == "application/vnd.ollama.image.model"`` contains the + GGUF weights. Vision models also have a projector layer + (``application/vnd.ollama.image.projector``). We read the config + layer to extract family/size info. + + Since Ollama blobs lack a ``.gguf`` extension (which the GGUF + loading pipeline requires), we create ``.gguf``-named links + pointing at the blobs so the existing ``detect_gguf_model`` and + ``llama-server -m`` paths work unchanged. Each model gets its + own subdirectory under the links dir (keyed by a short hash of + the manifest path) so that ``detect_mmproj_file`` only sees the + projector for *that* model. Links are created as symlinks when + possible, falling back to hardlinks (Windows without Developer + Mode) as a last resort. The link dir lives under + ``/.studio_links/`` when writable, otherwise under + Studio's own cache directory. + """ + manifests_root = ollama_dir / "manifests" + if not manifests_root.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + blobs_dir = ollama_dir / "blobs" + links_root = _ollama_links_dir(ollama_dir) + if links_root is None: + logger.warning( + "Skipping Ollama scan for %s: no writable location for .gguf links", + ollama_dir, + ) + return [] + + def _make_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]: + """Create a .gguf-named link to an Ollama blob. + + Tries symlink first, then hardlink (works on Windows without + Developer Mode when target is on the same filesystem). Skips + the model if neither works -- a full file copy of a multi-GB + GGUF inside a synchronous API request would block the backend. + + Idempotent: skips recreation when a valid link already exists. + """ + link_dir.mkdir(parents = True, exist_ok = True) + link_path = link_dir / link_name + resolved = target.resolve() + + # Skip if the link already points at the exact same blob. + # Only use samefile -- size-based checks can reuse stale links + # after `ollama pull` updates a tag to a same-sized blob. + try: + if link_path.exists() and os.path.samefile(str(link_path), str(resolved)): + return str(link_path) + except OSError as e: + logger.debug("Error checking existing link %s: %s", link_path, e) + + tmp_path = link_dir / f".{link_name}.tmp-{uuid.uuid4().hex[:8]}" + try: + if tmp_path.is_symlink() or tmp_path.exists(): + tmp_path.unlink() + try: + tmp_path.symlink_to(resolved) + except OSError: + try: + os.link(str(resolved), str(tmp_path)) + except OSError: + logger.warning( + "Could not create link for Ollama blob %s " + "(symlinks and hardlinks both failed). " + "Skipping model to avoid blocking the API.", + target, + ) + return None + os.replace(str(tmp_path), str(link_path)) + return str(link_path) + except OSError as e: + logger.debug("Could not create Ollama link %s: %s", link_path, e) + try: + if tmp_path.is_symlink() or tmp_path.exists(): + tmp_path.unlink() + except OSError as cleanup_err: + logger.debug( + "Could not clean up tmp path %s: %s", tmp_path, cleanup_err + ) + return None + + try: + for tag_file in manifests_root.rglob("*"): + if not tag_file.is_file(): + continue + + rel = tag_file.relative_to(manifests_root) + parts = rel.parts + if len(parts) < 3: + continue + + host = parts[0] + repo_parts = list(parts[1:-1]) + tag = parts[-1] + + if ( + host == "registry.ollama.ai" + and repo_parts + and repo_parts[0] == "library" + ): + repo_name = "/".join(repo_parts[1:]) + elif host == "registry.ollama.ai": + repo_name = "/".join(repo_parts) + else: + repo_name = "/".join([host] + repo_parts) + + if not repo_name: + continue + + display = f"{repo_name}:{tag}" + + manifest_key = rel.as_posix() + stem_hash = hashlib.sha256(manifest_key.encode()).hexdigest()[:10] + + try: + manifest = json.loads(tag_file.read_text()) + except (json.JSONDecodeError, OSError) as e: + logger.debug( + "Skipping unreadable/invalid Ollama manifest %s: %s", + tag_file, + e, + ) + continue + + config_digest = manifest.get("config", {}).get("digest", "") + model_type = "" + file_type = "" + if config_digest and blobs_dir.is_dir(): + config_blob = blobs_dir / config_digest.replace(":", "-") + if config_blob.is_file(): + try: + cfg = json.loads(config_blob.read_text()) + model_type = cfg.get("model_type", "") + file_type = cfg.get("file_type", "") + except (json.JSONDecodeError, OSError) as e: + logger.debug( + "Could not parse Ollama config blob %s: %s", + config_blob, + e, + ) + + model_link_dir = links_root / stem_hash + + gguf_link_path: Optional[str] = None + quant = f"-{file_type}" if file_type else "" + safe_name = repo_name.replace("/", "-") + for layer in manifest.get("layers", []): + media = layer.get("mediaType", "") + digest = layer.get("digest", "") + if not digest: + continue + + if media == "application/vnd.ollama.image.model": + candidate = blobs_dir / digest.replace(":", "-") + if candidate.is_file(): + link_name = f"{safe_name}-{tag}{quant}.gguf" + gguf_link_path = _make_link( + model_link_dir, link_name, candidate + ) + + elif media == "application/vnd.ollama.image.projector": + candidate = blobs_dir / digest.replace(":", "-") + if candidate.is_file(): + mmproj_name = f"{safe_name}-{tag}-mmproj.gguf" + _make_link(model_link_dir, mmproj_name, candidate) + + if not gguf_link_path: + continue + + suffix = "" + if model_type: + suffix += f" ({model_type}" + if file_type: + suffix += f" {file_type}" + suffix += ")" + + try: + updated_at = tag_file.stat().st_mtime + except OSError: + updated_at = None + + found.append( + LocalModelInfo( + id = gguf_link_path, + model_id = f"ollama/{repo_name}:{tag}", + display_name = display + suffix, + path = gguf_link_path, + source = "custom", + updated_at = updated_at, + ), + ) + if limit is not None and len(found) >= limit: + return found + except OSError as e: + logger.warning("Error scanning Ollama directory %s: %s", ollama_dir, e) + return found + + @router.get("/local", response_model = LocalModelListResponse) async def list_local_models( models_dir: str = Query( @@ -493,11 +757,27 @@ async def list_local_models( for folder in custom_folders: folder_path = Path(folder["path"]) try: - custom_models = ( - _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) - + _scan_hf_cache(folder_path) - + _scan_lmstudio_dir(folder_path) - )[:_MAX_MODELS_PER_FOLDER] + # Ollama scanner creates .studio_links/ with .gguf symlinks. + # Filter those from the generic scanners to avoid duplicates + # and leaking internal paths into the UI. + _generic = [ + m + for m in ( + _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) + + _scan_hf_cache(folder_path) + + _scan_lmstudio_dir(folder_path) + ) + if not any( + p in (".studio_links", "ollama_links") + for p in Path(m.path).parts + ) + ] + custom_models = _generic + if len(custom_models) < _MAX_MODELS_PER_FOLDER: + custom_models += _scan_ollama_dir( + folder_path, + limit = _MAX_MODELS_PER_FOLDER - len(custom_models), + ) except OSError as e: logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) continue @@ -575,6 +855,57 @@ async def remove_scan_folder_endpoint( return {"ok": True} +@router.get("/recommended-folders") +async def get_recommended_folders( + current_subject: str = Depends(get_current_subject), +): + """Return well-known model directories that exist on this machine. + + Lightweight alternative to ``browse-folders`` for showing quick-pick + chips without the overhead of enumerating a directory tree. Returns + paths that actually exist on disk (HF cache, LM Studio, Ollama, + ``~/models``, etc.) so the frontend can offer them as one-click + "Recommended" shortcuts in the Custom Folders section. + """ + from utils.paths.storage_roots import lmstudio_model_dirs + + folders: list[str] = [] + seen: set[str] = set() + + def _add(p: Optional[Path]) -> None: + if p is None: + return + try: + resolved = str(p.resolve()) + except OSError: + return + if resolved in seen: + return + if Path(resolved).is_dir() and os.access(resolved, os.R_OK | os.X_OK): + seen.add(resolved) + folders.append(resolved) + + # LM Studio model directories + try: + for p in lmstudio_model_dirs(): + _add(p) + except Exception as e: + logger.warning("Failed to scan for LM Studio model directories: %s", e) + + # Ollama model directories + ollama_env = os.environ.get("OLLAMA_MODELS") + if ollama_env: + _add(Path(ollama_env).expanduser()) + for candidate in ( + Path.home() / ".ollama" / "models", + Path("/usr/share/ollama/.ollama/models"), + Path("/var/lib/ollama/.ollama/models"), + ): + _add(candidate) + + return {"folders": folders} + + # Heuristic ceiling on how many children to stat when checking whether a # directory "looks like" it contains models. Keeps the browser snappy # even when a directory has thousands of unrelated entries. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 44754520e3..a2d48cf009 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -959,6 +959,20 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional scan_order.append(resolved) _add(start_dir) + + # When ``path`` is a symlink (e.g. Ollama's ``.studio_links/...gguf`` + # -> ``blobs/sha256-...``), the symlink's parent directory rarely + # contains the mmproj sibling; the real mmproj file lives next to + # the symlink target. Add the target's parent to the scan so vision + # GGUFs that are surfaced via symlinks are still recognised as + # vision models. + try: + if p.is_symlink() and p.is_file(): + target_parent = p.resolve().parent + if target_parent.is_dir(): + _add(target_parent) + except OSError: + pass if search_root is not None: try: root_resolved = Path(search_root).resolve() @@ -1006,7 +1020,10 @@ def detect_gguf_model(path: str) -> Optional[str]: if p.suffix.lower() == ".gguf" and p.is_file(): if _is_mmproj(p.name): return None - return str(p.resolve()) + # Use absolute (not resolve) to preserve symlink names -- e.g. + # Ollama .studio_links/model.gguf -> blobs/sha256-... should + # keep the readable symlink name, not the opaque blob hash. + return str(p.absolute()) # Case 2: directory containing .gguf files (skip mmproj) if p.is_dir(): 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 2f661c2e72..8f318a8292 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -27,6 +27,7 @@ import { listCachedModels, listGgufVariants, listLocalModels, + listRecommendedFolders, listScanFolders, removeScanFolder, } from "@/features/chat/api/chat-api"; @@ -49,7 +50,7 @@ import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; import { Add01Icon, Cancel01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { FolderBrowser } from "./folder-browser"; -import { Trash2Icon } from "lucide-react"; +import { ChevronDownIcon, ChevronRightIcon, DownloadIcon, StarIcon, Trash2Icon } from "lucide-react"; import { type ReactNode, useCallback, @@ -73,10 +74,35 @@ function normalizeForSearch(s: string): string { return s.toLowerCase().replace(/[\s\-_\.]/g, ""); } -function ListLabel({ children }: { children: ReactNode }) { +function ListLabel({ + children, + icon, + collapsed, + onToggle, +}: { + children: ReactNode; + icon?: ReactNode; + collapsed?: boolean; + onToggle?: () => void; +}) { return ( -
- {children} +
+ + {icon} + {children} + + {onToggle && ( + + )}
); } @@ -489,6 +515,9 @@ export function HubModelPicker({ // Delete confirmation dialog state const [deleteTarget, setDeleteTarget] = useState(null); const [deleting, setDeleting] = useState(false); + const [downloadedCollapsed, setDownloadedCollapsed] = useState(false); + const [customFoldersCollapsed, setCustomFoldersCollapsed] = useState(false); + const [recommendedCollapsed, setRecommendedCollapsed] = useState(false); // Cached (already downloaded) repos -- use module-level cache so // re-mounting the popover does not flash an empty "Downloaded" section. @@ -514,6 +543,7 @@ export function HubModelPicker({ const [showFolderInput, setShowFolderInput] = useState(false); const [folderLoading, setFolderLoading] = useState(false); const [showFolderBrowser, setShowFolderBrowser] = useState(false); + const [recommendedFolders, setRecommendedFolders] = useState([]); const refreshLocalModelsList = useCallback(() => { listLocalModels() @@ -616,6 +646,9 @@ export function HubModelPicker({ // Always refresh LM Studio + custom folder models (not gated by alreadyCached) refreshLocalModelsList(); refreshScanFolders(); + listRecommendedFolders() + .then(setRecommendedFolders) + .catch(() => {}); // Always refetch cached GGUF/model lists. The module-level caches give // an instant render with stale data (no spinner flash), but newly @@ -893,8 +926,12 @@ export function HubModelPicker({ (cachedGguf.length > 0 || (!chatOnly && cachedModels.length > 0)) ? ( <> - Downloaded - {cachedGguf.map((c) => ( + } + collapsed={downloadedCollapsed} + onToggle={() => setDownloadedCollapsed((v) => !v)} + >Downloaded + {!downloadedCollapsed && cachedGguf.map((c) => (
))} - {!chatOnly && + {!downloadedCollapsed && !chatOnly && cachedModels.map((c) => (
@@ -1001,20 +1038,12 @@ export function HubModelPicker({ {!showHfSection ? ( <> -
- +
+ + Custom Folders
- + +
+
+
{/* Folder paths */} - {scanFolders.map((f) => ( + {!customFoldersCollapsed && scanFolders.map((f) => (
))} + {/* Recommended folders */} + {!customFoldersCollapsed && (() => { + const registered = new Set(scanFolders.map((f) => f.path)); + const unregistered = recommendedFolders.filter((p) => !registered.has(p)); + if (unregistered.length === 0) return null; + return ( +
+ {unregistered.map((p) => ( + + ))} +
+ ); + })()} + {/* Add folder input */} - {showFolderInput && ( + {!customFoldersCollapsed && showFolderInput && (
@@ -1114,11 +1188,15 @@ export function HubModelPicker({ {/* Models from custom folders */} - {customFolderModels.map((m) => { + {!customFoldersCollapsed && customFolderModels.map((m) => { + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); const isGguf = + isGgufFile || isGgufRepo(m.id) || - isGgufRepo(m.display_name) || - m.path.toLowerCase().endsWith(".gguf"); + isGgufRepo(m.display_name); + // Single .gguf files (e.g. Ollama blobs) load directly; + // GGUF repos/directories expand to pick a variant. + const isDirectGguf = isGgufFile; return (
{ - if (isGguf) { + if (isDirectGguf) { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + }); + } else if (isGguf) { setExpandedGguf((prev) => prev === m.id ? null : m.id, ); @@ -1158,8 +1242,12 @@ export function HubModelPicker({ {!showHfSection && cachedReady ? ( <> - Recommended - {visibleRecommendedIds.length === 0 ? ( + } + collapsed={recommendedCollapsed} + onToggle={() => setRecommendedCollapsed((v) => !v)} + >Recommended + {recommendedCollapsed ? null : visibleRecommendedIds.length === 0 ? (
No default models.
@@ -1203,7 +1291,7 @@ export function HubModelPicker({ ); }) )} - {hasMoreRecommended && ( + {!recommendedCollapsed && hasMoreRecommended && ( <>
@@ -1216,7 +1304,7 @@ export function HubModelPicker({ {showHfSection && filteredRecommendedIds.length > 0 ? ( <> - Recommended + }>Recommended {filteredRecommendedIds.map((id) => { const vram = recommendedVramMap.get(id); return ( diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 9aacfc5af4..15ac8748f0 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -262,6 +262,12 @@ export interface BrowseFoldersResponse { model_files_here?: number; } +export async function listRecommendedFolders(): Promise { + const response = await authFetch("/api/models/recommended-folders"); + const data = await parseJsonOrThrow<{ folders: string[] }>(response); + return data.folders; +} + export async function browseFolders( path?: string, showHidden = false, 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 53f0d1d352..037d3182ac 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 @@ -437,9 +437,10 @@ export function useChatModelRuntime() { const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType } = useChatRuntimeStore.getState(); // GGUF: use custom context length, or 0 = model's native context // Non-GGUF: use the Max Seq Length slider value + const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf"); const effectiveMaxSeqLength = customContextLength != null ? customContextLength - : ggufVariant != null ? (ggufContextLength ?? 0) : maxSeqLength; + : (ggufVariant != null || isDirectGgufFile) ? (ggufContextLength ?? 0) : maxSeqLength; const loadResponse = await loadModel({ model_path: modelId, hf_token: hfToken, diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 1ba41360c7..218be7e045 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -73,7 +73,109 @@ function Refresh-Environment { } $machinePath = [System.Environment]::GetEnvironmentVariable('Path', 'Machine') $userPath = [System.Environment]::GetEnvironmentVariable('Path', 'User') - $env:Path = "$machinePath;$userPath" + # Merge: venv Scripts (if active) > Machine > User > current $env:Path. Dedup raw+expanded. + $venvScripts = if ($env:VIRTUAL_ENV) { Join-Path $env:VIRTUAL_ENV 'Scripts' } else { $null } + $sources = @() + if ($venvScripts) { $sources += $venvScripts } + $sources += @($machinePath, $userPath, $env:Path) + $merged = ($sources | Where-Object { $_ }) -join ';' + $seen = @{} + $unique = New-Object System.Collections.Generic.List[string] + foreach ($p in $merged -split ";") { + $rawKey = $p.Trim().Trim('"').TrimEnd("\").ToLowerInvariant() + $expKey = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd("\").ToLowerInvariant() + if ($rawKey -and -not $seen.ContainsKey($rawKey) -and -not $seen.ContainsKey($expKey)) { + $seen[$rawKey] = $true + if ($expKey -and $expKey -ne $rawKey) { $seen[$expKey] = $true } + $unique.Add($p) + } + } + $env:Path = $unique -join ";" +} + +# ── Helper: safely add a directory to the persistent User PATH ── +# Direct registry access preserves REG_EXPAND_SZ (avoids dotnet/runtime#1442). +# Append (default) keeps existing tools first; Prepend for must-win entries. +function Add-ToUserPath { + param( + [Parameter(Mandatory = $true)][string]$Directory, + [ValidateSet('Append','Prepend')] + [string]$Position = 'Append' + ) + try { + $regKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Environment') + try { + $rawPath = $regKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames) + [string[]]$entries = if ($rawPath) { $rawPath -split ';' } else { @() } # string[] prevents scalar collapse + $normalDir = $Directory.Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $expNormalDir = [Environment]::ExpandEnvironmentVariables($Directory).Trim().Trim('"').TrimEnd('\').ToLowerInvariant() + $kept = New-Object System.Collections.Generic.List[string] + $matchIndices = New-Object System.Collections.Generic.List[int] + for ($i = 0; $i -lt $entries.Count; $i++) { + $stripped = $entries[$i].Trim().Trim('"') + $rawNorm = $stripped.TrimEnd('\').ToLowerInvariant() + $expNorm = [Environment]::ExpandEnvironmentVariables($stripped).TrimEnd('\').ToLowerInvariant() + $isMatch = ($rawNorm -and ($rawNorm -eq $normalDir -or $rawNorm -eq $expNormalDir)) -or + ($expNorm -and ($expNorm -eq $normalDir -or $expNorm -eq $expNormalDir)) + if ($isMatch) { + $matchIndices.Add($i) + continue + } + $kept.Add($entries[$i]) + } + $alreadyPresent = $matchIndices.Count -gt 0 + if ($alreadyPresent -and $Position -eq 'Append') { # Append: idempotent no-op + return $false + } + if ($alreadyPresent -and $Position -eq 'Prepend' -and # Prepend: no-op if already at front + $matchIndices.Count -eq 1 -and $matchIndices[0] -eq 0) { + return $false + } + # One-time backup under HKCU\Software\Unsloth\PathBackup + if ($rawPath) { + try { + $backupKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Software\Unsloth') + try { + $existingBackup = $backupKey.GetValue('PathBackup', $null) + if (-not $existingBackup) { + $backupKey.SetValue('PathBackup', $rawPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + } + } finally { + $backupKey.Close() + } + } catch { } + } + if (-not $rawPath) { + Write-Host "[WARN] User PATH is empty - initializing with $Directory" -ForegroundColor Yellow + } + $newPath = if ($rawPath) { + if ($Position -eq 'Prepend') { + (@($Directory) + $kept) -join ';' + } else { + ($kept + @($Directory)) -join ';' + } + } else { + $Directory + } + if ($newPath -ceq $rawPath) { # no actual change + return $false + } + $regKey.SetValue('Path', $newPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + # Broadcast WM_SETTINGCHANGE via dummy env-var roundtrip. + # [NullString]::Value avoids PS 7.5+/.NET 9 $null-to-"" coercion. + try { + $d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))" + [Environment]::SetEnvironmentVariable($d, '1', 'User') + [Environment]::SetEnvironmentVariable($d, [NullString]::Value, 'User') + } catch { } + return $true + } finally { + $regKey.Close() + } + } catch { + Write-Host "[WARN] Could not update User PATH: $($_.Exception.Message)" -ForegroundColor Yellow + return $false + } } # PowerShell 5.1 compatibility helper: avoid relying on New-TemporaryFile. @@ -493,6 +595,31 @@ if ($script:StudioVtOk -and -not $env:NO_COLOR) { Write-Host " $Rule" -ForegroundColor DarkGray } +# Back up User PATH under HKCU\Software\Unsloth before any modifications. +try { + $envKey = [Microsoft.Win32.Registry]::CurrentUser.OpenSubKey('Environment', $false) + if ($envKey) { + try { + $rawPath = $envKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames) + } finally { + $envKey.Close() + } + if ($rawPath) { + $backupKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Software\Unsloth') + try { + $existingBackup = $backupKey.GetValue('PathBackup', $null) + if (-not $existingBackup) { + $backupKey.SetValue('PathBackup', $rawPath, [Microsoft.Win32.RegistryValueKind]::ExpandString) + } + } finally { + $backupKey.Close() + } + } + } +} catch { + Write-Host "[DEBUG] Could not back up User PATH: $($_.Exception.Message)" -ForegroundColor DarkGray +} + # ========================================================================== # PHASE 1: System-level prerequisites (winget installs, env vars) # All heavy system tool installs happen here BEFORE touching Python. @@ -626,11 +753,8 @@ if (-not $HasCmake) { foreach ($d in $cmakeDefaults) { if (Test-Path (Join-Path $d "cmake.exe")) { $env:Path = "$d;$env:Path" - # Persist to user PATH so Refresh-Environment does not drop it later - $userPath = [Environment]::GetEnvironmentVariable('Path', 'User') - if (-not $userPath -or $userPath -notlike "*$d*") { - [Environment]::SetEnvironmentVariable('Path', "$d;$userPath", 'User') - } + # Persist to user PATH (Prepend so this cmake wins over older ones). + Add-ToUserPath -Directory $d -Position 'Prepend' | Out-Null $HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) if ($HasCmake) { Write-Host " Found cmake at $d (added to PATH)" -ForegroundColor Gray @@ -896,14 +1020,8 @@ $nvccBinDir = Split-Path $NvccPath -Parent if ($env:PATH -notlike "*$nvccBinDir*") { [Environment]::SetEnvironmentVariable('PATH', "$nvccBinDir;$env:PATH", 'Process') } -# Persist nvcc bin dir to User PATH so it works in new terminals -$userPath = [Environment]::GetEnvironmentVariable('Path', 'User') -if (-not $userPath -or $userPath -notlike "*$nvccBinDir*") { - if ($userPath) { - [Environment]::SetEnvironmentVariable('Path', "$nvccBinDir;$userPath", 'User') - } else { - [Environment]::SetEnvironmentVariable('Path', "$nvccBinDir", 'User') - } +# Persist nvcc bin dir (Prepend so the driver-compatible toolkit wins). +if (Add-ToUserPath -Directory $nvccBinDir -Position 'Prepend') { substep "Persisted CUDA bin dir to user PATH" } @@ -1061,15 +1179,11 @@ if ($HasPython) { $PythonOk = $true } -# Ensure Python Scripts dir is on PATH (so 'unsloth' command works in new terminals) -$ScriptsDir = python -c "import sysconfig; print(sysconfig.get_path('scripts', 'nt_user') if __import__('os').path.exists(sysconfig.get_path('scripts', 'nt_user')) else sysconfig.get_path('scripts'))" +# Add user-scheme Python Scripts dir to PATH (nt_user only, no venv fallback). +$ScriptsDir = python -c "import os, sysconfig; p = sysconfig.get_path('scripts', 'nt_user'); print(p if os.path.exists(p) else '')" if ($LASTEXITCODE -eq 0 -and $ScriptsDir -and (Test-Path $ScriptsDir)) { - $UserPath = [Environment]::GetEnvironmentVariable('Path', 'User') - $UserPathEntries = if ($UserPath) { $UserPath.Split(';') } else { @() } - if (-not ($UserPathEntries | Where-Object { $_.TrimEnd('\') -eq $ScriptsDir })) { - $newUserPath = if ($UserPath) { "$ScriptsDir;$UserPath" } else { $ScriptsDir } - [Environment]::SetEnvironmentVariable('Path', $newUserPath, 'User') - + # Append (not Prepend) -- this dir has other pip scripts; shim handles unsloth. + if (Add-ToUserPath -Directory $ScriptsDir) { # Also add to current process so it's available immediately $ProcessPathEntries = $env:PATH.Split(';') if (-not ($ProcessPathEntries | Where-Object { $_.TrimEnd('\') -eq $ScriptsDir })) { diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1e0b015c44..820a111a9b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -65,7 +65,9 @@ __all__ = [ "patch_compiled_autograd", "process_vision_info", "unsloth_compile_transformers", - "determine_attention_implementation", + "resolve_model_class", + "resolve_attention_implementation", + "resolve_encoder_attention_implementation", "_set_attn_impl", "patch_fast_lora", "validate_loftq_config", @@ -233,7 +235,7 @@ def apply_unsloth_gradient_checkpointing( # access on some GPU architectures (B200). Falls back to eager safely. _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _EAGER_ONLY_PREFIXES = ("gemma3n",) -_FLASH_ATTENTION_DISABLED_MODELS = ("gemma4", "gemma4_text") +_FLASH_ATTENTION_MAX_HEAD_DIM = 256 _FLASH_ATTENTION_DISABLED_WARNED = set() @@ -245,8 +247,102 @@ def _is_eager_only(model_type): return any(model_type.startswith(p) for p in _EAGER_ONLY_PREFIXES) -def _is_flash_attention_disabled(model_type): - return model_type in _FLASH_ATTENTION_DISABLED_MODELS +def _config_items(config): + if isinstance(config, dict): + return config.items() + if hasattr(config, "__dict__"): + return vars(config).items() + return () + + +def _config_get(config, field_name, default = None): + if isinstance(config, dict): + return config.get(field_name, default) + return getattr(config, field_name, default) + + +def _config_set(config, field_name, value): + if isinstance(config, dict): + config[field_name] = value + elif config is not None: + setattr(config, field_name, value) + + +def _iter_attention_configs(config, seen = None): + if config is None or ( + not isinstance(config, dict) and not hasattr(config, "__dict__") + ): + return + if seen is None: + seen = set() + config_id = id(config) + if config_id in seen: + return + seen.add(config_id) + yield config + + for field_name, child_config in _config_items(config): + if not isinstance(field_name, str) or not field_name.endswith("_config"): + continue + if isinstance(child_config, dict) or hasattr(child_config, "__dict__"): + yield from _iter_attention_configs(child_config, seen) + + +def _collect_attention_head_dims(config): + explicit_head_dims = [] + + for field_name in ( + "head_dim", + "global_head_dim", + "local_head_dim", + "kv_head_dim", + ): + value = _config_get(config, field_name, None) + if isinstance(value, int) and value > 0: + explicit_head_dims.append(value) + + if len(explicit_head_dims) != 0: + return explicit_head_dims + + head_dims = [] + + hidden_size_names = ("hidden_size", "d_model", "embed_dim", "dim") + num_heads_names = ("num_attention_heads", "num_heads", "n_heads") + for hidden_size_name in hidden_size_names: + hidden_size = _config_get(config, hidden_size_name, None) + if not isinstance(hidden_size, int) or hidden_size <= 0: + continue + for num_heads_name in num_heads_names: + num_heads = _config_get(config, num_heads_name, None) + if ( + isinstance(num_heads, int) + and num_heads > 0 + and (hidden_size % num_heads) == 0 + ): + head_dims.append(hidden_size // num_heads) + + return head_dims + + +def _get_max_attention_head_dim(config): + head_dims = [] + for attention_config in _iter_attention_configs(config): + head_dims.extend(_collect_attention_head_dims(attention_config)) + return max(head_dims) if len(head_dims) != 0 else None + + +def _get_flash_attention_disable_reason(config): + max_head_dim = _get_max_attention_head_dim(config) + if max_head_dim is not None and max_head_dim > _FLASH_ATTENTION_MAX_HEAD_DIM: + return ( + f"max attention head dim {max_head_dim} exceeds the Flash Attention 2 " + f"limit of {_FLASH_ATTENTION_MAX_HEAD_DIM}" + ) + return None + + +def _is_flash_attention_disabled(config): + return _get_flash_attention_disable_reason(config) is not None def _is_flash_attention_requested(attn_implementation): @@ -256,20 +352,24 @@ def _is_flash_attention_requested(attn_implementation): def _disable_flash_attention_if_needed( - model_type, config, attn_implementation = None, supports_sdpa = False, would_use_flash_attention = False, + disable_reason = None, ): - if not _is_flash_attention_disabled(model_type): + if disable_reason is None: + disable_reason = _get_flash_attention_disable_reason(config) + if disable_reason is None: return attn_implementation requested_attn_implementation = attn_implementation if requested_attn_implementation is None: - requested_attn_implementation = getattr(config, "_attn_implementation", None) + requested_attn_implementation = _config_get( + config, "_attn_implementation", None + ) if requested_attn_implementation is None: - requested_attn_implementation = getattr(config, "attn_implementation", None) + requested_attn_implementation = _config_get(config, "attn_implementation", None) if requested_attn_implementation == "eager": return _set_attn_impl(config, "eager") @@ -284,16 +384,18 @@ def _disable_flash_attention_if_needed( if _is_flash_attention_requested(requested_attn_implementation) else "flash_attention_2" ) + model_type = _config_get(config, "model_type", "") warning_key = ( model_type, logged_attn_implementation, fallback_attn_implementation, + disable_reason, ) if warning_key not in _FLASH_ATTENTION_DISABLED_WARNED: _FLASH_ATTENTION_DISABLED_WARNED.add(warning_key) print( f"Unsloth: `{logged_attn_implementation}` is not supported " - "for Gemma 4 - " + f"for `{model_type}` because {disable_reason} - " f"defaulting to `{fallback_attn_implementation}`." ) @@ -301,69 +403,125 @@ def _disable_flash_attention_if_needed( def _set_attn_impl(config, impl): - """Helper function to set attention implementation on config and return it.""" if config is not None: - setattr(config, "_attn_implementation", impl) - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", impl) + _config_set(config, "_attn_implementation", impl) + if isinstance(config, dict) or hasattr(config, "attn_implementation"): + _config_set(config, "attn_implementation", impl) return impl -def determine_attention_implementation(model_class, config): - model_type = getattr(config, "model_type", "").lower() +def resolve_model_class(auto_model, config): + mapping = getattr(auto_model, "_model_mapping", {}) + try: + result = mapping[config.__class__] + except Exception: + for config_class, model_class in mapping.items(): + if isinstance(config, config_class): + result = model_class + break + else: + return None - # Eager-only models (e.g. gemma3n timm vision towers) - if _is_eager_only(model_type): - _set_attn_impl(config, "eager") - return "eager" + return result[0] if isinstance(result, (list, tuple)) else result - # Models with known Flash Attention incompatibilities. Gemma 4 full-attention - # layers use global_head_dim=512, which exceeds Flash Attention's dense - # head-dim support. Keep explicit eager requests, otherwise prefer SDPA. - if _is_flash_attention_disabled(model_type): + +def resolve_attention_implementation( + model_class, + config, + requested_attn_implementation = None, + supports_sdpa = None, +): + model_type_name = _config_get(config, "model_type", "") + model_type = model_type_name.lower() + if supports_sdpa is None: supports_sdpa = model_class is not None and getattr( model_class, "_supports_sdpa", False ) - return _disable_flash_attention_if_needed( - model_type, + supports_flash_attention = model_class is not None and ( + getattr(model_class, "_supports_flash_attn_2", False) + or getattr(model_class, "_supports_flash_attn", False) + ) + disable_reason = _get_flash_attention_disable_reason(config) + flash_attention_disabled = disable_reason is not None + + if model_class is None: + attn_impl = _set_attn_impl(config, "sdpa" if supports_sdpa else "eager") + else: + if _is_eager_only(model_type): + attn_impl = _set_attn_impl(config, "eager") + elif flash_attention_disabled: + attn_impl = _disable_flash_attention_if_needed( + config, + supports_sdpa = supports_sdpa, + would_use_flash_attention = ( + HAS_FLASH_ATTENTION and supports_flash_attention + ), + disable_reason = disable_reason, + ) + elif HAS_FLASH_ATTENTION and supports_flash_attention: + attn_impl = _set_attn_impl(config, "flash_attention_2") + elif supports_sdpa: + attn_impl = _set_attn_impl(config, "sdpa") + else: + attn_impl = "eager" + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": + try: + from transformers.utils.import_utils import ( + is_torch_flex_attn_available, + ) + + if ( + is_torch_flex_attn_available() + and getattr(model_class, "_supports_flex_attn", False) + and not _is_flex_excluded(model_type) + ): + attention_dropout = ( + _config_get(config, "attention_dropout", 0) or 0 + ) + if attention_dropout == 0: + attn_impl = _set_attn_impl(config, "flex_attention") + except Exception: + pass + if attn_impl == "eager": + attn_impl = _set_attn_impl(config, "eager") + + if requested_attn_implementation is None: + final_attn_impl = attn_impl + elif flash_attention_disabled: + final_attn_impl = _disable_flash_attention_if_needed( config, + requested_attn_implementation, supports_sdpa = supports_sdpa, + disable_reason = disable_reason, ) + else: + final_attn_impl = requested_attn_implementation + _set_attn_impl(config, final_attn_impl) - # Flash Attention 2 - if HAS_FLASH_ATTENTION and model_class is not None: - supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr( - model_class, "_supports_flash_attn", False + if not supports_sdpa and final_attn_impl == "sdpa": + print( + f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager." ) - if supports_fa2: - _set_attn_impl(config, "flash_attention_2") - return "flash_attention_2" + final_attn_impl = _set_attn_impl(config, "eager") - # Flex Attention - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": - try: - from transformers.utils.import_utils import is_torch_flex_attn_available + return final_attn_impl - if ( - is_torch_flex_attn_available() - and model_class is not None - and getattr(model_class, "_supports_flex_attn", False) - and not _is_flex_excluded(model_type) - ): - attention_dropout = getattr(config, "attention_dropout", 0) or 0 - if attention_dropout == 0: - _set_attn_impl(config, "flex_attention") - return "flex_attention" - except Exception: - pass - # SDPA - if model_class is not None and getattr(model_class, "_supports_sdpa", False): - _set_attn_impl(config, "sdpa") +def resolve_encoder_attention_implementation( + auto_model, + config, + model_type = "", + disable_sdpa_model_names = (), +): + model_class = resolve_model_class(auto_model, config) + supports_sdpa = model_class is not None and getattr( + model_class, "_supports_sdpa", False + ) + if any(name in model_type.lower() for name in disable_sdpa_model_names): + return "eager" + if supports_sdpa: return "sdpa" - - _set_attn_impl(config, "eager") - return "eager" + return None def _run_temporary_patches(phase): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d145f08bd5..cc021f7b97 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2346,7 +2346,7 @@ class FastLlamaModel: model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__] IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1") - preferred_attn_impl = determine_attention_implementation( + preferred_attn_impl = resolve_attention_implementation( model_function, model_config ) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index cd12544ae9..fc91178d88 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -1151,9 +1151,6 @@ class FastModel(FastBaseModel): ) os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1" - # Disable flex_attention for Gemma-4: flex compile overhead is 2.7x slower - # than SDPA. Our attention patch ensures Q/K/V dtype alignment for SDPA. - os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0" # Gemma 3N must be before Gemma 3 elif "gemma3n" in model_types_all: if transformers_version < Version("4.53.0"): diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index ad59165a50..541875a3fc 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -15,7 +15,11 @@ import logging from .loader import FastModel, DISABLE_SDPA_MODEL_NAMES -from ._utils import SUPPORTS_BFLOAT16 +from ._utils import ( + SUPPORTS_BFLOAT16, + resolve_model_class, + resolve_encoder_attention_implementation, +) import inspect import json import os @@ -31,7 +35,6 @@ import transformers from packaging.version import Version import re from transformers import AutoModel, AutoConfig -from transformers.models.auto.auto_factory import _get_model_class import tempfile from huggingface_hub import HfApi, get_token from ..save import unsloth_save_pretrained_torchao, unsloth_save_pretrained_gguf @@ -870,7 +873,7 @@ class FastSentenceTransformer(FastModel): if auto_model_class is None: auto_model_class = AutoModel # try to resolve the class - model_class = _get_model_class(config, auto_model_class._model_mapping) + model_class = resolve_model_class(auto_model_class, config) if model_class: sig = inspect.signature(model_class.__init__) @@ -1446,32 +1449,18 @@ class FastSentenceTransformer(FastModel): ): st_device = "cuda" - # Check if model supports SDPA (Scaled Dot Product Attention) for extra speedup - supports_sdpa = False - if config is not None: - try: - model_class = _get_model_class( - config, kwargs.get("auto_model", AutoModel)._model_mapping - ) - supports_sdpa = getattr(model_class, "_supports_sdpa", False) - except: - pass - # Build model_kwargs for SentenceTransformer model_kwargs = {"torch_dtype": dtype} - # Enable SDPA if supported (1.2x extra speedup on top of torch.compile) - # But disable for models with known SDPA + torch.compile backward issues - _force_eager = False - for _sdpa_model in DISABLE_SDPA_MODEL_NAMES: - if _sdpa_model in model_type.lower(): - supports_sdpa = False - _force_eager = True - break - if supports_sdpa: - model_kwargs["attn_implementation"] = "sdpa" - elif _force_eager: - model_kwargs["attn_implementation"] = "eager" + encoder_attn_impl = resolve_encoder_attention_implementation( + kwargs.get("auto_model", AutoModel), + config, + model_type = model_type, + disable_sdpa_model_names = DISABLE_SDPA_MODEL_NAMES, + ) + supports_sdpa = encoder_attn_impl == "sdpa" + if encoder_attn_impl is not None: + model_kwargs["attn_implementation"] = encoder_attn_impl # Print optimization status sdpa_str = " + SDPA" if supports_sdpa else "" diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index ba23a575c8..1bdfa9b583 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -33,8 +33,8 @@ from ._utils import ( __version__, importlib_version, _prepare_model_for_qat, - _is_flash_attention_disabled, - _disable_flash_attention_if_needed, + resolve_model_class, + resolve_attention_implementation, ) from ._utils import * from .loader_utils import _get_fp8_mode_and_check_settings @@ -739,55 +739,18 @@ class FastBaseModel: token = token, trust_remote_code = trust_remote_code, ) - user_attn_implementation = kwargs.get("attn_implementation", None) - try: - model_class = auto_model._model_mapping[auto_config.__class__] - except Exception: - model_class = None - if model_class is None: - # When model_class cannot be resolved (remote-code or unmapped - # configs), preserve the old fallback of sdpa when supported. - attn_impl = _set_attn_impl( - auto_config, "sdpa" if supports_sdpa else "eager" - ) - else: - attn_impl = determine_attention_implementation(model_class, auto_config) + model_class = resolve_model_class(auto_model, auto_config) + attn_impl = resolve_attention_implementation( + model_class, + auto_config, + requested_attn_implementation = kwargs.get("attn_implementation", None), + supports_sdpa = supports_sdpa, + ) # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with # FP8 weights. We just need to update it here for sanity. auto_config.model_name = model_name - # Re-resolve model_class after potential config change - try: - model_class = auto_model._model_mapping[auto_config.__class__] - except Exception: - model_class = None - - if not ("attn_implementation" in kwargs): - kwargs["attn_implementation"] = attn_impl - model_type = getattr(auto_config, "model_type", "").lower() - if _is_flash_attention_disabled(model_type): - supports_fa2 = model_class is not None and ( - getattr(model_class, "_supports_flash_attn_2", False) - or getattr(model_class, "_supports_flash_attn", False) - ) - kwargs["attn_implementation"] = _disable_flash_attention_if_needed( - model_type, - auto_config, - kwargs.get("attn_implementation"), - supports_sdpa = supports_sdpa, - would_use_flash_attention = ( - user_attn_implementation is None - and HAS_FLASH_ATTENTION - and supports_fa2 - ), - ) - if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa": - print( - f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." - ) - del kwargs["attn_implementation"] - # Re-stamp config so it stays consistent with the actual impl - _set_attn_impl(auto_config, "eager") + kwargs["attn_implementation"] = attn_impl bnb_config = None user_quantization_config = kwargs.get("quantization_config", None) @@ -930,9 +893,7 @@ class FastBaseModel: token = token, trust_remote_code = trust_remote_code, ) - setattr(auto_config, "_attn_implementation", config_attn_impl) - if hasattr(auto_config, "attn_implementation"): - setattr(auto_config, "attn_implementation", config_attn_impl) + _set_attn_impl(auto_config, config_attn_impl) model_config = auto_config verify_fp8_support_if_applicable(model_config) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 4fc09ed76b..85b3c5308b 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -636,173 +636,639 @@ def load_correct_tokenizer( return tokenizer -def _find_end_position(template, endfor, endif): - where_endfor = template.find(endfor) - where_endif = template.find(endif) - if where_endfor == where_endif == -1: +# All four Jinja whitespace-control variants of endfor/endif: +# {% endfor %} {%- endfor %} {% endfor -%} {%- endfor -%} +_RE_ENDFOR = re.compile(r"\{%(-?)\s*endfor\s*(-?)%\}") +_RE_ENDIF = re.compile(r"\{%(-?)\s*endif\s*(-?)%\}") +_RE_JINJA_COMMENT = re.compile(r"\{#.*?#\}", flags = re.DOTALL) + + +def _find_end_position(template, endfor = None, endif = None): + """Rightmost {% endfor %}/{% endif %} (any dash variant), as a dict + with start/end/text/dash_left/dash_right. Tokens inside Jinja comments + are ignored. `endfor`/`endif` kwargs kept for back-compat, ignored.""" + # Space-pad comments so positions still map 1:1 to the original. + scrubbed = _RE_JINJA_COMMENT.sub(lambda m: " " * len(m.group(0)), template) + endfor_matches = list(_RE_ENDFOR.finditer(scrubbed)) + endif_matches = list(_RE_ENDIF.finditer(scrubbed)) + last_endfor = endfor_matches[-1] if endfor_matches else None + last_endif = endif_matches[-1] if endif_matches else None + candidates = [m for m in (last_endfor, last_endif) if m is not None] + if not candidates: return None - elif where_endfor > where_endif: - return endfor + m = max(candidates, key = lambda x: x.end()) + return { + "start": m.start(), + "end": m.end(), + "text": m.group(0), + "dash_left": bool(m.group(1)), + "dash_right": bool(m.group(2)), + } + + +def _template_ends_with_toplevel_for(chat_template): + """Return True if the last structural node at the template's top level is + a For (message-iteration) loop, ignoring trailing pure-whitespace Output + nodes. Unwraps benign outer-If guards (no else branch, not testing + add_generation_prompt) so that templates like + ``{% if messages %}{% for ... %}{% endfor %}{% endif %}`` are still + repairable. Rejects real structural wrappers (e.g. Qwen3-Guard with + else branches).""" + try: + import jinja2 + import jinja2.nodes + + ast = jinja2.Environment().parse(chat_template) + except Exception: + return False + + def _last_structural(nodes): + for node in reversed(nodes): + if isinstance(node, jinja2.nodes.Output): + only_ws = all( + isinstance(child, jinja2.nodes.TemplateData) + and child.data.strip() == "" + for child in node.nodes + ) + if only_ws: + continue + return node + return None + + node = _last_structural(ast.body) + while isinstance(node, jinja2.nodes.If) and not node.else_: + names = [] + if isinstance(node.test, jinja2.nodes.Name): + names.append(node.test) + names.extend(node.test.find_all(jinja2.nodes.Name)) + if any(n.name == "add_generation_prompt" for n in names): + break + node = _last_structural(node.body) + + return isinstance(node, jinja2.nodes.For) + + +def _if_body_emits_content(if_node): + """True if the If's body contains any Output node (directly or nested). + Distinguishes a real generation block from a header guard that only + does `{% set ... %}`.""" + import jinja2.nodes + + for node in if_node.body: + if isinstance(node, jinja2.nodes.Output): + return True + if any( + isinstance(d, jinja2.nodes.Output) + for d in node.find_all(jinja2.nodes.Output) + ): + return True + return False + + +def _has_add_generation_prompt_block(chat_template): + """True if the template has a *positive* `{% if add_generation_prompt %}` + gate whose body emits output. Rejects header guards like + `{% if not add_generation_prompt is defined %}{% set ... %}{% endif %}` + that reference the name but emit nothing. AST-based; string-scan + fallback if Jinja fails to parse.""" + try: + import jinja2 + import jinja2.nodes + + ast = jinja2.Environment().parse(chat_template) + except Exception: + return "if add_generation_prompt" in chat_template and "%}" in chat_template + for if_node in ast.find_all(jinja2.nodes.If): + test = if_node.test + # Reject negated gates: `{% if not add_generation_prompt %}` fires + # when agp=False, so it's not a generation block even if it emits. + if isinstance(test, jinja2.nodes.Not): + continue + # find_all skips the test root, so check bare Name tests explicitly. + references_agp = False + if isinstance(test, jinja2.nodes.Name) and test.name == "add_generation_prompt": + references_agp = True + else: + for name_node in test.find_all(jinja2.nodes.Name): + if name_node.name == "add_generation_prompt": + references_agp = True + break + if references_agp and _if_body_emits_content(if_node): + return True + return False + + +# Sentinels for _derive_assistant_prefix_by_render. Diverge at char 0 so +# commonprefix can't absorb them; long random tail makes collision with real +# template literals negligible (see T18). +_RENDER_DIFF_SENTINEL_A = "AAAA_0123456789_UNSLOTH_RENDER_DIFF_SENTINEL" +_RENDER_DIFF_SENTINEL_B = "BBBB_0123456789_UNSLOTH_RENDER_DIFF_SENTINEL" +_RENDER_DIFF_SENTINEL_C = "CCCC_0123456789_UNSLOTH_RENDER_DIFF_SENTINEL" + + +def _derive_assistant_prefix_by_render(chat_template, is_sharegpt = False): + """Return the assistant-turn prefix the template emits, derived by + rendering two dialogs that differ only in assistant content: the common + prefix of their tails (after the base [user]-only render) is what the + template emits for an assistant turn. None if any guard fails. + + Works for Llama-3 / Gemma / Phi-3 and other non-ChatML shapes; the + template is its own ground truth. + + Known limitation: an `eos-on-non-last` pattern (turn-end sentinel only + emitted for non-last messages) would produce a consistent but wrong + prefix that `_validate_patched_template` can't catch. No real-world + template is known to use this. + """ + try: + from jinja2.sandbox import SandboxedEnvironment + except Exception: + return None + + if is_sharegpt: + base_msgs = [{"from": "human", "value": "Hi"}] + sent_a_msgs = base_msgs + [{"from": "gpt", "value": _RENDER_DIFF_SENTINEL_A}] + sent_b_msgs = base_msgs + [{"from": "gpt", "value": _RENDER_DIFF_SENTINEL_B}] + # User-role cross-check (Guard C below). + sent_c_msgs = base_msgs + [{"from": "human", "value": _RENDER_DIFF_SENTINEL_C}] else: - return endif + base_msgs = [{"role": "user", "content": "Hi"}] + sent_a_msgs = base_msgs + [ + {"role": "assistant", "content": _RENDER_DIFF_SENTINEL_A} + ] + sent_b_msgs = base_msgs + [ + {"role": "assistant", "content": _RENDER_DIFF_SENTINEL_B} + ] + sent_c_msgs = base_msgs + [{"role": "user", "content": _RENDER_DIFF_SENTINEL_C}] + + # Strip trailing whitespace/comments after the last endfor/endif: they + # appear after the message loop and would break Guard A. The splice in + # `_fix_chat_template` drops them too. + probe_template = chat_template + end = _find_end_position(chat_template) + if end is not None: + after = chat_template[end["end"] :] + if _RE_JINJA_COMMENT.sub("", after).strip() == "": + probe_template = chat_template[: end["end"]] + + # Sandboxed: probe renders at load time, before user calls + # apply_chat_template. SandboxedEnvironment blocks attribute-chain exploits. + try: + env = SandboxedEnvironment( + autoescape = False, + keep_trailing_newline = True, + ) + tmpl = env.from_string(probe_template) + out_base = tmpl.render(messages = base_msgs, add_generation_prompt = False) + out_a = tmpl.render(messages = sent_a_msgs, add_generation_prompt = False) + out_b = tmpl.render(messages = sent_b_msgs, add_generation_prompt = False) + except Exception: + return None + + # Best-effort: alternation-enforcing templates (e.g. Gemma's + # raise_exception) fail on [user, user]; that's a positive signal + # for Guard C, not a probe failure. + out_user_c = None + try: + out_user_c = tmpl.render(messages = sent_c_msgs, add_generation_prompt = False) + except Exception: + pass + + # Guard A: assistant renders extend base (no reordering). + if not (out_a.startswith(out_base) and out_b.startswith(out_base)): + return None + + tail_a = out_a[len(out_base) :] + tail_b = out_b[len(out_base) :] + if not tail_a or not tail_b: + return None + + prefix = os.path.commonprefix([tail_a, tail_b]) + + # Guard B: divergence is exactly at the content-insertion site. + if not ( + tail_a[len(prefix) :].startswith(_RENDER_DIFF_SENTINEL_A) + and tail_b[len(prefix) :].startswith(_RENDER_DIFF_SENTINEL_B) + ): + return None + + # Guard C: reject if a [user, user] render also emits the same prefix + # (role-insensitive template, e.g. `{% set greeting='Hi' %}...`). + if out_user_c is not None and out_user_c.startswith(out_base): + tail_c = out_user_c[len(out_base) :] + if tail_c.startswith(prefix) and prefix != "": + return None + + if not prefix: + return None + + return prefix -def _fix_chat_template(chat_template): - endfor = "{% endfor %}" - endif = "{% endif %}" - chosen_end = _find_end_position(chat_template, endfor, endif) - if chosen_end is None: - endfor = "{%- endfor %}" - endif = "{%- endif %}" - chosen_end = _find_end_position(chat_template, endfor, endif) - if chosen_end is None: +def _fix_chat_template(chat_template, is_sharegpt = False): + # Fast path: already has an {% if add_generation_prompt %} block, nothing + # to do. This catches cases the old string-based check would miss (e.g. + # templates that use {%- if add_generation_prompt -%} with both-side dash, + # or that sneak the block into a nested If/For). + if _has_add_generation_prompt_block(chat_template): return chat_template - where = chat_template.find(chosen_end) + end = _find_end_position(chat_template) + if end is None: + return chat_template - after_endfor = chat_template[where + len(chosen_end) :] - - dash = "-" if chosen_end.startswith("{%-") else "" + after_endfor = chat_template[end["end"] :] + dash_l = "-" if end["dash_left"] else "" + dash_r = "-" if end["dash_right"] else "" + open_tag = lambda body: "{%" + dash_l + " " + body + " " + dash_r + "%}" + # Case 1 (pre-existing base case): template ends with a single trailing + # {{ expr }} that is the generation prefix. Wrap it in an + # {% if add_generation_prompt %} ... {% endif %}. if ( - "{%" + dash + " if" not in after_endfor - and "{%" + dash + " set " not in after_endfor + "{%" + dash_l + " if" not in after_endfor + and "{%" + dash_l + " set " not in after_endfor and after_endfor.startswith("{{") and after_endfor.endswith("}}") and after_endfor.count("{{") == 1 and after_endfor.count("}}") == 1 ): - after_endfor = ( - "{%" + dash + " if add_generation_prompt %}" + after_endfor + endif + wrapped = ( + open_tag("if add_generation_prompt") + after_endfor + open_tag("endif") ) + return chat_template[: end["end"]] + wrapped - chat_template = chat_template[: where + len(chosen_end)] + after_endfor - - elif re.sub(r"\{#.*?#\}", "", after_endfor, flags = re.DOTALL).strip() == "": - # GH#4150: ChatML templates ending at {% endfor %} without an - # add_generation_prompt block. Scrub Jinja `{# ... #}` comments so - # tokens inside comments cannot fool the guard below. - scrubbed = re.sub(r"\{#.*?#\}", "", chat_template, flags = re.DOTALL) - if ( - "<|im_start|>" in scrubbed - and "<|im_end|>" in scrubbed - and "add_generation_prompt" not in scrubbed - ): - # Infer the assistant-turn separator. Prefer an explicit - # '<|im_start|>assistant' literal; else the unique - # `message['role'] + ''` from role concatenations; else - # '<|im_sep|>' if present (Phi-4-mini uses '\n' for system and - # '<|im_sep|>' for user/assistant); else '\n'. - assistant_match = re.search( - r"""(['"])<\|im_start\|>assistant([^'"]*)\1""", - scrubbed, + # Case 2 (GH#4150): template ends at {% endfor %} with only whitespace + # or comments left. Inject an {% if add_generation_prompt %} block with + # the assistant prefix derived by render-diff. The top-level-For gate + # keeps us out of outer-If wrappers (e.g. Qwen3-Guard). + if _RE_JINJA_COMMENT.sub( + "", after_endfor + ).strip() == "" and _template_ends_with_toplevel_for(chat_template): + # No redundant "agp not in scrubbed" check: the fast path already + # confirmed no *positive* block, and a mere reference (header + # guard) should still get repaired. + assistant_prefix = _derive_assistant_prefix_by_render( + chat_template, is_sharegpt + ) + # Dual-probe: dict/list callers don't know the shape up front. + if assistant_prefix is None and not is_sharegpt: + assistant_prefix = _derive_assistant_prefix_by_render( + chat_template, is_sharegpt = True ) - role_seps = [ - m.group(2) - for m in re.finditer( - r"""message(?:\[['"]role['"]\]|\.role)\s*\+\s*(['"])([^'"]*)\1""", - scrubbed, - ) - ] - unique_role_seps = list(dict.fromkeys(role_seps)) - if assistant_match is not None and assistant_match.group(2): - separator = assistant_match.group(2) - elif len(unique_role_seps) == 1: - separator = unique_role_seps[0] - elif "<|im_sep|>" in scrubbed: - separator = "<|im_sep|>" - else: - separator = "\\n" - # Emit a double-quoted Jinja literal so a single quote in the - # separator cannot break the block. Drop trailing whitespace/ - # comments after endfor: they would render as stray output - # after the generation prefix. - assistant_prefix = "<|im_start|>assistant" + separator - generation_block = ( - "{%" + dash + " if add_generation_prompt %}" - '{{ "' + assistant_prefix.replace('"', '\\"') + '" }}' - "{%" + dash + " endif %}" - ) - chat_template = chat_template[: where + len(chosen_end)] + generation_block + if assistant_prefix is None: + return chat_template + # Escape for a double-quoted Jinja string literal. + escaped = ( + assistant_prefix.replace("\\", "\\\\") + .replace('"', '\\"') + .replace("\n", "\\n") + .replace("\r", "\\r") + ) + generation_block = ( + open_tag("if add_generation_prompt") + + '{{ "' + + escaped + + '" }}' + + open_tag("endif") + ) + return chat_template[: end["end"]] + generation_block return chat_template +def _is_strict_chat_template_mode(): + """Opt-in strict mode restores the pre-warn RuntimeError behavior.""" + val = os.environ.get("UNSLOTH_STRICT_CHAT_TEMPLATE", "0") + return str(val).strip().lower() in ("1", "true", "yes", "on") + + +def _name_is_local_path(name_or_path): + """True if name_or_path refers to an existing local directory. Used to + tailor the warning message: for local paths the user cannot 'file a bug + report to the maintainers of ' since that path is their own.""" + if not name_or_path: + return False + try: + return os.path.isdir(str(name_or_path)) + except Exception: + return False + + +def _format_chat_template_message( + name_or_path, + repaired, + has_generation_block = False, + local_path_source = None, + strict = False, +): + """Build a user-facing warning/error message that points at the right + responsible party (user's downstream tool vs. upstream model maintainer).""" + local = _name_is_local_path( + local_path_source if local_path_source is not None else name_or_path + ) + if local: + source_hint = ( + "This tokenizer was loaded from a local path. The likely cause is a " + "downstream tool (LlamaFactory, Axolotl, etc.) that re-serialized " + "the tokenizer during save and stripped the generation-prompt " + "block. Either re-save with the original template, or set " + "`tokenizer.chat_template` manually before loading." + ) + else: + source_hint = ( + "The chat_template shipped with `{name}` appears incomplete. " + "Consider filing a bug report with the model maintainers." + ).format(name = name_or_path) + strict_suffix = ( + "" + if strict + else (" Set UNSLOTH_STRICT_CHAT_TEMPLATE=1 to raise instead of warn.") + ) + if repaired: + return ( + "Unsloth: Patched the chat_template on `{name}` to add a " + "{{% if add_generation_prompt %}} block. {hint}" + ).format(name = name_or_path, hint = source_hint) + if has_generation_block: + return ( + "Unsloth: The tokenizer `{name}` has a " + "{{% if add_generation_prompt %}} block, but it does not change " + "the rendered output. {hint}{suffix}" + ).format(name = name_or_path, hint = source_hint, suffix = strict_suffix) + load_clause = ( + "Loading is blocked in strict mode." + if strict + else "The model will still load, but " + "`apply_chat_template(add_generation_prompt=True)` may not produce a " + "correct assistant-turn marker." + ) + return ( + "Unsloth: The tokenizer `{name}` does not have a " + "{{% if add_generation_prompt %}} block for generation purposes, and " + "automatic repair was not possible. {load_clause} {hint}{suffix}" + ).format( + name = name_or_path, + load_clause = load_clause, + hint = source_hint, + suffix = strict_suffix, + ) + + +def _validate_patched_template(tokenizer, patched_template, is_sharegpt): + """Render the just-patched template with and without + add_generation_prompt, and confirm the patched output responds to the + flag by appending (not replacing) content. Returns True if validation + passes.""" + msgs = ( + [{"from": "human", "value": "Hi"}] + if is_sharegpt + else [{"role": "user", "content": "Hi"}] + ) + original = getattr(tokenizer, "chat_template", None) + try: + try: + tokenizer.chat_template = patched_template + except Exception: + return False # read-only tokenizer, skip validation + try: + yes = tokenizer.apply_chat_template( + msgs, + add_generation_prompt = True, + tokenize = False, + ) + no = tokenizer.apply_chat_template( + msgs, + add_generation_prompt = False, + tokenize = False, + ) + except Exception: + return False + finally: + try: + tokenizer.chat_template = original + except Exception: + pass # best-effort restore + # Contract after a successful repair: the two renders differ, and the + # "yes" render is a strict extension of the "no" render (we only + # appended content inside the new add_generation_prompt block). + return yes != no and yes.startswith(no) + + +def _repair_string_template(tokenizer, chat_template, is_sharegpt): + """Core string-template repair. Returns the repaired template on success, + or None if repair was not possible / failed validation.""" + candidate = _fix_chat_template(chat_template, is_sharegpt = is_sharegpt) + if not _has_add_generation_prompt_block(candidate): + return None + # Validate with the caller's is_sharegpt first. If that fails, the + # dual-probe in _fix_chat_template may have fallen back to the other + # schema internally -- try validating with the opposite schema before + # giving up. + if _validate_patched_template(tokenizer, candidate, is_sharegpt): + return candidate + if _validate_patched_template(tokenizer, candidate, not is_sharegpt): + return candidate + return None + + +def _fix_chat_template_for_tokenizer(tokenizer, chat_template): + """Entry point for a string chat_template. Runs the no==yes diagnostic, + attempts repair if needed, and returns the (possibly patched) template. + + On repair failure, the behavior is controlled by + UNSLOTH_STRICT_CHAT_TEMPLATE: warn + return original (default) or raise + RuntimeError (strict).""" + name = getattr(tokenizer, "name_or_path", "unknown") + source_path = getattr(tokenizer, "_source_path", name) + + # Detect ShareGPT vs HF style by probing apply_chat_template. + is_sharegpt = None + try: + tokenizer.apply_chat_template( + [{"role": "user", "content": "Who are you?"}], + add_generation_prompt = False, + tokenize = False, + ) + is_sharegpt = False + except Exception: + try: + tokenizer.apply_chat_template( + [{"from": "human", "value": "Who are you?"}], + add_generation_prompt = False, + tokenize = False, + ) + is_sharegpt = True + except Exception: + is_sharegpt = None + + if is_sharegpt is None: + return chat_template + + messages = ( + [{"from": "human", "value": "Who are you?"}] + if is_sharegpt + else [{"role": "user", "content": "Who are you?"}] + ) + try: + no = tokenizer.apply_chat_template( + messages, + add_generation_prompt = False, + tokenize = False, + ) + yes = tokenizer.apply_chat_template( + messages, + add_generation_prompt = True, + tokenize = False, + ) + except Exception: + return chat_template + + if no != yes: + # Template already responds to the flag; leave as is. + return chat_template + + # no == yes: template ignores add_generation_prompt. Try to repair. + if _has_add_generation_prompt_block(chat_template): + # Template has the block but it does not change output. This is the + # "wasn't provided correctly" case from the pre-warn code path. + strict = _is_strict_chat_template_mode() + msg = _format_chat_template_message( + name, + repaired = False, + has_generation_block = True, + local_path_source = source_path, + strict = strict, + ) + if strict: + raise RuntimeError(msg) + logger.warning_once(msg) + return chat_template + + repaired = _repair_string_template(tokenizer, chat_template, is_sharegpt) + if repaired is not None: + logger.warning_once( + _format_chat_template_message( + name, + repaired = True, + local_path_source = source_path, + ) + ) + return repaired + + strict = _is_strict_chat_template_mode() + msg = _format_chat_template_message( + name, + repaired = False, + local_path_source = source_path, + strict = strict, + ) + if strict: + raise RuntimeError(msg) + logger.warning_once(msg) + return chat_template + + +class _VariantTokenizerProxy: + """Single-variant view of a multi-variant tokenizer. Routes each variant + through `_fix_chat_template_for_tokenizer` so the full contract + (is_sharegpt probe, no==yes, warn/strict, `_validate_patched_template`) + applies instead of jumping straight to structural repair. + + `apply_chat_template` swaps `base.chat_template` to the variant before + calling so tokenizer globals (bos_token, filters, raise_exception) are + preserved; falls back to bare Jinja for read-only stubs. + """ + + def __init__(self, base_tokenizer, variant_template, variant_label = ""): + self._base = base_tokenizer + self._template = variant_template + base_name = getattr(base_tokenizer, "name_or_path", "unknown") + self._source_path = base_name + self.name_or_path = ( + f"{base_name} ({variant_label})" if variant_label else base_name + ) + + @property + def chat_template(self): + return self._template + + @chat_template.setter + def chat_template(self, value): + self._template = value + + def apply_chat_template(self, *args, **kwargs): + base_original = getattr(self._base, "chat_template", None) + swapped = False + try: + try: + self._base.chat_template = self._template + swapped = True + except Exception: + swapped = False + if swapped: + return self._base.apply_chat_template(*args, **kwargs) + # Read-only base: fall back to sandboxed Jinja. + from jinja2.sandbox import SandboxedEnvironment + + env = SandboxedEnvironment( + autoescape = False, + keep_trailing_newline = True, + ) + messages = args[0] if args else kwargs.get("messages", []) + add_generation_prompt = kwargs.get("add_generation_prompt", False) + return env.from_string(self._template).render( + messages = messages, + add_generation_prompt = add_generation_prompt, + ) + finally: + if swapped: + try: + self._base.chat_template = base_original + except Exception: + pass # best-effort restore + + def fix_chat_template(tokenizer): chat_template = getattr(tokenizer, "chat_template", None) if chat_template is None: return None - ### 1. Check if add_generation_prompt works - # Check for ShareGPT style first - is_sharegpt = None - try: - messages = [ - {"role": "user", "content": "Who are you?"}, - ] - tokenizer.apply_chat_template( - messages, add_generation_prompt = False, tokenize = False - ) - is_sharegpt = False - except: - try: - messages = [ - {"from": "human", "value": "Who are you?"}, - ] - tokenizer.apply_chat_template( - messages, add_generation_prompt = False, tokenize = False + # Multi-variant dict (e.g. Hermes-3 {default, tool_use}): route each + # variant through the full repair contract via _VariantTokenizerProxy. + if isinstance(chat_template, dict): + fixed = {} + for key, tmpl in chat_template.items(): + if not isinstance(tmpl, str): + fixed[key] = tmpl + continue + proxy = _VariantTokenizerProxy( + tokenizer, tmpl, variant_label = f"variant={key!r}" ) - is_sharegpt = True - except: - is_sharegpt = None + fixed[key] = _fix_chat_template_for_tokenizer(proxy, tmpl) + return fixed - # Not ShareGPT or HF style - just return - if is_sharegpt is None: - return chat_template - - # Tokenize - messages = [ - {"role": "user", "content": "Who are you?"} - if not is_sharegpt - else {"from": "human", "value": "Who are you?"} - ] - no = tokenizer.apply_chat_template( - messages, add_generation_prompt = False, tokenize = False - ) - yes = tokenizer.apply_chat_template( - messages, add_generation_prompt = True, tokenize = False - ) - - if no == yes: - # SAME?! That's not good! We check for add_generation_prompt - if ( - "{% if add_generation_prompt %}" not in chat_template - and "{%- if add_generation_prompt %}" not in chat_template - ): - # Try fixing it by adding it - new_chat_template = _fix_chat_template(chat_template) - if ( - "{% if add_generation_prompt %}" not in new_chat_template - and "{%- if add_generation_prompt %}" not in new_chat_template - ): - raise RuntimeError( - f"Unsloth: The tokenizer `{tokenizer.name_or_path}`\n" - "does not have a {% if add_generation_prompt %} for generation purposes.\n" - f"Please file a bug report to the maintainers of `{tokenizer.name_or_path}` - thanks!" - ) + # List-of-dicts form (older HF multi-template style). + if isinstance(chat_template, list): + fixed = [] + for item in chat_template: + if not isinstance(item, dict) or "template" not in item: + fixed.append(item) + continue + tmpl = item["template"] + if not isinstance(tmpl, str): + fixed.append(item) + continue + label = f"variant={item.get('name', '?')!r}" + proxy = _VariantTokenizerProxy(tokenizer, tmpl, variant_label = label) + new_tmpl = _fix_chat_template_for_tokenizer(proxy, tmpl) + if new_tmpl is tmpl or new_tmpl == tmpl: + fixed.append(item) else: - logger.warning_once( - "Unsloth: We successfully patched the tokenizer to add a {% if add_generation_prompt %} to the chat_template.\n" - f"This is not a bug, but please notify the maintainers of `{tokenizer.name_or_path}` - thanks!" - ) - chat_template = new_chat_template - else: - raise RuntimeError( - f"Unsloth: The tokenizer `{tokenizer.name_or_path}`\n" - "has a {% if add_generation_prompt %} for generation purposes, but wasn't provided correctly.\n" - "Please file a bug report immediately - thanks!" - ) - return chat_template + fixed.append({**item, "template": new_tmpl}) + return fixed + + return _fix_chat_template_for_tokenizer(tokenizer, chat_template) def check_tokenizer(