diff --git a/core/matcher.py b/core/matcher.py index 695ae30..42c17ef 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -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]: diff --git a/core/type_utils.py b/core/type_utils.py index a0a83c0..f12c63f 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -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) diff --git a/tests/test_matcher.py b/tests/test_matcher.py index 0383d6e..51f64f8 100644 --- a/tests/test_matcher.py +++ b/tests/test_matcher.py @@ -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))