Improve model matching by normalizing technical suffixes and family signatures

This commit is contained in:
Dariusz L
2026-07-20 14:15:41 +02:00
parent f0889a6c7f
commit 7ee358b2c5
3 changed files with 435 additions and 21 deletions
+150 -11
View File
@@ -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
View File
@@ -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
View File
@@ -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))