Download tests/test_hub.py from AEUPH/hdcstorage: direct link, hf CLI and curl.
- Browser
- Download file 16.9 kB
-
https://huggingface.co/AEUPH/hdcstorage/resolve/main/tests/test_hub.py
- Command line
-
hf download hf://AEUPH/hdcstorage/tests/test_hub.py
-
curl -L -o test_hub.py https://huggingface.co/AEUPH/hdcstorage/resolve/main/tests/test_hub.py
16.9 kB
| """Tests for model discovery: the Hub search client and its measurements. | |
| Nothing here touches the network. Every function that would make a request is | |
| monkeypatched, and the safetensors fixtures are packed in the test — the same | |
| approach ``tests/test_semantic.py`` uses for the projection tests. | |
| The dtype cases matter: the code-tuned sibling this project settled on, | |
| ``minishlab/potion-code-16M-v2``, is stored as **F16** (63 457 x 256), while the | |
| general ``potion-base-8M`` is F32. Reading an F16 file with the F32 path would | |
| silently return garbage rows, so the reader is pinned here. | |
| """ | |
| import json | |
| import os | |
| import struct | |
| import tempfile | |
| import unittest | |
| import urllib.parse | |
| from unittest import mock | |
| from hdc import hub | |
| from hdc.semantic import StaticEmbeddings | |
| RAW_ONE = { | |
| "id": "minishlab/potion-base-8M", | |
| "downloads": 447_769, | |
| "likes": 12, | |
| "library_name": "model2vec", | |
| "pipeline_tag": "feature-extraction", | |
| "gated": False, | |
| "tags": ["model2vec", "license:mit", "feature-extraction"], | |
| "siblings": [{"rfilename": "model.safetensors"}, {"rfilename": "tokenizer.json"}], | |
| "sha": "abc123", | |
| } | |
| RAW_TWO = { | |
| "id": "jinaai/jina-embeddings-v2-base-code", | |
| "downloads": 156_469, | |
| "likes": 40, | |
| "library_name": "sentence-transformers", | |
| "pipeline_tag": "feature-extraction", | |
| "gated": True, | |
| "tags": ["license:apache-2.0", "sentence-transformers"], | |
| "siblings": [{"rfilename": "pytorch_model.bin"}], | |
| "sha": "def456", | |
| } | |
| DTYPE_BYTES = {"F32": 4, "F16": 2, "BF16": 2, "I8": 1} | |
| def header(tensors: dict) -> dict: | |
| """A parsed safetensors header keyed the way ``hub.safetensors`` returns it. | |
| ``bytes`` is derived from shape x dtype, because that is exactly what the | |
| real header's ``data_offsets`` gives without reading the tensor buffer. | |
| """ | |
| return {"repo": "x", "tensors": tensors, | |
| "bytes": sum(hub._shape_product(t["shape"]) * DTYPE_BYTES.get(t["dtype"], 4) | |
| for t in tensors.values()), | |
| "params": sum(hub._shape_product(t["shape"]) for t in tensors.values())} | |
| def write_static(directory: str, dtype: str, rows: list[list[float]], | |
| vocab: dict | None = None) -> str: | |
| """Pack a tiny safetensors file in ``dtype`` plus a tokenizer, no libraries.""" | |
| flat = [value for row in rows for value in row] | |
| if dtype == "F16": | |
| payload = struct.pack("<%de" % len(flat), *flat) | |
| elif dtype == "BF16": | |
| # bf16 is the top half of the f32 word | |
| halves = [struct.unpack("<I", struct.pack("<f", value))[0] >> 16 for value in flat] | |
| payload = struct.pack("<%dH" % len(flat), *halves) | |
| else: | |
| payload = struct.pack("<%df" % len(flat), *flat) | |
| meta = json.dumps({ | |
| "embeddings": {"dtype": dtype, "shape": [len(rows), len(rows[0])], | |
| "data_offsets": [0, len(payload)]} | |
| }).encode("utf-8") | |
| meta += b" " * ((8 - len(meta) % 8) % 8) | |
| with open(os.path.join(directory, "model.safetensors"), "wb") as fh: | |
| fh.write(struct.pack("<Q", len(meta))) | |
| fh.write(meta) | |
| fh.write(payload) | |
| with open(os.path.join(directory, "tokenizer.json"), "w", encoding="utf-8") as fh: | |
| json.dump({"model": {"vocab": vocab or {"a": 0}, "type": "WordPiece", | |
| "continuing_subword_prefix": "##"}}, fh) | |
| return directory | |
| class TestNormalise(unittest.TestCase): | |
| def test_record_is_reduced_and_licence_is_unpacked(self): | |
| record = hub.normalise(RAW_ONE) | |
| self.assertEqual(record["id"], "minishlab/potion-base-8M") | |
| self.assertEqual(record["downloads"], 447_769) | |
| self.assertEqual(record["likes"], 12) | |
| self.assertEqual(record["library"], "model2vec") | |
| self.assertEqual(record["license"], "mit") | |
| self.assertEqual(record["files"], | |
| ["model.safetensors", "tokenizer.json"]) | |
| self.assertFalse(record["gated"]) | |
| self.assertTrue(hub.normalise(RAW_TWO)["gated"]) | |
| def test_missing_fields_do_not_raise(self): | |
| record = hub.normalise({}) | |
| self.assertEqual(record["id"], "") | |
| self.assertEqual(record["downloads"], 0) | |
| self.assertEqual(record["license"], "") | |
| class TestStaticTable(unittest.TestCase): | |
| def test_single_rank_two_tensor_is_static(self): | |
| is_static, rows, dim = hub.static_table(header({"embeddings": {"dtype": "F16", | |
| "shape": [63457, 256]}})) | |
| self.assertTrue(is_static) | |
| self.assertEqual((rows, dim), (63457, 256)) | |
| def test_many_tensors_are_a_transformer(self): | |
| many = {name: {"dtype": "F32", "shape": [8, 8]} | |
| for name in ("q_proj", "k_proj", "v_proj")} | |
| self.assertEqual(hub.static_table(header(many)), (False, 0, 0)) | |
| def test_rank_one_and_rank_three_are_rejected(self): | |
| rank_one = header({"embeddings": {"dtype": "F32", "shape": [256]}}) | |
| rank_three = header({"embeddings": {"dtype": "F32", "shape": [2, 4, 8]}}) | |
| self.assertEqual(hub.static_table(rank_one), (False, 0, 0)) | |
| self.assertEqual(hub.static_table(rank_three), (False, 0, 0)) | |
| def test_no_header_is_not_static(self): | |
| self.assertEqual(hub.static_table(None), (False, 0, 0)) | |
| class TestCodeCoverage(unittest.TestCase): | |
| def test_full_vocabulary_covers_everything_whole(self): | |
| vocab = {token: i for i, token in enumerate(hub.MACHINE_CODE)} | |
| coverage = hub.code_coverage(vocab) | |
| self.assertEqual(coverage["coverage"], 1.0) | |
| self.assertEqual(coverage["whole_fraction"], 1.0) | |
| self.assertEqual(coverage["any_fraction"], 1.0) | |
| self.assertEqual(coverage["missing"], []) | |
| def test_fragments_are_worth_less_than_whole_tokens(self): | |
| # every token is reachable, but only as first-char + "##" remainder | |
| vocab = {} | |
| for token in hub.MACHINE_CODE: | |
| vocab[token[0]] = len(vocab) | |
| vocab["##" + token[1:]] = len(vocab) | |
| coverage = hub.code_coverage(vocab) | |
| self.assertEqual(coverage["whole_fraction"], 0.0) | |
| self.assertEqual(coverage["any_fraction"], 1.0) | |
| # the 0.7 / 0.3 weighting is visible in the number itself | |
| self.assertAlmostEqual(coverage["coverage"], 0.3 * coverage["any_fraction"], places=6) | |
| self.assertLess(coverage["coverage"], 0.5) | |
| def test_token_with_no_pieces_is_missing(self): | |
| coverage = hub.code_coverage({}, tokens=["qqq", "mov"]) | |
| self.assertEqual(coverage["missing"], ["qqq", "mov"]) | |
| self.assertEqual(coverage["coverage"], 0.0) | |
| class TestScore(unittest.TestCase): | |
| def setUp(self): | |
| self.coverage = hub.code_coverage({token: i for i, token in enumerate(hub.MACHINE_CODE)}) | |
| def test_components_match_the_weights(self): | |
| graded = hub.score(hub.normalise(RAW_ONE), | |
| header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}), | |
| self.coverage) | |
| self.assertEqual(set(graded["components"]), set(hub.WEIGHTS)) | |
| def test_weighted_sum_is_the_total(self): | |
| graded = hub.score(hub.normalise(RAW_ONE), | |
| header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}), | |
| self.coverage) | |
| expected = sum(hub.WEIGHTS[name] * value | |
| for name, value in graded["components"].items()) | |
| self.assertAlmostEqual(graded["total"], expected, places=2) | |
| def test_transformer_scores_zero_on_runtime_free(self): | |
| graded = hub.score(hub.normalise(RAW_TWO), None, self.coverage) | |
| self.assertEqual(graded["components"]["runtime_free"], 0.0) | |
| self.assertFalse(graded["static"]) | |
| def test_small_static_beats_large_transformer_at_equal_coverage(self): | |
| small = hub.score(hub.normalise(RAW_ONE), | |
| header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}), | |
| self.coverage) | |
| large = hub.score(hub.normalise(RAW_TWO), | |
| header({"model.embed_tokens": {"dtype": "F16", "shape": [250000, 1024]}, | |
| "layers.0.q_proj": {"dtype": "F16", "shape": [1024, 1024]}}), | |
| self.coverage) | |
| self.assertGreater(large["bytes"], 400_000_000) | |
| self.assertGreater(small["total"], large["total"]) | |
| class TestClamps(unittest.TestCase): | |
| def test_size_score_is_clamped_and_monotonic(self): | |
| self.assertEqual(hub._size_score(30_000_000), 1.0) | |
| self.assertEqual(hub._size_score(600_000_000), 0.0) | |
| self.assertEqual(hub._size_score(0), 0.0) | |
| self.assertGreater(hub._size_score(60_000_000), hub._size_score(200_000_000)) | |
| def test_dim_score_prefers_256_and_clamps(self): | |
| self.assertEqual(hub._dim_score(256), 1.0) | |
| self.assertEqual(hub._dim_score(0), 0.0) | |
| self.assertLess(hub._dim_score(128), hub._dim_score(256)) | |
| self.assertGreaterEqual(hub._dim_score(2048), 0.0) | |
| class TestSearch(unittest.TestCase): | |
| def setUp(self): | |
| self.calls: list[str] = [] | |
| def fake_get_json(self, url, timeout=45): | |
| self.calls.append(url) | |
| return [RAW_ONE, RAW_TWO] | |
| def query(self) -> dict: | |
| parsed = urllib.parse.urlparse(self.calls[-1]) | |
| return urllib.parse.parse_qs(parsed.query) | |
| def test_text_query_keeps_the_task_filter(self): | |
| with mock.patch.object(hub, "_get_json", self.fake_get_json): | |
| records = hub.search("code embeddings") | |
| self.assertEqual(len(records), 2) | |
| self.assertEqual(records[0]["id"], RAW_ONE["id"]) | |
| query = self.query() | |
| self.assertEqual(query["search"], ["code embeddings"]) | |
| self.assertEqual(query["filter"], ["feature-extraction"]) | |
| def test_author_query_drops_the_filter(self): | |
| # measured API quirk: author + filter returns zero rows, silently | |
| with mock.patch.object(hub, "_get_json", self.fake_get_json): | |
| hub.search(author="minishlab") | |
| query = self.query() | |
| self.assertNotIn("filter", query) | |
| self.assertEqual(query["author"], ["minishlab"]) | |
| def test_non_list_payload_is_an_empty_result(self): | |
| with mock.patch.object(hub, "_get_json", lambda url, timeout=45: {"error": "nope"}): | |
| self.assertEqual(hub.search("anything"), []) | |
| class TestCandidatePool(unittest.TestCase): | |
| def test_queries_are_de_duplicated_and_ordered_by_downloads(self): | |
| shared = hub.normalise(dict(RAW_ONE, id="a/shared", downloads=5)) | |
| more = hub.normalise(dict(RAW_ONE, id="b/bigger", downloads=9)) | |
| least = hub.normalise(dict(RAW_ONE, id="c/small", downloads=1)) | |
| def fake_search(text=None, **kwargs): | |
| return [shared, more, least, more] # the same repo from two queries | |
| with mock.patch.object(hub, "search", fake_search): | |
| pool = hub.candidate_pool(limit=10) | |
| self.assertEqual([row["id"] for row in pool], ["b/bigger", "a/shared", "c/small"]) | |
| def test_limit_is_applied(self): | |
| rows = [hub.normalise(dict(RAW_ONE, id="r/%d" % i, downloads=i)) for i in range(5)] | |
| with mock.patch.object(hub, "search", lambda *a, **k: rows): | |
| self.assertEqual(len(hub.candidate_pool(limit=2)), 2) | |
| class TestBakeoff(unittest.TestCase): | |
| def test_offline_uses_the_cache_and_touches_no_network(self): | |
| entry = {"id": "minishlab/potion-code-16M-v2", "sha": "s", "downloads": 7, | |
| "components": {"runtime_free": 1.0}, "total": 1.0, | |
| "static": True, "rows": 63457, "dim": 256, "bytes": 32_489_984, | |
| "coverage": None, "license": "mit"} | |
| with mock.patch.object(hub, "load_cache", lambda: {"minishlab/potion-code-16M-v2@s": entry}), \ | |
| mock.patch.object(hub, "save_cache", lambda cache: None), \ | |
| mock.patch.object(hub, "candidate_pool", lambda **kwargs: []): | |
| result = hub.bakeoff(offline=True) | |
| self.assertEqual(result["pool"], 1) | |
| self.assertEqual(result["graded"][0]["id"], "minishlab/potion-code-16M-v2") | |
| self.assertTrue(result["graded"][0]["cached"]) | |
| self.assertEqual(hub.best(result)["total"], 1.0) | |
| class TestHeaderRange(unittest.TestCase): | |
| def test_two_range_requests_give_shape_size_and_params(self): | |
| meta = json.dumps({ | |
| "embeddings": {"dtype": "F16", "shape": [2000, 256], | |
| "data_offsets": [0, 1_024_000]}, | |
| }).encode("utf-8") | |
| seen: list[dict] = [] | |
| def fake_request(url, headers=None, timeout=45): | |
| seen.append(headers or {}) | |
| if (headers or {}).get("Range") == "bytes=0-7": | |
| return struct.pack("<Q", len(meta)) | |
| return meta | |
| with mock.patch.object(hub, "_request", fake_request): | |
| parsed = hub.safetensors("minishlab/potion-code-16M-v2") | |
| self.assertEqual(len(seen), 2) | |
| self.assertEqual(parsed["bytes"], 1_024_000) | |
| self.assertEqual(parsed["params"], 2000 * 256) | |
| self.assertEqual(parsed["tensors"]["embeddings"]["dtype"], "F16") | |
| self.assertEqual(hub.static_table(parsed), (True, 2000, 256)) | |
| class TestDtypes(unittest.TestCase): | |
| """F16 is what the code model actually ships; F32 is what the reader used to assume.""" | |
| ROWS = [ | |
| [1.0, 0.5, -1.5, 0.25], # every value exact in f16 and bf16 | |
| [2.0, 0.0, 1.875, -0.25], | |
| ] | |
| def build(self, dtype: str) -> StaticEmbeddings: | |
| tmp = tempfile.TemporaryDirectory() | |
| self.addCleanup(tmp.cleanup) | |
| return StaticEmbeddings(write_static(tmp.name, dtype, self.ROWS)) | |
| def test_float16_rows_decode_exactly(self): | |
| model = self.build("F16") | |
| self.assertEqual(model.dtype, "F16") | |
| self.assertEqual(model.itemsize, 2) | |
| self.assertEqual(model.row(0), (1.0, 0.5, -1.5, 0.25)) | |
| self.assertEqual(model.row(1), (2.0, 0.0, 1.875, -0.25)) | |
| def test_bfloat16_rows_decode_exactly(self): | |
| model = self.build("BF16") | |
| self.assertEqual(model.dtype, "BF16") | |
| self.assertEqual(model.row(0), (1.0, 0.5, -1.5, 0.25)) | |
| self.assertEqual(model.row(1), (2.0, 0.0, 1.875, -0.25)) | |
| def test_rows_do_not_overlap(self): | |
| # the offset is rows * dim * itemsize; a wrong stride returns row 0 twice | |
| model = self.build("F16") | |
| self.assertNotEqual(model.row(0), model.row(1)) | |
| def test_unsupported_dtype_is_rejected(self): | |
| tmp = tempfile.TemporaryDirectory() | |
| self.addCleanup(tmp.cleanup) | |
| write_static(tmp.name, "I8", self.ROWS) | |
| with self.assertRaises(ValueError): | |
| StaticEmbeddings(tmp.name) | |
| class TestCacheLifetime(unittest.TestCase): | |
| """The score cache outlives the code that wrote it. | |
| An older run of ``python -m hdc hub --offline`` raised ``KeyError`` on a | |
| coverage entry that predated a new field, so ``load_cache`` has to drop what | |
| it cannot understand instead of handing it to the report. | |
| """ | |
| def write(self, payload): | |
| tmp = tempfile.TemporaryDirectory() | |
| self.addCleanup(tmp.cleanup) | |
| path = os.path.join(tmp.name, "scores.json") | |
| with open(path, "w", encoding="utf-8") as fh: | |
| json.dump(payload, fh) | |
| return path | |
| def test_current_entries_are_kept(self): | |
| entry = {"id": "a/b", "coverage": {"any_fraction": 1.0, "whole_fraction": 0.5}} | |
| path = self.write({"a/b@sha": entry}) | |
| with mock.patch.object(hub, "cache_path", return_value=path): | |
| self.assertIn("a/b@sha", hub.load_cache()) | |
| def test_stale_entries_are_dropped(self): | |
| stale = {"id": "a/b", "coverage": {"coverage": 1.0, "missing": []}} | |
| fresh = {"id": "c/d", "coverage": {"any_fraction": 1.0, "whole_fraction": 0.4}} | |
| path = self.write({"a/b@old": stale, "c/d@new": fresh}) | |
| with mock.patch.object(hub, "cache_path", return_value=path): | |
| cache = hub.load_cache() | |
| self.assertEqual(sorted(cache), ["c/d@new"]) | |
| def test_missing_or_corrupt_cache_is_empty(self): | |
| with mock.patch.object(hub, "cache_path", | |
| return_value=os.path.join("does", "not", "exist.json")): | |
| self.assertEqual(hub.load_cache(), {}) | |
| def test_report_survives_a_coverage_without_any_fraction(self): | |
| # older cached rows reach report() when offline; it must not raise | |
| entry = {"id": "a/b", "total": 0.9, "static": True, "dim": 256, | |
| "rows": 10, "bytes": 30_000_000, "license": "mit", | |
| "components": {"code": 0.5}, | |
| "coverage": {"coverage": 1.0, "missing": [], "tokens": 42}} | |
| with mock.patch.object(hub, "bakeoff", return_value={"pool": 1, "graded": [entry]}): | |
| self.assertIn("a/b", hub.report()) | |
| if __name__ == "__main__": | |
| unittest.main() | |