Merge branch 'main' into fix/bnb-multidevice-inference-hooks
This commit is contained in:
Daniel Han
GitHub
commit
705ef146da
14 changed files with 1637 additions and 336 deletions
183
install.ps1
183
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.
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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 ``<ollama_dir>/.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::
|
||||
|
||||
<ollama_dir>/manifests/<host>/<namespace>/<model>/<tag>
|
||||
<ollama_dir>/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
|
||||
``<ollama_dir>/.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.
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="px-2.5 py-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
{children}
|
||||
<div className="flex items-center justify-between gap-1 px-2.5 py-1.5">
|
||||
<span className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
{icon}
|
||||
{children}
|
||||
</span>
|
||||
{onToggle && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={onToggle}
|
||||
aria-label={collapsed ? "Expand section" : "Collapse section"}
|
||||
className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
|
||||
>
|
||||
{collapsed
|
||||
? <ChevronRightIcon className="size-3" />
|
||||
: <ChevronDownIcon className="size-3" />}
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -489,6 +515,9 @@ export function HubModelPicker({
|
|||
// Delete confirmation dialog state
|
||||
const [deleteTarget, setDeleteTarget] = useState<string | null>(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<string[]>([]);
|
||||
|
||||
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)) ? (
|
||||
<>
|
||||
<ListLabel>Downloaded</ListLabel>
|
||||
{cachedGguf.map((c) => (
|
||||
<ListLabel
|
||||
icon={<DownloadIcon className="size-3" />}
|
||||
collapsed={downloadedCollapsed}
|
||||
onToggle={() => setDownloadedCollapsed((v) => !v)}
|
||||
>Downloaded</ListLabel>
|
||||
{!downloadedCollapsed && cachedGguf.map((c) => (
|
||||
<div key={c.repo_id}>
|
||||
<ModelRow
|
||||
label={c.repo_id}
|
||||
|
|
@ -922,7 +959,7 @@ export function HubModelPicker({
|
|||
)}
|
||||
</div>
|
||||
))}
|
||||
{!chatOnly &&
|
||||
{!downloadedCollapsed && !chatOnly &&
|
||||
cachedModels.map((c) => (
|
||||
<div key={c.repo_id} className="flex items-center gap-0.5">
|
||||
<div className="min-w-0 flex-1">
|
||||
|
|
@ -1001,20 +1038,12 @@ export function HubModelPicker({
|
|||
|
||||
{!showHfSection ? (
|
||||
<>
|
||||
<div className="flex items-center justify-between gap-1 px-2.5 py-1.5">
|
||||
<span className="text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
<div className="flex items-center gap-1 px-2.5 py-1.5">
|
||||
<span className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
<HugeiconsIcon icon={Folder02Icon} className="size-3" />
|
||||
Custom Folders
|
||||
</span>
|
||||
<div className="flex items-center gap-0.5">
|
||||
<button
|
||||
type="button"
|
||||
aria-label="Browse for a folder on the server"
|
||||
title="Browse folders on the server"
|
||||
onClick={() => setShowFolderBrowser(true)}
|
||||
className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon icon={Search01Icon} className="size-3" />
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
aria-label={showFolderInput ? "Cancel adding folder" : "Add scan folder by path"}
|
||||
|
|
@ -1029,11 +1058,33 @@ export function HubModelPicker({
|
|||
>
|
||||
<HugeiconsIcon icon={showFolderInput ? Cancel01Icon : Add01Icon} className="size-3" />
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
aria-label="Browse for a folder on the server"
|
||||
title="Browse folders on the server"
|
||||
onClick={() => setShowFolderBrowser(true)}
|
||||
className="shrink-0 rounded p-0.5 text-muted-foreground/60 transition-colors hover:text-foreground"
|
||||
>
|
||||
<HugeiconsIcon icon={Search01Icon} className="size-2.5" />
|
||||
</button>
|
||||
</div>
|
||||
<div className="ml-auto">
|
||||
<button
|
||||
type="button"
|
||||
aria-label={customFoldersCollapsed ? "Expand custom folders" : "Collapse custom folders"}
|
||||
title={customFoldersCollapsed ? "Expand" : "Collapse"}
|
||||
onClick={() => setCustomFoldersCollapsed((v) => !v)}
|
||||
className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
|
||||
>
|
||||
{customFoldersCollapsed
|
||||
? <ChevronRightIcon className="size-3" />
|
||||
: <ChevronDownIcon className="size-3" />}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Folder paths */}
|
||||
{scanFolders.map((f) => (
|
||||
{!customFoldersCollapsed && scanFolders.map((f) => (
|
||||
<div
|
||||
key={f.id}
|
||||
className="group flex items-center gap-1.5 px-2.5 py-0.5"
|
||||
|
|
@ -1056,8 +1107,31 @@ export function HubModelPicker({
|
|||
</div>
|
||||
))}
|
||||
|
||||
{/* 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 (
|
||||
<div className="flex flex-wrap gap-1 px-2.5 pb-0.5">
|
||||
{unregistered.map((p) => (
|
||||
<button
|
||||
key={p}
|
||||
type="button"
|
||||
onClick={() => void handleAddFolder(p)}
|
||||
disabled={folderLoading}
|
||||
title={`Add ${p}`}
|
||||
className="rounded-full border border-dashed border-border/50 px-2 py-0.5 font-mono text-[10px] text-muted-foreground/70 transition-colors hover:border-foreground/30 hover:bg-accent hover:text-foreground disabled:opacity-40"
|
||||
>
|
||||
<span className="text-[11px] font-semibold">+</span> {p.length > 30 ? `...${p.slice(-27)}` : p}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
);
|
||||
})()}
|
||||
|
||||
{/* Add folder input */}
|
||||
{showFolderInput && (
|
||||
{!customFoldersCollapsed && showFolderInput && (
|
||||
<div className="px-2.5 pb-1 pt-0.5">
|
||||
<div className="flex items-center gap-1">
|
||||
<HugeiconsIcon icon={Folder02Icon} className="size-3 shrink-0 text-muted-foreground/40" />
|
||||
|
|
@ -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 (
|
||||
<div key={m.id}>
|
||||
<ModelRow
|
||||
|
|
@ -1126,7 +1204,13 @@ export function HubModelPicker({
|
|||
meta={isGguf ? "GGUF" : "Local"}
|
||||
selected={value === m.id}
|
||||
onClick={() => {
|
||||
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 ? (
|
||||
<>
|
||||
<ListLabel>Recommended</ListLabel>
|
||||
{visibleRecommendedIds.length === 0 ? (
|
||||
<ListLabel
|
||||
icon={<StarIcon className="size-3" />}
|
||||
collapsed={recommendedCollapsed}
|
||||
onToggle={() => setRecommendedCollapsed((v) => !v)}
|
||||
>Recommended</ListLabel>
|
||||
{recommendedCollapsed ? null : visibleRecommendedIds.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No default models.
|
||||
</div>
|
||||
|
|
@ -1203,7 +1291,7 @@ export function HubModelPicker({
|
|||
);
|
||||
})
|
||||
)}
|
||||
{hasMoreRecommended && (
|
||||
{!recommendedCollapsed && hasMoreRecommended && (
|
||||
<>
|
||||
<div ref={recommendedSentinelRef} className="h-px" />
|
||||
<div className="flex items-center justify-center py-2">
|
||||
|
|
@ -1216,7 +1304,7 @@ export function HubModelPicker({
|
|||
|
||||
{showHfSection && filteredRecommendedIds.length > 0 ? (
|
||||
<>
|
||||
<ListLabel>Recommended</ListLabel>
|
||||
<ListLabel icon={<StarIcon className="size-3" />}>Recommended</ListLabel>
|
||||
{filteredRecommendedIds.map((id) => {
|
||||
const vram = recommendedVramMap.get(id);
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -262,6 +262,12 @@ export interface BrowseFoldersResponse {
|
|||
model_files_here?: number;
|
||||
}
|
||||
|
||||
export async function listRecommendedFolders(): Promise<string[]> {
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
158
studio/setup.ps1
158
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 })) {
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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<sep>' literal; else the unique
|
||||
# `message['role'] + '<sep>'` 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 <path>' 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(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue