Improve model matching by normalizing technical suffixes and family signatures
This commit is contained in:
+150
-11
@@ -292,10 +292,35 @@ def base_model_score(candidate: str, preferred: Optional[str]) -> float:
|
||||
|
||||
# ==================== CENTRALIZED MATCHING & CONFIDENCE HELPERS ====================
|
||||
|
||||
from .type_utils import MODEL_EXTENSIONS, PRECISION_FORMAT_SUFFIXES, get_version_sort_key
|
||||
from .type_utils import MODEL_EXTENSIONS, TECHNICAL_MODEL_SUFFIXES, get_version_sort_key
|
||||
|
||||
MODEL_FILE_EXTENSIONS = MODEL_EXTENSIONS
|
||||
|
||||
_TECHNICAL_VARIANT_TOKENS = {
|
||||
*TECHNICAL_MODEL_SUFFIXES,
|
||||
}
|
||||
_TECHNICAL_VARIANT_TOKEN_RE = re.compile(
|
||||
r"^(?:(?:bf|bfloat|fp|float|int|uint|f|i)\d+|"
|
||||
r"(?:i?q|tq)\d+|(?:mxfp|nvfp)\d+|e\d+m\d+(?:fnuz|fn)?|bnb\d+)$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_QUANTIZATION_HEAD_TOKEN_RE = re.compile(r"^(?:i?q|tq)\d+$", re.IGNORECASE)
|
||||
_QUANTIZATION_TAIL_TOKENS = {
|
||||
"0",
|
||||
"1",
|
||||
"4",
|
||||
"8",
|
||||
"k",
|
||||
"m",
|
||||
"s",
|
||||
"l",
|
||||
"xl",
|
||||
"xs",
|
||||
"xxs",
|
||||
"nl",
|
||||
}
|
||||
_MODEL_CAPACITY_TOKEN_RE = re.compile(r"^\d+(?:\.\d+)?[bmk]$", re.IGNORECASE)
|
||||
|
||||
|
||||
def strip_known_model_extension(filename: str) -> str:
|
||||
"""Strip only known model extensions, preserving names like v4.0."""
|
||||
@@ -321,6 +346,54 @@ def normalize_model_title(value: str) -> str:
|
||||
return re.sub(r"\s+", " ", value).strip()
|
||||
|
||||
|
||||
def normalize_model_family_filename(filename: str) -> str:
|
||||
"""Normalize a filename while ignoring precision and quantization variants."""
|
||||
normalized_filename = normalize_filename(filename)
|
||||
normalized_filename = re.sub(r"\bema\s+only\b", "emaonly", normalized_filename)
|
||||
tokens = normalized_filename.split()
|
||||
family_tokens = []
|
||||
quantization_tail = False
|
||||
|
||||
for token in tokens:
|
||||
token_lower = token.lower()
|
||||
is_quantization_token = bool(
|
||||
_TECHNICAL_VARIANT_TOKEN_RE.fullmatch(token_lower)
|
||||
)
|
||||
if token_lower in _TECHNICAL_VARIANT_TOKENS or is_quantization_token:
|
||||
quantization_tail = bool(
|
||||
_QUANTIZATION_HEAD_TOKEN_RE.fullmatch(token_lower)
|
||||
)
|
||||
continue
|
||||
if quantization_tail and token_lower in _QUANTIZATION_TAIL_TOKENS:
|
||||
continue
|
||||
|
||||
quantization_tail = False
|
||||
family_tokens.append(token_lower)
|
||||
|
||||
return " ".join(family_tokens)
|
||||
|
||||
|
||||
def _get_model_family_signature(normalized_family_name: str) -> Optional[Tuple[str, str]]:
|
||||
"""Return the architecture prefix and parameter count when both are available."""
|
||||
tokens = normalized_family_name.split()
|
||||
for index, token in enumerate(tokens):
|
||||
if not _MODEL_CAPACITY_TOKEN_RE.fullmatch(token):
|
||||
continue
|
||||
architecture = "".join(tokens[:index])
|
||||
if len(architecture) >= 3 and re.search(r"[a-z]", architecture):
|
||||
return architecture, token
|
||||
return None
|
||||
|
||||
|
||||
def _architectures_have_numeric_conflict(left: str, right: str) -> bool:
|
||||
"""Detect otherwise-identical architecture names with different generations."""
|
||||
left_digits = re.findall(r"\d+", left)
|
||||
right_digits = re.findall(r"\d+", right)
|
||||
if not left_digits or not right_digits or left_digits == right_digits:
|
||||
return False
|
||||
return re.sub(r"\d+", "", left) == re.sub(r"\d+", "", right)
|
||||
|
||||
|
||||
def calculate_model_title_confidence(query: str, model_name: str) -> float:
|
||||
"""Calculate confidence score (0.0 to 100.0) for a model title matching a query."""
|
||||
query_norm = normalize_model_title(query)
|
||||
@@ -351,9 +424,57 @@ def calculate_filename_confidence(target_filename: str, candidate_filename: str)
|
||||
if target_norm == candidate_norm:
|
||||
return 100.0
|
||||
|
||||
best_similarity = calculate_similarity_with_normalization(
|
||||
target_filename, candidate_filename
|
||||
)
|
||||
best_similarity = calculate_similarity(target_norm, candidate_norm)
|
||||
target_family = normalize_model_family_filename(target_filename)
|
||||
candidate_family = normalize_model_family_filename(candidate_filename)
|
||||
|
||||
if target_family and candidate_family:
|
||||
family_similarity = calculate_similarity(target_family, candidate_family)
|
||||
target_signature = _get_model_family_signature(target_family)
|
||||
candidate_signature = _get_model_family_signature(candidate_family)
|
||||
|
||||
if target_family == candidate_family:
|
||||
# Same model name with a different precision or quantization is a
|
||||
# strong alternative, but remains below an exact filename match.
|
||||
best_similarity = max(best_similarity, 0.94)
|
||||
else:
|
||||
best_similarity = max(best_similarity, family_similarity)
|
||||
|
||||
if (
|
||||
target_family != candidate_family
|
||||
and target_signature
|
||||
and target_signature == candidate_signature
|
||||
):
|
||||
target_tokens = set(target_family.split())
|
||||
candidate_tokens = set(candidate_family.split())
|
||||
if target_tokens.issubset(candidate_tokens) or candidate_tokens.issubset(
|
||||
target_tokens
|
||||
):
|
||||
# One filename adds a semantic variant such as "abliterated".
|
||||
best_similarity = max(best_similarity, 0.90)
|
||||
else:
|
||||
# The architecture and size agree, but both names contain
|
||||
# different semantic modifiers such as "base" and "heretic".
|
||||
best_similarity = max(best_similarity, 0.84)
|
||||
elif (
|
||||
target_family != candidate_family
|
||||
and target_signature
|
||||
and candidate_signature
|
||||
):
|
||||
target_architecture, target_capacity = target_signature
|
||||
candidate_architecture, candidate_capacity = candidate_signature
|
||||
if target_architecture == candidate_architecture and target_capacity != candidate_capacity:
|
||||
# A 4B model is not an interchangeable precision variant of 8B.
|
||||
best_similarity = min(best_similarity, 0.69)
|
||||
elif (
|
||||
target_capacity == candidate_capacity
|
||||
and _architectures_have_numeric_conflict(
|
||||
target_architecture, candidate_architecture
|
||||
)
|
||||
):
|
||||
# Keep adjacent generations such as Qwen2 and Qwen3 out of the
|
||||
# visible recommendation threshold unless another exact signal exists.
|
||||
best_similarity = min(best_similarity, 0.69)
|
||||
|
||||
# Cap similarity at 0.999 for non-exact matches to prevent false 100% scores
|
||||
if best_similarity >= 0.999:
|
||||
@@ -380,7 +501,7 @@ def calculate_archived_model_confidence(
|
||||
if query_norm == candidate_norm:
|
||||
return 100.0
|
||||
|
||||
score = calculate_similarity_with_normalization(query, candidate)
|
||||
score = calculate_filename_confidence(query, candidate) / 100.0
|
||||
if query_norm and candidate_norm and (
|
||||
query_norm in candidate_norm or candidate_norm in query_norm
|
||||
):
|
||||
@@ -413,15 +534,33 @@ def build_filename_search_queries(filename: str) -> List[str]:
|
||||
stem = os.path.splitext(basename)[0].strip()
|
||||
stem_lower = stem.lower()
|
||||
|
||||
# Pattern to match common precision/format suffixes
|
||||
pattern = r"[-_]?(" + "|".join(PRECISION_FORMAT_SUFFIXES) + r")"
|
||||
literal_suffixes = sorted(
|
||||
(re.escape(token) for token in TECHNICAL_MODEL_SUFFIXES),
|
||||
key=len,
|
||||
reverse=True,
|
||||
)
|
||||
technical_token_pattern = (
|
||||
r"(?:^|[-_.\s])(?:"
|
||||
+ "|".join(
|
||||
[
|
||||
*literal_suffixes,
|
||||
r"(?:bf|bfloat|fp|float|int|uint|f|i)\d+",
|
||||
r"(?:i?q|tq)\d+",
|
||||
r"(?:mxfp|nvfp)\d+",
|
||||
r"e\d+m\d+(?:fnuz|fn)?",
|
||||
r"bnb\d+",
|
||||
]
|
||||
)
|
||||
+ r")(?=$|[-_.\s])"
|
||||
)
|
||||
|
||||
# 1. Strip suffix followed by any other characters (like clean_filename_for_search)
|
||||
cleaned = re.sub(pattern + r".*$", "", stem, flags=re.IGNORECASE).strip(" -_")
|
||||
technical_match = re.search(technical_token_pattern, stem, flags=re.IGNORECASE)
|
||||
cleaned = (
|
||||
stem[: technical_match.start()] if technical_match else stem
|
||||
).strip(" -_.")
|
||||
cleaned_lower = cleaned.lower()
|
||||
|
||||
# 2. Strip suffix at the end of the stem (like simplified in search queries)
|
||||
simplified = re.sub(pattern + r"$", "", stem_lower, flags=re.IGNORECASE).strip(" -_")
|
||||
simplified = cleaned_lower
|
||||
|
||||
queries = []
|
||||
for query in [basename, stem, stem_lower, cleaned, cleaned_lower, simplified]:
|
||||
|
||||
+90
-8
@@ -62,26 +62,108 @@ MODEL_EXTENSIONS = {
|
||||
}
|
||||
|
||||
|
||||
PRECISION_FORMAT_SUFFIXES = [
|
||||
"fp16",
|
||||
NUMERIC_PRECISION_SUFFIXES = (
|
||||
"fp64",
|
||||
"float64",
|
||||
"f64",
|
||||
"fp32",
|
||||
"fp8",
|
||||
"fp4",
|
||||
"float32",
|
||||
"f32",
|
||||
"fp16",
|
||||
"float16",
|
||||
"f16",
|
||||
"half",
|
||||
"bf16",
|
||||
"bfloat16",
|
||||
"fp8",
|
||||
"float8",
|
||||
"e4m3",
|
||||
"e4m3fn",
|
||||
"e4m3fnuz",
|
||||
"e5m2",
|
||||
"e5m2fnuz",
|
||||
"fp4",
|
||||
"float4",
|
||||
"e2m1",
|
||||
"int8",
|
||||
"uint8",
|
||||
"int4",
|
||||
"uint4",
|
||||
"int3",
|
||||
"uint3",
|
||||
"int2",
|
||||
"uint2",
|
||||
"nf4",
|
||||
"mxfp4",
|
||||
"nvfp4",
|
||||
)
|
||||
|
||||
|
||||
QUANTIZATION_LEVEL_SUFFIXES = (
|
||||
"q1",
|
||||
"q2",
|
||||
"q3",
|
||||
"q4",
|
||||
"q5",
|
||||
"q6",
|
||||
"q7",
|
||||
"q8",
|
||||
"iq1",
|
||||
"iq2",
|
||||
"iq3",
|
||||
"iq4",
|
||||
"tq1",
|
||||
"tq2",
|
||||
)
|
||||
|
||||
|
||||
QUANTIZATION_SCHEME_SUFFIXES = (
|
||||
"gptq",
|
||||
"awq",
|
||||
"exl2",
|
||||
"hqq",
|
||||
"aqlm",
|
||||
"bnb4",
|
||||
"bnb8",
|
||||
)
|
||||
|
||||
|
||||
MODEL_CONTAINER_SUFFIXES = (
|
||||
"ggml",
|
||||
"gguf",
|
||||
)
|
||||
|
||||
|
||||
MODEL_VARIANT_SUFFIXES = (
|
||||
"mixed",
|
||||
"scaled",
|
||||
"pruned",
|
||||
"ema",
|
||||
"emaonly",
|
||||
"q4",
|
||||
"q8",
|
||||
]
|
||||
"quant",
|
||||
"quantized",
|
||||
"converted",
|
||||
"conversion",
|
||||
"compressed",
|
||||
"conv",
|
||||
)
|
||||
|
||||
|
||||
TECHNICAL_MODEL_SUFFIXES = tuple(
|
||||
dict.fromkeys(
|
||||
NUMERIC_PRECISION_SUFFIXES
|
||||
+ QUANTIZATION_LEVEL_SUFFIXES
|
||||
+ QUANTIZATION_SCHEME_SUFFIXES
|
||||
+ MODEL_CONTAINER_SUFFIXES
|
||||
+ MODEL_VARIANT_SUFFIXES
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_generic_filename_tokens() -> set:
|
||||
"""Return a set of generic filename tokens (extensions and suffixes) in lowercase."""
|
||||
dotless_extensions = {ext.lstrip(".") for ext in MODEL_EXTENSIONS}
|
||||
return dotless_extensions | set(PRECISION_FORMAT_SUFFIXES)
|
||||
return dotless_extensions | set(TECHNICAL_MODEL_SUFFIXES)
|
||||
|
||||
|
||||
|
||||
|
||||
+195
-2
@@ -1,8 +1,26 @@
|
||||
import unittest
|
||||
from core.matcher import normalize_filename, calculate_similarity, find_matches, build_filename_search_queries
|
||||
from core.matcher import (
|
||||
build_filename_search_queries,
|
||||
calculate_archived_model_confidence,
|
||||
calculate_filename_confidence,
|
||||
calculate_similarity,
|
||||
find_matches,
|
||||
normalize_filename,
|
||||
normalize_model_family_filename,
|
||||
)
|
||||
from core.sources.model_list import search_model_list, search_model_list_multiple
|
||||
from core.progress import report_progress
|
||||
from core.type_utils import as_dict, as_list, format_size_bytes
|
||||
from core.type_utils import (
|
||||
MODEL_CONTAINER_SUFFIXES,
|
||||
MODEL_VARIANT_SUFFIXES,
|
||||
NUMERIC_PRECISION_SUFFIXES,
|
||||
QUANTIZATION_LEVEL_SUFFIXES,
|
||||
QUANTIZATION_SCHEME_SUFFIXES,
|
||||
TECHNICAL_MODEL_SUFFIXES,
|
||||
as_dict,
|
||||
as_list,
|
||||
format_size_bytes,
|
||||
)
|
||||
class MatcherTests(unittest.TestCase):
|
||||
def test_normalize_filename_removes_extension_and_converts_lowercase(self):
|
||||
self.assertEqual("test model", normalize_filename("Test_Model.safetensors"))
|
||||
@@ -32,6 +50,120 @@ class MatcherTests(unittest.TestCase):
|
||||
matches_filtered = find_matches("sd_xl_base_1.0.safetensors", candidates, threshold=0.7)
|
||||
self.assertEqual(1, len(matches_filtered))
|
||||
|
||||
def test_model_family_normalization_ignores_precision_and_quantization(self):
|
||||
self.assertEqual(
|
||||
"qwen3vl 4b abliterated",
|
||||
normalize_model_family_filename(
|
||||
"qwen3vl-4b-abliterated_fp8_e4m3fn_scaled.safetensors"
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
"qwen3vl 4b abliterated",
|
||||
normalize_model_family_filename(
|
||||
"qwen3vl-4b-abliterated_int8_conv.safetensors"
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
"qwen3vl 4b",
|
||||
normalize_model_family_filename("qwen3vl_4b_iq4_nl.gguf"),
|
||||
)
|
||||
|
||||
def test_model_family_normalization_handles_common_technical_families(self):
|
||||
suffixes = [
|
||||
"float16",
|
||||
"half",
|
||||
"e4m3fnuz",
|
||||
"e5m2fnuz",
|
||||
"e2m1",
|
||||
"nf4",
|
||||
"q4_0",
|
||||
"q3_k_m",
|
||||
"iq2_xxs",
|
||||
"iq4_nl",
|
||||
"tq1_0",
|
||||
"mxfp4",
|
||||
"nvfp4",
|
||||
"gptq",
|
||||
"awq",
|
||||
"exl2",
|
||||
"hqq",
|
||||
"aqlm",
|
||||
"bnb4",
|
||||
"ema-only",
|
||||
"pruned",
|
||||
]
|
||||
|
||||
for suffix in suffixes:
|
||||
with self.subTest(suffix=suffix):
|
||||
self.assertEqual(
|
||||
"qwen3vl 4b",
|
||||
normalize_model_family_filename(
|
||||
f"qwen3vl_4b_{suffix}.safetensors"
|
||||
),
|
||||
)
|
||||
|
||||
def test_filename_confidence_prioritizes_family_and_size_over_precision(self):
|
||||
target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors"
|
||||
|
||||
self.assertEqual(
|
||||
94.0,
|
||||
calculate_filename_confidence(
|
||||
target, "qwen3vl-4b-abliterated_int8.safetensors"
|
||||
),
|
||||
)
|
||||
self.assertGreaterEqual(
|
||||
calculate_filename_confidence(target, "qwen3vl_4b_int8.safetensors"),
|
||||
90.0,
|
||||
)
|
||||
self.assertGreaterEqual(
|
||||
calculate_filename_confidence(
|
||||
target, "qwen3vl-4b-heretic_int8.safetensors"
|
||||
),
|
||||
84.0,
|
||||
)
|
||||
self.assertLess(
|
||||
calculate_filename_confidence(
|
||||
target, "qwen_image_fp8_e4m3fn.safetensors"
|
||||
),
|
||||
calculate_filename_confidence(
|
||||
target, "qwen3vl-4b-heretic_int8.safetensors"
|
||||
),
|
||||
)
|
||||
|
||||
def test_filename_confidence_does_not_merge_different_sizes_or_generations(self):
|
||||
target = "qwen3vl-4b_fp8_e4m3fn.safetensors"
|
||||
|
||||
self.assertLess(
|
||||
calculate_filename_confidence(target, "qwen3vl-8b_int8.safetensors"),
|
||||
70.0,
|
||||
)
|
||||
self.assertLess(
|
||||
calculate_filename_confidence(target, "qwen2vl-4b_int8.safetensors"),
|
||||
70.0,
|
||||
)
|
||||
|
||||
def test_find_matches_includes_compatible_qwen_precision_variants(self):
|
||||
target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors"
|
||||
candidates = [
|
||||
{"filename": "qwen3vl-4b-abliterated_int8.safetensors"},
|
||||
{"filename": "qwen3vl_4b_int8.safetensors"},
|
||||
{"filename": "qwen3vl-4b-heretic_int8.safetensors"},
|
||||
{"filename": "qwen3vl-8b_int8.safetensors"},
|
||||
{"filename": "qwen_image_fp8_e4m3fn.safetensors"},
|
||||
]
|
||||
|
||||
matches = find_matches(target, candidates, threshold=0.7)
|
||||
|
||||
self.assertEqual(
|
||||
[
|
||||
"qwen3vl-4b-abliterated_int8.safetensors",
|
||||
"qwen3vl_4b_int8.safetensors",
|
||||
"qwen3vl-4b-heretic_int8.safetensors",
|
||||
"qwen_image_fp8_e4m3fn.safetensors",
|
||||
],
|
||||
[match["filename"] for match in matches],
|
||||
)
|
||||
|
||||
def test_search_model_list_unification(self):
|
||||
try:
|
||||
search_model_list("nonexistent_model_name_xyz.safetensors", exact_only=True)
|
||||
@@ -50,6 +182,45 @@ class MatcherTests(unittest.TestCase):
|
||||
self.assertIn("model.ckpt", queries_simple)
|
||||
self.assertIn("model", queries_simple)
|
||||
|
||||
queries_int8 = build_filename_search_queries(
|
||||
"qwen3vl_4b_abliterated_int8.safetensors"
|
||||
)
|
||||
self.assertIn("qwen3vl_4b_abliterated", queries_int8)
|
||||
|
||||
for filename in [
|
||||
"qwen3vl_4b_q4_k_m.gguf",
|
||||
"qwen3vl_4b_iq2_xxs.gguf",
|
||||
"qwen3vl_4b_e4m3fnuz.safetensors",
|
||||
"qwen3vl_4b_exl2.safetensors",
|
||||
]:
|
||||
with self.subTest(filename=filename):
|
||||
self.assertIn("qwen3vl_4b", build_filename_search_queries(filename))
|
||||
|
||||
def test_archived_model_confidence_uses_model_family_variants(self):
|
||||
target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors"
|
||||
|
||||
self.assertEqual(
|
||||
94.0,
|
||||
calculate_archived_model_confidence(
|
||||
target,
|
||||
filename="qwen3vl-4b-abliterated_int8.safetensors",
|
||||
),
|
||||
)
|
||||
self.assertGreaterEqual(
|
||||
calculate_archived_model_confidence(
|
||||
target,
|
||||
filename="qwen3vl-4b-heretic_int8.safetensors",
|
||||
),
|
||||
84.0,
|
||||
)
|
||||
self.assertLess(
|
||||
calculate_archived_model_confidence(
|
||||
target,
|
||||
filename="qwen3vl-8b_int8.safetensors",
|
||||
),
|
||||
70.0,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -95,6 +266,28 @@ class ProgressTests(unittest.TestCase):
|
||||
|
||||
|
||||
class TypeUtilsTests(unittest.TestCase):
|
||||
def test_technical_suffix_categories_are_disjoint_and_combined(self):
|
||||
precision = set(NUMERIC_PRECISION_SUFFIXES)
|
||||
quantization_levels = set(QUANTIZATION_LEVEL_SUFFIXES)
|
||||
quantization_schemes = set(QUANTIZATION_SCHEME_SUFFIXES)
|
||||
containers = set(MODEL_CONTAINER_SUFFIXES)
|
||||
variants = set(MODEL_VARIANT_SUFFIXES)
|
||||
categories = [
|
||||
precision,
|
||||
quantization_levels,
|
||||
quantization_schemes,
|
||||
containers,
|
||||
variants,
|
||||
]
|
||||
|
||||
for index, category in enumerate(categories):
|
||||
for other_category in categories[index + 1 :]:
|
||||
self.assertTrue(category.isdisjoint(other_category))
|
||||
self.assertEqual(
|
||||
set().union(*categories),
|
||||
set(TECHNICAL_MODEL_SUFFIXES),
|
||||
)
|
||||
|
||||
def test_as_dict_returns_dict_when_dict(self):
|
||||
d = {"a": 1}
|
||||
self.assertEqual(d, as_dict(d))
|
||||
|
||||
Reference in New Issue
Block a user