From dca6bc93fbc27bac6e3d94f5aecfc1ed51c99d67 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Wed, 16 Sep 2026 19:55:09 +0200 Subject: [PATCH] Improve filename confidence and HuggingFace partial matching Use order-insensitive token comparison in calculate_filename_confidence so variant names with reordered tokens (e.g., int8_convrot vs convrot_int8) score higher. Replace first-match partial matching in HuggingFace source with best-match selection based on filename confidence, and surface confidence in search results. --- core/matcher.py | 12 ++++ core/sources/huggingface.py | 82 ++++++++++++++++---------- tests/test_huggingface.py | 114 ++++++++++++++++++++++++++++++++++++ tests/test_matcher.py | 9 +++ 4 files changed, 187 insertions(+), 30 deletions(-) diff --git a/core/matcher.py b/core/matcher.py index edbfa09..5af1755 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -425,6 +425,18 @@ def calculate_filename_confidence(target_filename: str, candidate_filename: str) return 100.0 best_similarity = calculate_similarity(target_norm, candidate_norm) + target_tokens = target_norm.split() + candidate_tokens = candidate_norm.split() + if target_tokens and candidate_tokens: + # Compare a canonical token order as well as the original spelling. + # This keeps names such as `convrot_int8` and `int8_convrot` together + # without introducing a second variant-matching implementation. + order_insensitive_similarity = calculate_similarity( + " ".join(sorted(target_tokens)), + " ".join(sorted(candidate_tokens)), + ) + best_similarity = max(best_similarity, order_insensitive_similarity) + target_family = normalize_model_family_filename(target_filename) candidate_family = normalize_model_family_filename(candidate_filename) diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 21a3f04..a1a4e7f 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -19,7 +19,10 @@ from ..progress import get_progress_reporter log = create_module_logger(__name__) -from ..matcher import build_filename_search_queries +from ..matcher import ( + build_filename_search_queries, + calculate_filename_confidence, +) from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic from ..type_utils import ( @@ -704,6 +707,7 @@ def _build_huggingface_result( file_info: Dict[str, Any], match_type: str, headers: Optional[Dict[str, str]] = None, + confidence: float = 0.0, ) -> SearchResult: sha256 = _extract_huggingface_file_sha256(file_info) download_url = get_huggingface_download_url(repo_id, file_path) @@ -721,6 +725,7 @@ def _build_huggingface_result( download_url=download_url, size=size, match_type=match_type, + confidence=confidence, sha256=sha256, normalize_hashes=True, repo_id=repo_id, @@ -794,7 +799,6 @@ def _find_matching_file_in_repo( ) -> Optional[SearchResult]: filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] - partial_match = None for file_info in files: file_path = file_info.get("path", "") @@ -803,7 +807,6 @@ def _find_matching_file_in_repo( file_name = get_filename_from_path(file_path) file_name_lower = file_name.lower() - file_base = os.path.splitext(file_name_lower)[0] if file_name_lower == filename_lower or file_path.lower().endswith(filename_lower): return _build_huggingface_result( @@ -812,11 +815,13 @@ def _find_matching_file_in_repo( file_info, "exact", headers=headers, + confidence=100.0, ) if exact_only: return None + best_match: Optional[SearchResult] = None for file_info in files: file_path = file_info.get("path", "") if not file_path: @@ -825,22 +830,28 @@ def _find_matching_file_in_repo( file_name = get_filename_from_path(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] + if not file_path.lower().endswith((".safetensors", ".ckpt")): + continue - if ( - not exact_only - and (filename_base in file_base or file_base in filename_base) - and (file_path.endswith(".safetensors") or file_path.endswith(".ckpt")) - ): - partial_match = _build_huggingface_result( - repo_id, - file_path, - file_info, - "partial", - headers=headers, - ) - break + confidence = calculate_filename_confidence(filename, file_name) + is_partial_filename = ( + filename_base in file_base or file_base in filename_base + ) + if not is_partial_filename and confidence < 70.0: + continue - return partial_match + candidate = _build_huggingface_result( + repo_id, + file_path, + file_info, + "partial", + headers=headers, + confidence=confidence, + ) + if best_match is None or candidate.confidence > best_match.confidence: + best_match = candidate + + return best_match def _find_matching_file_in_author_index( @@ -870,13 +881,13 @@ def _find_matching_file_in_author_index( file_info, "hash", headers=headers, + confidence=100.0, ) return None filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] - partial_match = None for file_info in files: repo_id = file_info.get("repo_id", "") @@ -895,11 +906,13 @@ def _find_matching_file_in_author_index( file_info, "exact", headers=headers, + confidence=100.0, ) if exact_only: return None + best_match: Optional[SearchResult] = None for file_info in files: repo_id = file_info.get("repo_id", "") file_path = file_info.get("path", "") @@ -911,19 +924,28 @@ def _find_matching_file_in_author_index( file_base = os.path.splitext(file_name_lower)[0] file_path_lower = file_path.lower() - if ( - filename_base in file_base or file_base in filename_base - ) and file_path_lower.endswith((".safetensors", ".ckpt")): - partial_match = _build_huggingface_result( - repo_id, - file_path, - file_info, - "partial", - headers=headers, - ) - break + if not file_path_lower.endswith((".safetensors", ".ckpt")): + continue - return partial_match + confidence = calculate_filename_confidence(filename, file_name) + is_partial_filename = ( + filename_base in file_base or file_base in filename_base + ) + if not is_partial_filename and confidence < 70.0: + continue + + candidate = _build_huggingface_result( + repo_id, + file_path, + file_info, + "partial", + headers=headers, + confidence=confidence, + ) + if best_match is None or candidate.confidence > best_match.confidence: + best_match = candidate + + return best_match diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py index 5acaa4e..648ccec 100644 --- a/tests/test_huggingface.py +++ b/tests/test_huggingface.py @@ -13,6 +13,7 @@ from core.sources.huggingface import ( _build_author_index_from_models, _fetch_author_index, _find_matching_file_in_author_index, + _find_matching_file_in_repo, _get_author_index, _is_author_index_fresh, _read_persistent_author_indexes, @@ -172,6 +173,119 @@ class HuggingFaceSourceTests(unittest.TestCase): self.assertEqual("hash", result["match_type"]) self.assertEqual(sha256, result["sha256"]) self.assertEqual("models/example.safetensors", result["path"]) + self.assertEqual(100.0, result["confidence"]) + + def test_huggingface_matches_reordered_variant_with_confidence(self): + files = [ + { + "path": "checkpoints/yue2_3b_bf16.safetensors", + "size": 100, + }, + { + "path": "checkpoints/yue2_3b_int8_convrot.safetensors", + "size": 100, + }, + ] + + repo_result = _find_matching_file_in_repo( + "Comfy-Org/YuE2", + files, + "yue2_convrot_int8.safetensors", + ) + index_result = _find_matching_file_in_author_index( + { + "files": [ + { + "repo_id": "Comfy-Org/YuE2", + "path": file["path"], + "size": file["size"], + } + for file in files + ] + }, + "yue2_convrot_int8.safetensors", + ) + + for result in (repo_result, index_result): + self.assertIsNotNone(result) + result = _result_dict(result) + self.assertEqual( + "checkpoints/yue2_3b_int8_convrot.safetensors", result["path"] + ) + self.assertEqual("partial", result["match_type"]) + self.assertGreater(result["confidence"], 0.0) + self.assertLess(result["confidence"], 100.0) + + def test_huggingface_variant_match_stays_strict_for_exact_only(self): + result = _find_matching_file_in_repo( + "Comfy-Org/YuE2", + [ + { + "path": "checkpoints/yue2_3b_int8_convrot.safetensors", + "size": 100, + } + ], + "yue2_convrot_int8.safetensors", + exact_only=True, + ) + + self.assertIsNone(result) + + def test_huggingface_variant_match_keeps_precision_alternatives(self): + result = _find_matching_file_in_author_index( + { + "files": [ + { + "repo_id": "Comfy-Org/YuE2", + "path": "checkpoints/yue2_3b_bf16.safetensors", + "size": 100, + } + ] + }, + "yue2_convrot_int8.safetensors", + ) + + self.assertIsNotNone(result) + result = _result_dict(result) + self.assertEqual( + "checkpoints/yue2_3b_bf16.safetensors", result["path"] + ) + self.assertGreater(result["confidence"], 0.0) + + def test_file_search_returns_variant_match_from_author_index(self): + index = { + "files": [ + { + "repo_id": "Comfy-Org/YuE2", + "path": "checkpoints/yue2_3b_int8_convrot.safetensors", + "size": 100, + } + ] + } + + clear_search_cache() + with ( + patch( + "core.sources.huggingface._get_author_index", + return_value=index, + ) as get_index, + patch("core.sources.huggingface.execute_provider_json_request") as request, + ): + result = search_huggingface_for_file( + "yue2_convrot_int8.safetensors", + use_api_search=False, + use_brave_fallback=False, + ) + + self.assertIsNotNone(result) + result = _result_dict(result) + self.assertEqual("Comfy-Org/YuE2", result["repo_id"]) + self.assertEqual( + "checkpoints/yue2_3b_int8_convrot.safetensors", result["path"] + ) + self.assertGreater(result["confidence"], 0.0) + get_index.assert_called_once() + request.assert_not_called() def test_hash_search_uses_author_index_before_filename_search(self): sha256 = "c" * 64 diff --git a/tests/test_matcher.py b/tests/test_matcher.py index afed5e8..f526daa 100644 --- a/tests/test_matcher.py +++ b/tests/test_matcher.py @@ -176,6 +176,15 @@ class MatcherTests(unittest.TestCase): ), ) + def test_filename_confidence_ignores_token_order(self): + target = "yue2_convrot_int8.safetensors" + + confidence = calculate_filename_confidence( + target, "yue2_3b_int8_convrot.safetensors" + ) + + self.assertGreaterEqual(confidence, 90.0) + def test_filename_confidence_does_not_merge_different_sizes_or_generations(self): target = "qwen3vl-4b_fp8_e4m3fn.safetensors"