375 lines
13 KiB
Python
375 lines
13 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
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.progress import report_progress
|
|
from core.sources.model_list import search_model_list, search_model_list_multiple
|
|
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"))
|
|
self.assertEqual("test model", normalize_filename("test-model.ckpt"))
|
|
|
|
def test_normalize_filename_replaces_dots_and_separators(self):
|
|
self.assertEqual("test model v1 0", normalize_filename("test.model.v1.0.safetensors"))
|
|
self.assertEqual("my cool model 1 2 3", normalize_filename("My_Cool-Model.1.2.3.ckpt"))
|
|
|
|
def test_calculate_similarity_identical(self):
|
|
self.assertEqual(1.0, calculate_similarity("hello", "hello"))
|
|
|
|
def test_calculate_similarity_different(self):
|
|
self.assertLess(calculate_similarity("hello", "world"), 0.5)
|
|
|
|
def test_find_matches_exact_and_fuzzy(self):
|
|
candidates = [
|
|
{"filename": "sd_xl_base_1.0.safetensors"},
|
|
{"filename": "pony Diffusion.safetensors"},
|
|
]
|
|
matches = find_matches("sd_xl_base_1.0.safetensors", candidates)
|
|
self.assertEqual(2, len(matches))
|
|
self.assertEqual("sd_xl_base_1.0.safetensors", matches[0]["filename"])
|
|
self.assertEqual(1.0, matches[0]["similarity"])
|
|
|
|
# With higher threshold, only the best match should be returned
|
|
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)
|
|
search_model_list_multiple("nonexistent_model_name_xyz.safetensors")
|
|
except Exception as e:
|
|
self.fail(f"search_model_list failed with exception: {e}")
|
|
|
|
def test_model_list_result_fields_preserve_match_mode_contracts(self):
|
|
models = [
|
|
{
|
|
"filename": "exact.safetensors",
|
|
"url": "https://example.test/exact",
|
|
"name": "Exact model",
|
|
"type": "checkpoints",
|
|
"save_path": "checkpoints",
|
|
"size": "12 MB",
|
|
},
|
|
{
|
|
"filename": "target-model.safetensors",
|
|
"url": "https://example.test/target",
|
|
"name": "Target model",
|
|
"type": "loras",
|
|
"save_path": "loras",
|
|
"size": 24,
|
|
},
|
|
]
|
|
|
|
with patch("core.sources.model_list._load_model_list", return_value=models):
|
|
exact = search_model_list("exact.safetensors", exact_only=True)
|
|
fuzzy = search_model_list("target-model-extra.safetensors")
|
|
multiple = search_model_list_multiple("exact", limit=1)
|
|
|
|
self.assertEqual(
|
|
exact,
|
|
{
|
|
"source": "model_list",
|
|
"filename": "exact.safetensors",
|
|
"url": "https://example.test/exact",
|
|
"name": "Exact model",
|
|
"type": "checkpoints",
|
|
"directory": "checkpoints",
|
|
"size": "12 MB",
|
|
"match_type": "exact",
|
|
},
|
|
)
|
|
self.assertEqual(fuzzy["filename"], "target-model.safetensors")
|
|
self.assertEqual(fuzzy["match_type"], "fuzzy")
|
|
self.assertIn("confidence", fuzzy)
|
|
self.assertEqual(multiple[0]["filename"], "exact.safetensors")
|
|
self.assertIn("confidence", multiple[0])
|
|
self.assertNotIn("match_type", multiple[0])
|
|
|
|
def test_build_filename_search_queries(self):
|
|
# Verify that common precisions are extracted correctly
|
|
queries = build_filename_search_queries("sdxl_ep10_fp16.safetensors")
|
|
self.assertIn("sdxl_ep10_fp16.safetensors", queries)
|
|
self.assertIn("sdxl_ep10_fp16", queries)
|
|
self.assertIn("sdxl_ep10", queries)
|
|
|
|
queries_simple = build_filename_search_queries("model.ckpt")
|
|
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,
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def test_report_progress_calls_callback(self):
|
|
called_payload = None
|
|
|
|
def callback(payload):
|
|
nonlocal called_payload
|
|
called_payload = payload
|
|
|
|
report_progress(
|
|
callback,
|
|
stage="test_stage",
|
|
message="test_message",
|
|
percent=50.0,
|
|
foo="bar"
|
|
)
|
|
|
|
self.assertIsNotNone(called_payload)
|
|
self.assertEqual("test_stage", called_payload.get("stage"))
|
|
self.assertEqual("test_message", called_payload.get("message"))
|
|
self.assertEqual(50.0, called_payload.get("percent"))
|
|
self.assertEqual("bar", called_payload.get("foo"))
|
|
|
|
def test_report_progress_handles_exceptions(self):
|
|
def bad_callback(payload):
|
|
raise ValueError("Boom")
|
|
|
|
# Should not raise exception
|
|
try:
|
|
report_progress(bad_callback, "stage", "msg")
|
|
except Exception as e:
|
|
self.fail(f"report_progress raised exception: {e}")
|
|
|
|
def test_report_progress_handles_none_callback(self):
|
|
# Should do nothing and not raise exception
|
|
try:
|
|
report_progress(None, "stage", "msg")
|
|
except Exception as e:
|
|
self.fail(f"report_progress raised exception with None callback: {e}")
|
|
|
|
|
|
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))
|
|
|
|
def test_as_dict_returns_empty_dict_when_not_dict(self):
|
|
self.assertEqual({}, as_dict("not a dict"))
|
|
self.assertEqual({}, as_dict(None))
|
|
|
|
def test_as_list_returns_filtered_list(self):
|
|
self.assertEqual([1, 2], as_list([1, None, "", 2]))
|
|
|
|
def test_as_list_handles_tuples_and_sets(self):
|
|
self.assertEqual([1, 2], as_list((1, None, "", 2)))
|
|
self.assertEqual([1, 2], as_list({1, "", 2}))
|
|
|
|
def test_as_list_handles_comma_separated_strings(self):
|
|
self.assertEqual(["a", "b", "c"], as_list("a, b,, c"))
|
|
|
|
def test_as_list_returns_empty_list_otherwise(self):
|
|
self.assertEqual([], as_list(None))
|
|
self.assertEqual([], as_list(123))
|
|
|
|
def test_format_size_bytes(self):
|
|
self.assertEqual("1.5 GB", format_size_bytes(1.5 * 1024 * 1024 * 1024))
|
|
self.assertEqual("1.5GB", format_size_bytes(1.5 * 1024 * 1024 * 1024, include_space=False))
|
|
self.assertEqual("0 B", format_size_bytes(0))
|
|
self.assertEqual("0B", format_size_bytes(0, include_space=False))
|
|
self.assertIsNone(format_size_bytes(None))
|
|
self.assertIsNone(format_size_bytes(""))
|
|
self.assertEqual("invalid_str", format_size_bytes("invalid_str"))
|
|
self.assertEqual("500.0 KB", format_size_bytes("512000"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|