* fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) When a flat local GGUF directory holds several unrelated models with their own mmproj siblings, detect_mmproj_file() returned the first projector it walked into. For the layout reported in #5347 (Qwen weights + a Gemma mmproj in the same dir) that meant llama-server was launched with --mmproj pointing at the Gemma projector, which fails to load and surfaces as a confusing crash. Disambiguation rules: - Drop candidates whose family token (qwen/gemma/llama/mistral/phi/...) disagrees with the model's family. Candidates with no recognised family token (e.g. the HF-convention 'mmproj-F16.gguf') are kept. - Among same-family candidates, prefer the one whose stem shares the longest prefix with the model (Qwen3.5-9B mmproj beats Qwen3.5-35B mmproj for a Qwen3.5-9B model). - If every candidate is dropped, return None — better than attaching a wrong projector and getting a server-launch failure. Tests cover the cross-family block, multi-candidate prefix tie-break, HF-convention 'mmproj-F16.gguf', unrecognised families, and the existing search_root walk. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: word-bounded family match, expanded token list, launcher guard Tighten the family-token detector to match only on word boundaries so substring collisions stop tagging false families: phi no longer matches sapphire, yi no longer matches yip, mimo no longer matches mimosa, and mistral does not bleed into ministral/magistral/devstral. Pick the token whose first occurrence is leftmost in the filename rather than the first hit in tuple order, so merge models disambiguate predictably (llama-phi tags llama; phi-llama tags phi). Expand _MODEL_FAMILY_TOKENS with the families an audit of the unsloth HF org turned up that the previous list missed: devstral, ministral, magistral (Mistral-derivative naming), nemotron, kimi, nanonets, cosmos, mimo, apriel, lfm. Without these, a flat local GGUF directory containing one of these weights plus an unrelated renamed projector still hit the original #5347 failure. Add mmproj_matches_model_family() and call it at the llama-server launch site in core/inference/llama_cpp.py. detect_mmproj_file already drops cross-family candidates at discovery time, but mmproj_path can also reach the launcher via config injection or future overrides; this guard keeps those paths from silently loading a known-wrong projector. Tests: 12 new cases covering substring rejection, leftmost-position selection, new family tokens, a new flat-dir Nemotron + Gemma rejection case, and the launcher-level guard. All 21 detect_mmproj_file tests and the existing 106 llama_cpp tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: pair via GGUF general.* metadata, not just filenames Real Unsloth vision GGUFs carry rich identity metadata that has been ignored by the discovery path. Every projector under the unsloth org has general.type='mmproj' plus general.base_model.0.repo_url pointing at the same upstream HF repo as its weight, and the equivalent basename, base_model.0.name, and base_model.0.organization fields. A flat-dir mismatch is therefore decidable from the headers alone, no matter how the user has renamed the files. Add utils/models/gguf_metadata.py with read_gguf_general_metadata(): a fast (~30 ms) header walk that pulls only the general.* string fields and skips everything else, cached by (resolved path, mtime_ns, size). Mirrors the parser shape already used by LlamaCppBackend._read_gguf_metadata so the format handling is consistent. is_mmproj_by_metadata() returns True/False/None from general.type, and pairing_score() returns 100 for an exact base_model URL match, 80 for basename plus organization match, 60 for basename only, -1 for definitive metadata disagreement, and 0 when neither side has enough metadata to decide. Rewire detect_mmproj_file() to a two-stage selector: 1. Detect projectors via metadata (general.type) when present, else fall back to the filename substring heuristic. This recovers headerless projectors AND projectors whose name does not contain 'mmproj' but whose header advertises one. 2. Score each candidate against the weight via pairing_score. Drop candidates with score -1 (definitive metadata disagreement). For candidates with score 0 (no usable metadata) fall back to the existing filename family-token check, dropping recognised-family mismatches. Pick the survivor with the highest (score, longest_prefix, -len(stem)) tuple, so a metadata URL match always wins over a filename-prefix match. Tests: 16 new cases. tests/test_gguf_metadata.py covers the parser (missing file, non-GGUF, string extraction, walking past arrays and uint32s, cache invalidation by mtime/size) and the score helpers. tests/test_detect_mmproj_file.py adds end-to-end cases that synthesise real on-disk GGUF headers: URL match wins over a longer-prefix sibling, URL mismatch returns None even when filenames match, a projector named 'vision-projector.gguf' is still discovered via general.type, and a 100-score header match outranks a near-perfect filename prefix on a headerless candidate. All 75 tests across detect_mmproj_file, gguf_metadata, llama_cpp load progress, cached gguf routes, trained model scan, and vision cache pass. * studio/mmproj: shorten comments and docstrings across the #5347 changes Trim verbose explanations to one-line statements of intent. The behaviour is unchanged: 161 tests across detect_mmproj_file, gguf_metadata, llama_cpp_load_progress (+ matrix), llama_server_args, llama_cpp_cache_aware_disk_check, trained_model_scan, and vision_cache all pass. * studio/mmproj: shorten remaining detect_mmproj_file body comments Trim the docstring and the dir-walking block comments inside detect_mmproj_file to one-liners. Behaviour unchanged; 44 mmproj + gguf_metadata + llama_cpp_load_progress tests pass. * studio/mmproj: cap gguf_metadata cache below ceiling on every insert The eviction branch popped exactly one entry when len >= max, so the cache size could only converge to the cap when entries were added slowly enough for natural growth. After a sandbox sim that reduced the cap mid-run, len stayed above the cap because each insert popped one and added one. Switch to a while loop so we evict until len is strictly below the cap before inserting. Steady-state behaviour at the default 4096 ceiling is unchanged. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Tai An
GitHub
parent
cb15a7a5b6
commit
63c6750532
5 changed files with 949 additions and 42 deletions
|
|
@ -2367,8 +2367,20 @@ class LlamaCppBackend:
|
|||
if not Path(mmproj_path).is_file():
|
||||
logger.warning(f"mmproj file not found: {mmproj_path}")
|
||||
else:
|
||||
cmd.extend(["--mmproj", mmproj_path])
|
||||
logger.info(f"Using mmproj for vision: {mmproj_path}")
|
||||
# #5347 guard for paths that bypass detect_mmproj_file.
|
||||
from utils.models.model_config import (
|
||||
mmproj_matches_model_family,
|
||||
)
|
||||
|
||||
if not mmproj_matches_model_family(model_path, mmproj_path):
|
||||
logger.warning(
|
||||
f"Skipping mmproj with mismatched family: "
|
||||
f"model={Path(model_path).name}, "
|
||||
f"mmproj={Path(mmproj_path).name}"
|
||||
)
|
||||
else:
|
||||
cmd.extend(["--mmproj", mmproj_path])
|
||||
logger.info(f"Using mmproj for vision: {mmproj_path}")
|
||||
|
||||
# Option C: add --api-key for direct client access when enabled
|
||||
import os as _os
|
||||
|
|
@ -3747,7 +3759,7 @@ class LlamaCppBackend:
|
|||
|
||||
except json.JSONDecodeError:
|
||||
logger.debug(
|
||||
f"Skipping malformed SSE line: " f"{line[:100]}"
|
||||
f"Skipping malformed SSE line: {line[:100]}"
|
||||
)
|
||||
if _stream_done:
|
||||
break # exit outer for
|
||||
|
|
|
|||
326
studio/backend/tests/test_detect_mmproj_file.py
Normal file
326
studio/backend/tests/test_detect_mmproj_file.py
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for :func:`utils.models.model_config.detect_mmproj_file` (#5347)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import struct
|
||||
|
||||
from utils.models.model_config import (
|
||||
_detect_family_token,
|
||||
detect_mmproj_file,
|
||||
mmproj_matches_model_family,
|
||||
)
|
||||
|
||||
|
||||
_GGUF_MAGIC = 0x46554747
|
||||
|
||||
|
||||
def _gguf_with_general(path: Path, fields: dict) -> Path:
|
||||
"""Write a minimal GGUF with only ``general.*`` string KVs."""
|
||||
body = b""
|
||||
for k, v in fields.items():
|
||||
kb = k.encode("utf-8")
|
||||
vb = v.encode("utf-8")
|
||||
body += struct.pack("<Q", len(kb)) + kb
|
||||
body += struct.pack("<I", 8) # STRING vtype
|
||||
body += struct.pack("<Q", len(vb)) + vb
|
||||
header = struct.pack("<IIQQ", _GGUF_MAGIC, 3, 0, len(fields))
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(header + body)
|
||||
return path
|
||||
|
||||
|
||||
def _touch(path: Path) -> Path:
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(b"")
|
||||
return path
|
||||
|
||||
|
||||
def test_returns_none_when_no_mmproj(tmp_path: Path):
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_single_matching_family_mmproj_picked(tmp_path: Path):
|
||||
"""Single same-family projector: returned (historical behaviour)."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
mmproj = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_hf_style_unprefixed_mmproj_still_works(tmp_path: Path):
|
||||
"""HF convention: weight + ``mmproj-F16.gguf`` sibling."""
|
||||
model = _touch(tmp_path / "model.gguf")
|
||||
mmproj = _touch(tmp_path / "mmproj-F16.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_blocks_single_cross_family_projector(tmp_path: Path):
|
||||
"""#5347 core: Qwen weight + lone Gemma mmproj returns None."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_picks_matching_family_among_mixed_candidates(tmp_path: Path):
|
||||
"""Mixed Qwen + Gemma projectors: pick Qwen, drop Gemma."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
qwen_mm = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(qwen_mm.resolve())
|
||||
|
||||
|
||||
def test_prefers_longest_prefix_within_same_family(tmp_path: Path):
|
||||
"""Same family, different sizes: longest shared stem prefix wins."""
|
||||
model = _touch(tmp_path / "Qwen3.5-35B-A3B-UD-Q4_K_L.gguf")
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
big_mm = _touch(tmp_path / "Qwen3.5-35B-A3B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(big_mm.resolve())
|
||||
|
||||
|
||||
def test_unrecognised_family_does_not_break_detection(tmp_path: Path):
|
||||
"""Unknown model family must not return None on a sole candidate."""
|
||||
model = _touch(tmp_path / "MyCustomBrand-7B-Q4_K_M.gguf")
|
||||
mmproj = _touch(tmp_path / "MyCustomBrand-7B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_directory_path_returns_first_candidate(tmp_path: Path):
|
||||
"""Directory path: no model stem to compare; legacy first-candidate."""
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
result = detect_mmproj_file(str(tmp_path))
|
||||
assert result is not None
|
||||
assert "mmproj" in Path(result).name.lower()
|
||||
|
||||
|
||||
def test_search_root_walk_still_works(tmp_path: Path):
|
||||
"""Snapshot layout: weight in quant subdir, mmproj at snapshot root."""
|
||||
snapshot = tmp_path / "snapshot"
|
||||
weight = _touch(snapshot / "BF16" / "Qwen3.5-9B-BF16.gguf")
|
||||
mmproj = _touch(snapshot / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
result = detect_mmproj_file(str(weight), search_root = str(snapshot))
|
||||
assert result == str(mmproj.resolve())
|
||||
|
||||
|
||||
# -- Family token detection: word-bounded matching ----------------------
|
||||
|
||||
|
||||
def test_family_token_phi_does_not_match_sapphire():
|
||||
"""``phi`` substring inside ``sapphire`` must not tag Phi."""
|
||||
assert _detect_family_token("sapphire-7b-q4_k_m.gguf") is None
|
||||
|
||||
|
||||
def test_family_token_yi_does_not_match_tinyish_names():
|
||||
"""``yi`` must not cross letter boundaries (``yip``)."""
|
||||
assert _detect_family_token("yip-7b.gguf") is None
|
||||
assert _detect_family_token("yi-vl-6b.gguf") == "yi"
|
||||
|
||||
|
||||
def test_family_token_mimo_does_not_match_mimosa():
|
||||
"""``mimo`` must not tag ``mimosa``."""
|
||||
assert _detect_family_token("mimosa-rosa-7b.gguf") is None
|
||||
assert _detect_family_token("MiMo-VL-7B-RL-BF16.gguf") == "mimo"
|
||||
|
||||
|
||||
def test_family_token_mistral_does_not_match_ministral():
|
||||
"""Pin Mistral-derivative tagging."""
|
||||
assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral"
|
||||
assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral"
|
||||
assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral"
|
||||
assert (
|
||||
_detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
|
||||
== "devstral"
|
||||
)
|
||||
|
||||
|
||||
def test_family_token_picks_leftmost_when_multiple_present():
|
||||
"""Leftmost family token wins, not tuple order."""
|
||||
assert _detect_family_token("llama-phi-merge.gguf") == "llama"
|
||||
assert _detect_family_token("phi-llama-merge.gguf") == "phi"
|
||||
assert _detect_family_token("llama3-3b-instruct.gguf") == "llama"
|
||||
|
||||
|
||||
def test_family_token_new_families_recognised():
|
||||
"""Catalogue-audit additions tag correctly."""
|
||||
assert _detect_family_token("NVIDIA-Nemotron-3-Nano-Omni-30B.gguf") == "nemotron"
|
||||
assert _detect_family_token("Kimi-K2.6-BF16.gguf") == "kimi"
|
||||
assert _detect_family_token("Nanonets-OCR-s-BF16.gguf") == "nanonets"
|
||||
assert _detect_family_token("Cosmos-Reason1-7B-BF16.gguf") == "cosmos"
|
||||
assert _detect_family_token("Apriel-1.5-15b-Thinker-BF16.gguf") == "apriel"
|
||||
assert _detect_family_token("LFM2.5-VL-1.6B-BF16.gguf") == "lfm"
|
||||
|
||||
|
||||
# -- Cross-family rejection with the expanded token list ----------------
|
||||
|
||||
|
||||
def test_blocks_cross_family_for_new_token_pair(tmp_path: Path):
|
||||
"""Nemotron weight + lone Gemma projector returns None."""
|
||||
model = _touch(
|
||||
tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf"
|
||||
)
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_picks_devstral_mmproj_in_mixed_dir(tmp_path: Path):
|
||||
"""Devstral weight + Devstral mmproj + a Qwen mmproj: pick Devstral."""
|
||||
model = _touch(tmp_path / "Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
|
||||
dev_mm = _touch(tmp_path / "Devstral-Small-2-mmproj-bf16.gguf")
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(dev_mm.resolve())
|
||||
|
||||
|
||||
# -- Launcher-level family guard ----------------------------------------
|
||||
|
||||
|
||||
def test_mmproj_family_guard_blocks_cross_family():
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/gemma-4-26B-A4B-it.mmproj-q8_0.gguf",
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_same_family():
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/Qwen3.5-9B-BF16-mmproj.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_generic_hf_mmproj():
|
||||
"""No family token on the projector: wildcard."""
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/mmproj-F16.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_unrecognised_model_family():
|
||||
"""No family token on the model: wildcard."""
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Apriel-1.5-15b-Thinker-BF16.gguf",
|
||||
"/models/mmproj-F16.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# -- Metadata-primary pairing in detect_mmproj_file ---------------------
|
||||
|
||||
|
||||
def test_metadata_url_match_picked_over_filename_lookalike(tmp_path: Path):
|
||||
"""URL match beats a longer-prefix sibling."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
# Closer filename prefix, wrong upstream.
|
||||
_gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-mmproj-bf16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-1.5B",
|
||||
},
|
||||
)
|
||||
# Matching upstream.
|
||||
correct = _gguf_with_general(
|
||||
tmp_path / "mmproj-BF16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(correct.resolve())
|
||||
|
||||
|
||||
def test_metadata_url_mismatch_dropped(tmp_path: Path):
|
||||
"""Filenames match family but metadata disagrees: returns None."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "qwen-9b.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
_gguf_with_general(
|
||||
tmp_path / "qwen-9b-mmproj.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) is None
|
||||
|
||||
|
||||
def test_metadata_identifies_mmproj_without_filename_hint(tmp_path: Path):
|
||||
"""Projector named ``vision-projector.gguf`` discovered via header."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
projector = _gguf_with_general(
|
||||
tmp_path / "vision-projector.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(projector.resolve())
|
||||
|
||||
|
||||
def test_metadata_score_outranks_filename_prefix(tmp_path: Path):
|
||||
"""Score 100 (URL match) beats score 0 (long filename prefix)."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
# Headerless: long shared stem, score 0.
|
||||
_touch(tmp_path / "Qwen3.5-9B-Q4_K_M-mmproj.gguf")
|
||||
# Headered: generic name, score 100.
|
||||
correct = _gguf_with_general(
|
||||
tmp_path / "mmproj-BF16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(correct.resolve())
|
||||
216
studio/backend/tests/test_gguf_metadata.py
Normal file
216
studio/backend/tests/test_gguf_metadata.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for :mod:`utils.models.gguf_metadata`. Synthesise small GGUF
|
||||
headers in tmp dirs so we never depend on real model files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Mapping
|
||||
|
||||
from utils.models.gguf_metadata import (
|
||||
is_mmproj_by_metadata,
|
||||
pairing_score,
|
||||
read_gguf_general_metadata,
|
||||
)
|
||||
|
||||
|
||||
_GGUF_MAGIC = 0x46554747
|
||||
_VTYPE_STRING = 8
|
||||
_VTYPE_UINT32 = 4
|
||||
_VTYPE_ARRAY = 9
|
||||
|
||||
|
||||
def _enc_string(s: str) -> bytes:
|
||||
b = s.encode("utf-8")
|
||||
return struct.pack("<Q", len(b)) + b
|
||||
|
||||
|
||||
def _enc_kv_string(key: str, value: str) -> bytes:
|
||||
return _enc_string(key) + struct.pack("<I", _VTYPE_STRING) + _enc_string(value)
|
||||
|
||||
|
||||
def _enc_kv_uint32(key: str, value: int) -> bytes:
|
||||
return (
|
||||
_enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
|
||||
)
|
||||
|
||||
|
||||
def _enc_kv_string_array(key: str, values: Iterable[str]) -> bytes:
|
||||
vals = list(values)
|
||||
out = _enc_string(key) + struct.pack("<I", _VTYPE_ARRAY)
|
||||
out += struct.pack("<I", _VTYPE_STRING) + struct.pack("<Q", len(vals))
|
||||
for v in vals:
|
||||
out += _enc_string(v)
|
||||
return out
|
||||
|
||||
|
||||
def _write_synthetic_gguf(
|
||||
path: Path,
|
||||
general_strings: Mapping[str, str],
|
||||
*,
|
||||
extra_uint32: Mapping[str, int] | None = None,
|
||||
extra_string_arrays: Mapping[str, Iterable[str]] | None = None,
|
||||
) -> Path:
|
||||
"""Minimal GGUF: header + KV body, no tensors."""
|
||||
extra_uint32 = extra_uint32 or {}
|
||||
extra_string_arrays = extra_string_arrays or {}
|
||||
kv_count = len(general_strings) + len(extra_uint32) + len(extra_string_arrays)
|
||||
body = b""
|
||||
for k, v in general_strings.items():
|
||||
body += _enc_kv_string(k, v)
|
||||
for k, v in extra_uint32.items():
|
||||
body += _enc_kv_uint32(k, v)
|
||||
for k, v in extra_string_arrays.items():
|
||||
body += _enc_kv_string_array(k, v)
|
||||
header = struct.pack(
|
||||
"<IIQQ",
|
||||
_GGUF_MAGIC,
|
||||
3, # version
|
||||
0, # tensor_count
|
||||
kv_count,
|
||||
)
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(header + body)
|
||||
return path
|
||||
|
||||
|
||||
# --- read_gguf_general_metadata ----------------------------------------
|
||||
|
||||
|
||||
def test_returns_none_for_missing_file(tmp_path: Path):
|
||||
assert read_gguf_general_metadata(str(tmp_path / "nope.gguf")) is None
|
||||
|
||||
|
||||
def test_returns_none_for_non_gguf(tmp_path: Path):
|
||||
p = tmp_path / "garbage.gguf"
|
||||
p.write_bytes(b"not a gguf file at all, just bytes")
|
||||
assert read_gguf_general_metadata(str(p)) is None
|
||||
|
||||
|
||||
def test_extracts_general_string_fields(tmp_path: Path):
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.organization": "Qwen",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
"general.base_model.0.name": "Qwen3.5 9B",
|
||||
"general.base_model.0.organization": "Qwen",
|
||||
},
|
||||
)
|
||||
meta = read_gguf_general_metadata(str(p))
|
||||
assert meta is not None
|
||||
assert meta["general.architecture"] == "qwen2vl"
|
||||
assert meta["general.basename"] == "Qwen3.5"
|
||||
assert (
|
||||
meta["general.base_model.0.repo_url"]
|
||||
== "https://huggingface.co/Qwen/Qwen3.5-9B"
|
||||
)
|
||||
|
||||
|
||||
def test_skips_unrelated_fields_without_breaking(tmp_path: Path):
|
||||
"""Skip unwanted arrays and uint32s without losing position."""
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "Foo"},
|
||||
extra_uint32 = {"qwen2vl.context_length": 32768},
|
||||
extra_string_arrays = {"tokenizer.ggml.tokens": ["a", "bc", "def"]},
|
||||
)
|
||||
meta = read_gguf_general_metadata(str(p))
|
||||
assert meta == {"general.basename": "Foo"}
|
||||
|
||||
|
||||
def test_metadata_is_cached(tmp_path: Path):
|
||||
"""Cache invalidates on size change."""
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "First"},
|
||||
)
|
||||
first = read_gguf_general_metadata(str(p))
|
||||
assert first == {"general.basename": "First"}
|
||||
# Force size change so the (path, mtime, size) key invalidates.
|
||||
_write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "Second", "general.organization": "X"},
|
||||
)
|
||||
second = read_gguf_general_metadata(str(p))
|
||||
assert second == {"general.basename": "Second", "general.organization": "X"}
|
||||
|
||||
|
||||
# --- is_mmproj_by_metadata --------------------------------------------
|
||||
|
||||
|
||||
def test_is_mmproj_by_metadata_signals():
|
||||
assert is_mmproj_by_metadata({"general.type": "mmproj"}) is True
|
||||
assert is_mmproj_by_metadata({"general.type": "MMProj"}) is True
|
||||
assert is_mmproj_by_metadata({"general.type": "model"}) is False
|
||||
assert is_mmproj_by_metadata({"general.basename": "foo"}) is None
|
||||
assert is_mmproj_by_metadata({}) is None
|
||||
assert is_mmproj_by_metadata(None) is None
|
||||
|
||||
|
||||
# --- pairing_score -----------------------------------------------------
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_match():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 100
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_mismatch():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == -1
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_trailing_slash_normalised():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B/",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 100
|
||||
|
||||
|
||||
def test_pairing_score_basename_plus_org_fallback():
|
||||
weight = {
|
||||
"general.basename": "Nanonets-Ocr-S",
|
||||
"general.base_model.0.organization": "Nanonets",
|
||||
}
|
||||
mmproj = {
|
||||
"general.basename": "Nanonets-Ocr-S",
|
||||
"general.base_model.0.organization": "Nanonets",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 80
|
||||
|
||||
|
||||
def test_pairing_score_basename_only_fallback():
|
||||
assert (
|
||||
pairing_score(
|
||||
{"general.basename": "Nanonets-Ocr-S"},
|
||||
{"general.basename": "Nanonets-Ocr-S"},
|
||||
)
|
||||
== 60
|
||||
)
|
||||
|
||||
|
||||
def test_pairing_score_no_overlap_returns_zero():
|
||||
"""One side empty: scorer punts to filename fallback."""
|
||||
assert pairing_score({"general.basename": "Foo"}, {}) == 0
|
||||
assert pairing_score({}, {"general.basename": "Foo"}) == 0
|
||||
assert pairing_score(None, {"general.basename": "Foo"}) == 0
|
||||
233
studio/backend/utils/models/gguf_metadata.py
Normal file
233
studio/backend/utils/models/gguf_metadata.py
Normal file
|
|
@ -0,0 +1,233 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Free-function ``general.*`` reader for GGUF headers, used by
|
||||
``detect_mmproj_file`` to pair weights and projectors via
|
||||
``general.base_model.0.repo_url``. ~30 ms per file, cached by
|
||||
(path, mtime, size)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import struct
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
_GGUF_MAGIC = 0x46554747 # b"GGUF" LE u32
|
||||
|
||||
_WANTED_GENERAL_KEYS: frozenset[str] = frozenset(
|
||||
{
|
||||
"general.architecture",
|
||||
"general.type",
|
||||
"general.name",
|
||||
"general.basename",
|
||||
"general.organization",
|
||||
"general.size_label",
|
||||
"general.finetune",
|
||||
"general.base_model.0.name",
|
||||
"general.base_model.0.organization",
|
||||
"general.base_model.0.repo_url",
|
||||
"general.repo_url",
|
||||
"general.source.url",
|
||||
"general.source.repo_url",
|
||||
"general.source.huggingface.repository",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Cache failed parses too so a broken file is not retried each scan.
|
||||
_CacheKey = Tuple[str, int, int]
|
||||
_METADATA_CACHE: Dict[_CacheKey, Optional[Dict[str, str]]] = {}
|
||||
_CACHE_LOCK = threading.Lock()
|
||||
_CACHE_MAX_ENTRIES = 4096
|
||||
|
||||
|
||||
def _cache_key(path: str) -> Optional[_CacheKey]:
|
||||
try:
|
||||
st = os.stat(path)
|
||||
except OSError:
|
||||
return None
|
||||
try:
|
||||
resolved = str(Path(path).resolve())
|
||||
except OSError:
|
||||
resolved = str(path)
|
||||
return (resolved, st.st_mtime_ns, st.st_size)
|
||||
|
||||
|
||||
def read_gguf_general_metadata(path: str) -> Optional[Dict[str, str]]:
|
||||
"""Return ``general.*`` strings from a GGUF header, or ``None`` if
|
||||
the file is missing, unreadable, or not a GGUF. ``{}`` means the
|
||||
file is valid but carries none of the wanted keys."""
|
||||
key = _cache_key(path)
|
||||
if key is None:
|
||||
return None
|
||||
with _CACHE_LOCK:
|
||||
if key in _METADATA_CACHE:
|
||||
return _METADATA_CACHE[key]
|
||||
result = _parse_gguf_header(path)
|
||||
with _CACHE_LOCK:
|
||||
# Arbitrary eviction; header reads are cheap so true LRU is overkill.
|
||||
while len(_METADATA_CACHE) >= _CACHE_MAX_ENTRIES:
|
||||
try:
|
||||
_METADATA_CACHE.pop(next(iter(_METADATA_CACHE)))
|
||||
except StopIteration:
|
||||
break
|
||||
_METADATA_CACHE[key] = result
|
||||
return result
|
||||
|
||||
|
||||
def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]:
|
||||
out: Dict[str, str] = {}
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
head = f.read(24)
|
||||
if len(head) < 24:
|
||||
return None
|
||||
magic, _version, _tcount, kv_count = struct.unpack("<IIQQ", head)
|
||||
if magic != _GGUF_MAGIC:
|
||||
return None
|
||||
|
||||
for _ in range(kv_count):
|
||||
try:
|
||||
klen_bytes = f.read(8)
|
||||
if len(klen_bytes) < 8:
|
||||
break
|
||||
klen = struct.unpack("<Q", klen_bytes)[0]
|
||||
if klen > 1 << 20: # 1 MB sanity bound
|
||||
break
|
||||
kbytes = f.read(klen)
|
||||
if len(kbytes) < klen:
|
||||
break
|
||||
key = kbytes.decode("utf-8", "replace")
|
||||
vt_bytes = f.read(4)
|
||||
if len(vt_bytes) < 4:
|
||||
break
|
||||
vtype = struct.unpack("<I", vt_bytes)[0]
|
||||
|
||||
if vtype == 8 and key in _WANTED_GENERAL_KEYS:
|
||||
slen_bytes = f.read(8)
|
||||
if len(slen_bytes) < 8:
|
||||
break
|
||||
slen = struct.unpack("<Q", slen_bytes)[0]
|
||||
if slen > 1 << 22: # 4 MB sanity bound
|
||||
break
|
||||
sbytes = f.read(slen)
|
||||
if len(sbytes) < slen:
|
||||
break
|
||||
out[key] = sbytes.decode("utf-8", "replace")
|
||||
else:
|
||||
if not _skip_gguf_value(f, vtype):
|
||||
break
|
||||
except (struct.error, UnicodeDecodeError):
|
||||
break
|
||||
except OSError as e:
|
||||
logger.debug(f"read_gguf_general_metadata: cannot open {path}: {e}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug(f"read_gguf_general_metadata: parse failure on {path}: {e}")
|
||||
return None
|
||||
return out
|
||||
|
||||
|
||||
# Strings (8) and arrays (9) are handled inline.
|
||||
_FIXED_VTYPE_SIZES: Dict[int, int] = {
|
||||
0: 1, # uint8
|
||||
1: 1, # int8
|
||||
2: 2, # uint16
|
||||
3: 2, # int16
|
||||
4: 4, # uint32
|
||||
5: 4, # int32
|
||||
6: 4, # float32
|
||||
7: 1, # bool
|
||||
10: 8, # uint64
|
||||
11: 8, # int64
|
||||
12: 8, # float64
|
||||
}
|
||||
|
||||
|
||||
def _skip_gguf_value(f, vtype: int) -> bool:
|
||||
"""Advance past one GGUF value. False on truncation or unknown type."""
|
||||
if vtype == 8: # STRING
|
||||
slen_bytes = f.read(8)
|
||||
if len(slen_bytes) < 8:
|
||||
return False
|
||||
slen = struct.unpack("<Q", slen_bytes)[0]
|
||||
if slen > 1 << 30: # 1 GB sanity bound
|
||||
return False
|
||||
return len(f.read(slen)) == slen
|
||||
if vtype == 9: # ARRAY
|
||||
head = f.read(12)
|
||||
if len(head) < 12:
|
||||
return False
|
||||
atype, alen = struct.unpack("<IQ", head)
|
||||
if alen > 1 << 30:
|
||||
return False
|
||||
if atype == 8:
|
||||
for _ in range(alen):
|
||||
slen_bytes = f.read(8)
|
||||
if len(slen_bytes) < 8:
|
||||
return False
|
||||
slen = struct.unpack("<Q", slen_bytes)[0]
|
||||
if slen > 1 << 30:
|
||||
return False
|
||||
if len(f.read(slen)) != slen:
|
||||
return False
|
||||
return True
|
||||
sz = _FIXED_VTYPE_SIZES.get(atype)
|
||||
if sz is None:
|
||||
return False
|
||||
total = sz * alen
|
||||
return len(f.read(total)) == total
|
||||
sz = _FIXED_VTYPE_SIZES.get(vtype)
|
||||
if sz is None:
|
||||
return False
|
||||
return len(f.read(sz)) == sz
|
||||
|
||||
|
||||
def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]:
|
||||
"""True/False from ``general.type``; None means fall back to filename."""
|
||||
if not meta:
|
||||
return None
|
||||
t = meta.get("general.type")
|
||||
if t is None:
|
||||
return None
|
||||
return t.lower() == "mmproj"
|
||||
|
||||
|
||||
def pairing_score(
|
||||
weight_meta: Optional[Dict[str, str]],
|
||||
mmproj_meta: Optional[Dict[str, str]],
|
||||
) -> int:
|
||||
"""Pairing confidence: 100 = base_model URL match, 80 = basename + org,
|
||||
60 = basename, -1 = definitive mismatch, 0 = decide from filename."""
|
||||
if not weight_meta or not mmproj_meta:
|
||||
return 0
|
||||
|
||||
w_url = weight_meta.get("general.base_model.0.repo_url")
|
||||
p_url = mmproj_meta.get("general.base_model.0.repo_url")
|
||||
if w_url and p_url:
|
||||
return 100 if w_url.strip().rstrip("/") == p_url.strip().rstrip("/") else -1
|
||||
|
||||
w_base = weight_meta.get("general.basename")
|
||||
p_base = mmproj_meta.get("general.basename")
|
||||
w_org = weight_meta.get("general.base_model.0.organization") or weight_meta.get(
|
||||
"general.organization"
|
||||
)
|
||||
p_org = mmproj_meta.get("general.base_model.0.organization") or mmproj_meta.get(
|
||||
"general.organization"
|
||||
)
|
||||
if w_base and p_base and w_org and p_org:
|
||||
if w_base.lower() == p_base.lower() and w_org.lower() == p_org.lower():
|
||||
return 80
|
||||
return -1
|
||||
|
||||
if w_base and p_base:
|
||||
return 60 if w_base.lower() == p_base.lower() else -1
|
||||
|
||||
return 0
|
||||
|
|
@ -19,6 +19,11 @@ from utils.paths import (
|
|||
resolve_export_dir,
|
||||
)
|
||||
from utils.utils import without_hf_auth
|
||||
from utils.models.gguf_metadata import (
|
||||
is_mmproj_by_metadata,
|
||||
pairing_score,
|
||||
read_gguf_general_metadata,
|
||||
)
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
import os
|
||||
|
|
@ -801,12 +806,15 @@ _AUDIO_TOKEN_PATTERNS = {
|
|||
"whisper": lambda tokens: "<|startoftranscript|>" in tokens,
|
||||
"audio_vlm": lambda tokens: "<audio_soft_token>" in tokens,
|
||||
"bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
|
||||
"dac": lambda tokens: "<|audio_start|>" in tokens
|
||||
and "<|audio_end|>" in tokens
|
||||
and "<|text_start|>" in tokens
|
||||
and "<|text_end|>" in tokens,
|
||||
"snac": lambda tokens: sum(1 for t in tokens if t.startswith("<custom_token_"))
|
||||
> 10000,
|
||||
"dac": lambda tokens: (
|
||||
"<|audio_start|>" in tokens
|
||||
and "<|audio_end|>" in tokens
|
||||
and "<|text_start|>" in tokens
|
||||
and "<|text_end|>" in tokens
|
||||
),
|
||||
"snac": lambda tokens: (
|
||||
sum(1 for t in tokens if t.startswith("<custom_token_")) > 10000
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -913,6 +921,85 @@ def _is_mmproj(filename: str) -> bool:
|
|||
return "mmproj" in filename.lower()
|
||||
|
||||
|
||||
# Family tokens for #5347's filename fallback. Lowercase. Order does not
|
||||
# matter (see ``_detect_family_token``).
|
||||
_MODEL_FAMILY_TOKENS: tuple[str, ...] = (
|
||||
"qwen",
|
||||
"gemma",
|
||||
"llama",
|
||||
"mistral",
|
||||
"ministral",
|
||||
"magistral",
|
||||
"devstral",
|
||||
"phi",
|
||||
"deepseek",
|
||||
"internvl",
|
||||
"minicpm",
|
||||
"llava",
|
||||
"glm",
|
||||
"yi",
|
||||
"command-r",
|
||||
"molmo",
|
||||
"pixtral",
|
||||
"smolvlm",
|
||||
"moondream",
|
||||
"granite",
|
||||
"ovis",
|
||||
"nemotron",
|
||||
"kimi",
|
||||
"nanonets",
|
||||
"cosmos",
|
||||
"mimo",
|
||||
"apriel",
|
||||
"lfm",
|
||||
)
|
||||
|
||||
|
||||
# Word-bounded match: any letter on either side disqualifies. Stops
|
||||
# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc.
|
||||
_FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {}
|
||||
|
||||
|
||||
def _family_token_re(token: str) -> "_re.Pattern[str]":
|
||||
pat = _FAMILY_TOKEN_RE_CACHE.get(token)
|
||||
if pat is None:
|
||||
pat = _re.compile(rf"(?:^|[^a-z])({_re.escape(token)})(?:[^a-z]|$)")
|
||||
_FAMILY_TOKEN_RE_CACHE[token] = pat
|
||||
return pat
|
||||
|
||||
|
||||
def _detect_family_token(filename: str) -> Optional[str]:
|
||||
"""Leftmost-position match; ties prefer the longer token."""
|
||||
name = filename.lower()
|
||||
best: Optional[tuple[int, int, str]] = None # (start, -len, token)
|
||||
for token in _MODEL_FAMILY_TOKENS:
|
||||
m = _family_token_re(token).search(name)
|
||||
if m is None:
|
||||
continue
|
||||
key = (m.start(1), -len(token), token)
|
||||
if best is None or key < best:
|
||||
best = key
|
||||
return None if best is None else best[2]
|
||||
|
||||
|
||||
def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool:
|
||||
"""Defense-in-depth guard for the launcher: True unless both filenames
|
||||
carry recognised family tokens that disagree."""
|
||||
model_fam = _detect_family_token(Path(model_path).name)
|
||||
mmproj_fam = _detect_family_token(Path(mmproj_path).name)
|
||||
if model_fam is None or mmproj_fam is None:
|
||||
return True
|
||||
return model_fam == mmproj_fam
|
||||
|
||||
|
||||
def _shared_prefix_len(a: str, b: str) -> int:
|
||||
n = min(len(a), len(b))
|
||||
for i in range(n):
|
||||
if a[i] != b[i]:
|
||||
return i
|
||||
return n
|
||||
|
||||
|
||||
def _is_gguf_filename(filename: str) -> bool:
|
||||
return filename.lower().endswith(".gguf")
|
||||
|
||||
|
|
@ -927,33 +1014,18 @@ def _iter_gguf_files(directory: Path, recursive: bool = False):
|
|||
|
||||
|
||||
def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
|
||||
"""
|
||||
Find the mmproj (vision projection) GGUF file for a given model.
|
||||
"""Find the mmproj GGUF for a model.
|
||||
|
||||
Args:
|
||||
path: Directory to search — or a .gguf file (uses its parent dir
|
||||
as the starting point).
|
||||
search_root: Optional outer directory that should also be scanned
|
||||
(and any directory between it and ``path``). This handles
|
||||
local layouts where the model weights live in a quant-named
|
||||
subdir (``snapshot/BF16/foo.gguf``) but the mmproj sits at
|
||||
the snapshot root (``snapshot/mmproj-BF16.gguf``). When
|
||||
``None``, only the immediate parent dir is scanned, matching
|
||||
the historical behavior.
|
||||
|
||||
Returns:
|
||||
Full path to the mmproj .gguf file, or None if not found.
|
||||
"""
|
||||
``path``: directory or a .gguf file. ``search_root``: optional ancestor
|
||||
to also walk (snapshot layouts where the weight is in ``snapshot/BF16/``
|
||||
but the projector sits at ``snapshot/``). Returns the projector path or
|
||||
``None``."""
|
||||
p = Path(path)
|
||||
start_dir = p.parent if p.is_file() else p
|
||||
if not start_dir.is_dir():
|
||||
return None
|
||||
|
||||
# Build the list of dirs to scan: immediate dir first, then walk up
|
||||
# to (and including) ``search_root`` if it is an ancestor. We walk
|
||||
# incrementally rather than recursing into ``search_root`` so we
|
||||
# don't accidentally pick up an mmproj from a sibling subdir
|
||||
# belonging to a different model variant.
|
||||
# Walk incrementally so a sibling subdir's mmproj cannot leak in.
|
||||
seen: set[Path] = set()
|
||||
scan_order: list[Path] = []
|
||||
|
||||
|
|
@ -969,12 +1041,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
|
|||
|
||||
_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.
|
||||
# Ollama's .studio_links/foo.gguf -> blobs/sha256-...: also scan target dir.
|
||||
try:
|
||||
if p.is_symlink() and p.is_file():
|
||||
target_parent = p.resolve().parent
|
||||
|
|
@ -986,14 +1053,12 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
|
|||
try:
|
||||
root_resolved = Path(search_root).resolve()
|
||||
start_resolved = start_dir.resolve()
|
||||
# Only walk if start_dir is inside (or equal to) search_root.
|
||||
if root_resolved == start_resolved or (
|
||||
start_resolved.is_relative_to(root_resolved)
|
||||
if hasattr(start_resolved, "is_relative_to")
|
||||
else str(start_resolved).startswith(str(root_resolved) + "/")
|
||||
):
|
||||
cur = start_resolved
|
||||
# Walk up from start_dir to (and including) root_resolved.
|
||||
while cur != root_resolved and cur.parent != cur:
|
||||
cur = cur.parent
|
||||
_add(cur)
|
||||
|
|
@ -1002,11 +1067,66 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
|
|||
except OSError:
|
||||
pass
|
||||
|
||||
candidates: list[Path] = []
|
||||
seen_resolved: set[Path] = set()
|
||||
for d in scan_order:
|
||||
for f in _iter_gguf_files(d):
|
||||
if _is_mmproj(f.name):
|
||||
return str(f.resolve())
|
||||
return None
|
||||
try:
|
||||
resolved = f.resolve()
|
||||
except OSError:
|
||||
continue
|
||||
if resolved in seen_resolved:
|
||||
continue
|
||||
# Prefer ``general.type=='mmproj'``; fall back to filename.
|
||||
meta = read_gguf_general_metadata(str(resolved))
|
||||
by_meta = is_mmproj_by_metadata(meta)
|
||||
if by_meta is True or (by_meta is None and _is_mmproj(f.name)):
|
||||
seen_resolved.add(resolved)
|
||||
candidates.append(resolved)
|
||||
|
||||
if not candidates:
|
||||
return None
|
||||
|
||||
# Directory path: no model name to compare against; legacy behaviour.
|
||||
if not p.is_file():
|
||||
return str(candidates[0])
|
||||
|
||||
# Stage 1: GGUF metadata. Stage 2: filename family token (#5347).
|
||||
model_stem = p.stem.lower()
|
||||
model_family = _detect_family_token(p.name)
|
||||
weight_meta = read_gguf_general_metadata(str(p))
|
||||
|
||||
scored: list[tuple[int, Path]] = []
|
||||
for c in candidates:
|
||||
cand_meta = read_gguf_general_metadata(str(c))
|
||||
meta_score = pairing_score(weight_meta, cand_meta)
|
||||
if meta_score == -1:
|
||||
logger.info(f"detect_mmproj_file: dropped {c.name} (metadata mismatch)")
|
||||
continue
|
||||
if meta_score == 0 and model_family is not None:
|
||||
# Unrecognised candidate family is a wildcard (``mmproj-F16.gguf``).
|
||||
cand_family = _detect_family_token(c.name)
|
||||
if cand_family is not None and cand_family != model_family:
|
||||
logger.info(
|
||||
f"detect_mmproj_file: dropped {c.name} "
|
||||
f"(filename family {cand_family!r} vs model {model_family!r})"
|
||||
)
|
||||
continue
|
||||
scored.append((meta_score, c))
|
||||
|
||||
if not scored:
|
||||
return None
|
||||
|
||||
# Score first, then longest shared prefix, then shorter stem.
|
||||
best = max(
|
||||
scored,
|
||||
key = lambda sc: (
|
||||
sc[0],
|
||||
_shared_prefix_len(model_stem, sc[1].stem.lower()),
|
||||
-len(sc[1].stem),
|
||||
),
|
||||
)
|
||||
return str(best[1])
|
||||
|
||||
|
||||
def detect_gguf_model(path: str) -> Optional[str]:
|
||||
|
|
@ -1360,7 +1480,7 @@ def detect_gguf_model_remote(
|
|||
if attempt < 2:
|
||||
time.sleep(2**attempt)
|
||||
logger.warning(
|
||||
f"Could not check GGUF files for '{repo_id}' after 3 attempts: " f"{last_err}"
|
||||
f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue