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:
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user