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.
This commit is contained in:
Dariusz L
2026-09-16 19:55:09 +02:00
parent 7fbcd6b6ee
commit dca6bc93fb
4 changed files with 187 additions and 30 deletions
+12
View File
@@ -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)
+52 -30
View File
@@ -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
+114
View File
@@ -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
+9
View File
@@ -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"