Spaces:
Sleeping
Sleeping
| """Offline importer tests: generated checkpoints and a stub tokenizer only.""" | |
| import json | |
| import struct | |
| import sys | |
| from types import SimpleNamespace | |
| import numpy as np | |
| import pytest | |
| from safetensors.numpy import save_file | |
| from cism import loader | |
| def tiny_config(model_type="llama", **overrides): | |
| config = { | |
| "model_type": model_type, | |
| "architectures": ["LlamaForCausalLM" if model_type == "llama" else "Qwen3ForCausalLM"], | |
| "vocab_size": 7, | |
| "hidden_size": 8, | |
| "intermediate_size": 12, | |
| "num_hidden_layers": 1, | |
| "num_attention_heads": 2, | |
| "num_key_value_heads": 1, | |
| "head_dim": 4, | |
| "max_position_embeddings": 32, | |
| } | |
| config.update(overrides) | |
| return config | |
| def tiny_weights(config, dtype=np.float32): | |
| # Independent shape fixture, so a broken production shape table cannot pass. | |
| h, m, d = config["hidden_size"], config["intermediate_size"], config["head_dim"] | |
| q, kv = config["num_attention_heads"] * d, config["num_key_value_heads"] * d | |
| shapes = { | |
| "model.embed_tokens.weight": (config["vocab_size"], h), | |
| "model.norm.weight": (h,), | |
| "lm_head.weight": (config["vocab_size"], h), | |
| } | |
| for i in range(config["num_hidden_layers"]): | |
| for name, shape in { | |
| "input_layernorm.weight": (h,), "post_attention_layernorm.weight": (h,), | |
| "self_attn.q_proj.weight": (q, h), "self_attn.k_proj.weight": (kv, h), | |
| "self_attn.v_proj.weight": (kv, h), "self_attn.o_proj.weight": (h, q), | |
| "mlp.gate_proj.weight": (m, h), "mlp.up_proj.weight": (m, h), | |
| "mlp.down_proj.weight": (h, m), | |
| }.items(): | |
| shapes[f"model.layers.{i}.{name}"] = shape | |
| if config["model_type"] == "qwen3": | |
| for name in ("q_norm", "k_norm"): | |
| shapes[f"model.layers.{i}.self_attn.{name}.weight"] = (d,) | |
| return {name: (np.arange(np.prod(shape)).reshape(shape) / 16 - 1).astype(dtype) | |
| for name, shape in shapes.items()} | |
| def write_checkpoint(path, config=None, weights=None, sharded=False): | |
| config = tiny_config() if config is None else config | |
| weights = tiny_weights(config) if weights is None else weights | |
| (path / "config.json").write_text(json.dumps(config), encoding="utf-8") | |
| if not sharded: | |
| save_file(weights, path / "model.safetensors") | |
| else: | |
| names = list(weights) | |
| weight_map = {} | |
| for i, group in enumerate((names[::2], names[1::2])): | |
| filename = f"model-{i + 1:05d}-of-00002.safetensors" | |
| save_file({name: weights[name] for name in group}, path / filename) | |
| weight_map.update({name: filename for name in group}) | |
| (path / "model.safetensors.index.json").write_text( | |
| json.dumps({"weight_map": weight_map}), encoding="utf-8" | |
| ) | |
| return weights | |
| def offline_tokenizer(monkeypatch): | |
| calls = [] | |
| token = object() | |
| def from_pretrained(*args, **kwargs): | |
| calls.append((args, kwargs)) | |
| return token | |
| monkeypatch.setitem(sys.modules, "transformers", SimpleNamespace( | |
| AutoTokenizer=SimpleNamespace(from_pretrained=from_pretrained) | |
| )) | |
| def no_network(**kwargs): | |
| raise AssertionError("Unexpected network access") | |
| monkeypatch.setattr(loader, "snapshot_download", no_network) | |
| return calls, token | |
| def test_import_local(tmp_path, offline_tokenizer, model_type, head_dim, sharded, dtype): | |
| config = tiny_config(model_type, head_dim=head_dim) | |
| weights = tiny_weights(config, dtype) | |
| write_checkpoint(tmp_path, config, weights, sharded) | |
| loaded = loader.import_model(tmp_path) | |
| assert isinstance(loaded, loader.LoadedModel) | |
| assert loaded.source == str(tmp_path.resolve()) | |
| assert loaded.revision is None | |
| assert loaded.config["head_dim"] == head_dim | |
| assert loaded.config["rope_theta"] == 10000.0 | |
| assert loaded.tokenizer is offline_tokenizer[1] | |
| assert offline_tokenizer[0] == [((str(tmp_path.resolve()),), { | |
| "trust_remote_code": False, "local_files_only": True, | |
| })] | |
| assert set(loaded.weights) == set(weights) | |
| for name, value in loaded.weights.items(): | |
| assert value.dtype == np.float32 | |
| assert value.flags.c_contiguous | |
| np.testing.assert_array_equal(value, weights[name].astype(np.float32)) | |
| def test_rejected_config(overrides, match): | |
| with pytest.raises(loader.UnsupportedModelError, match=match): | |
| loader.validate_config(tiny_config(**overrides)) | |
| def test_default_rope_normalization(rope): | |
| original = tiny_config(**rope) | |
| normalized = loader.validate_config(original) | |
| theta = rope.get("rope_theta", rope.get("rope_parameters", {}).get("rope_theta", 10000)) | |
| assert normalized["rope_theta"] == theta | |
| assert normalized["rope_scaling"] is None | |
| assert normalized["rope_parameters"] == {"rope_type": "default"} | |
| assert loader.validate_config(normalized) == normalized | |
| assert original == tiny_config(**rope) | |
| def test_optional_hf_defaults_and_old_qwen_window(): | |
| config = tiny_config() | |
| del config["head_dim"] | |
| del config["num_key_value_heads"] | |
| normalized = loader.validate_config(config) | |
| assert normalized["head_dim"] == 4 | |
| assert normalized["num_key_value_heads"] == 2 | |
| with pytest.raises(loader.UnsupportedModelError, match="divide evenly"): | |
| loader.validate_config({**config, "hidden_size": 7}) | |
| qwen = loader.validate_config(tiny_config("qwen3", use_sliding_window=False, sliding_window=4096)) | |
| assert qwen["sliding_window"] is None | |
| def test_tied_head_alias(tmp_path, present): | |
| config = tiny_config(tie_word_embeddings=True) | |
| weights = tiny_weights(config) | |
| if not present: | |
| del weights["lm_head.weight"] | |
| write_checkpoint(tmp_path, config, weights, sharded=True) | |
| loaded = loader.import_model(tmp_path) | |
| assert loaded.weights["lm_head.weight"] is loaded.weights["model.embed_tokens.weight"] | |
| def test_bad_weights(tmp_path, offline_tokenizer, problem, match): | |
| config = tiny_config(tie_word_embeddings=problem == "tied") | |
| weights = tiny_weights(config) | |
| name = "model.layers.0.self_attn.q_proj.weight" | |
| if problem == "shape": | |
| weights[name] = np.zeros((1, 1), dtype=np.float32) | |
| elif problem == "missing": | |
| del weights[name] | |
| elif problem == "extra": | |
| weights["model.custom_scale.weight"] = np.ones(8, dtype=np.float32) | |
| elif problem == "bias": | |
| weights["model.layers.0.self_attn.q_proj.bias"] = np.ones(8, dtype=np.float32) | |
| elif problem in ("nan", "inf"): | |
| weights[name][0, 0] = float(problem) | |
| elif problem == "overflow": | |
| weights[name] = np.full(weights[name].shape, 1e100, dtype=np.float64) | |
| elif problem == "integer": | |
| weights[name] = weights[name].astype(np.int32) | |
| elif problem == "tied": | |
| weights["lm_head.weight"][0, 0] += 1 | |
| write_checkpoint(tmp_path, config, weights) | |
| with pytest.raises(loader.InvalidCheckpointError, match=match): | |
| loader.import_model(tmp_path) | |
| assert not offline_tokenizer[0] | |
| def test_documented_rotary_buffers(tmp_path, valid): | |
| config = tiny_config() | |
| weights = tiny_weights(config) | |
| weights["model.rotary_emb.inv_freq"] = np.array([1, 0.01 if valid else 0.5], dtype=np.float32) | |
| write_checkpoint(tmp_path, config, weights) | |
| if valid: | |
| assert "model.rotary_emb.inv_freq" not in loader.import_model(tmp_path).weights | |
| else: | |
| with pytest.raises(loader.InvalidCheckpointError, match="does not match default RoPE"): | |
| loader.import_model(tmp_path) | |
| def write_bf16(path, weights): | |
| """Generate real BF16 safetensors without requiring torch or ml_dtypes.""" | |
| header, chunks, offset = {}, [], 0 | |
| for name, value in weights.items(): | |
| bits = (value.astype(np.float32).view(np.uint32) >> 16).astype("<u2") | |
| chunk = bits.tobytes() | |
| header[name] = {"dtype": "BF16", "shape": list(value.shape), | |
| "data_offsets": [offset, offset + len(chunk)]} | |
| chunks.append(chunk) | |
| offset += len(chunk) | |
| encoded = json.dumps(header).encode("utf-8") | |
| encoded += b" " * (-len(encoded) % 8) | |
| path.write_bytes(struct.pack("<Q", len(encoded)) + encoded + b"".join(chunks)) | |
| def test_real_bf16_without_torch(tmp_path, monkeypatch): | |
| config = tiny_config("qwen3", head_dim=6) | |
| weights = tiny_weights(config) | |
| weights["model.norm.weight"] = np.array([1, -2, 0, -0.0, 1e30, 1e-30, 3.14, -7.7], dtype=np.float32) | |
| (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") | |
| write_bf16(tmp_path / "model.safetensors", weights) | |
| monkeypatch.setitem(sys.modules, "torch", None) | |
| loaded = loader.import_model(tmp_path) | |
| for name, value in weights.items(): | |
| expected_bits = value.view(np.uint32) & np.uint32(0xFFFF0000) | |
| actual = loaded.weights[name] | |
| assert actual.flags.c_contiguous | |
| assert actual.dtype == np.float32 | |
| np.testing.assert_array_equal(actual.view(np.uint32), expected_bits) | |
| def test_nonfinite_bf16(tmp_path, value): | |
| config = tiny_config() | |
| weights = tiny_weights(config) | |
| weights["model.norm.weight"][0] = value | |
| (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") | |
| write_bf16(tmp_path / "model.safetensors", weights) | |
| with pytest.raises(loader.InvalidCheckpointError, match="Nonfinite"): | |
| loader.import_model(tmp_path) | |
| def test_shard_path_traversal(tmp_path, filename): | |
| (tmp_path / "config.json").write_text(json.dumps(tiny_config()), encoding="utf-8") | |
| (tmp_path / "model.safetensors.index.json").write_text( | |
| json.dumps({"weight_map": {"model.norm.weight": filename}}), encoding="utf-8" | |
| ) | |
| with pytest.raises(loader.InvalidCheckpointError, match="Unsafe shard"): | |
| loader.import_model(tmp_path) | |
| def test_shard_integrity(tmp_path, problem, match): | |
| weights = write_checkpoint(tmp_path, sharded=True) | |
| index_path = tmp_path / "model.safetensors.index.json" | |
| index = json.loads(index_path.read_text(encoding="utf-8")) | |
| mapping = index["weight_map"] | |
| names = list(weights) | |
| first, second = mapping[names[0]], mapping[names[1]] | |
| if problem == "duplicate": | |
| tensors = {name: weights[name] for name in names[1::2]} | |
| tensors[names[0]] = weights[names[0]] | |
| save_file(tensors, tmp_path / second) | |
| elif problem == "index_missing": | |
| del mapping[names[0]] | |
| elif problem == "index_extra": | |
| mapping["model.missing.weight"] = first | |
| elif problem == "wrong_shard": | |
| mapping[names[0]] = second | |
| elif problem == "missing_shard": | |
| (tmp_path / second).unlink() | |
| elif problem == "unreferenced": | |
| save_file({"extra": np.ones(1)}, tmp_path / "extra.safetensors") | |
| index_path.write_text(json.dumps(index), encoding="utf-8") | |
| with pytest.raises(loader.InvalidCheckpointError, match=match): | |
| loader.import_model(tmp_path) | |
| def test_duplicate_json_index_keys(tmp_path): | |
| write_checkpoint(tmp_path) | |
| (tmp_path / "model.safetensors.index.json").write_text( | |
| '{"weight_map":{"x":"model.safetensors","x":"model.safetensors"}}', encoding="utf-8" | |
| ) | |
| with pytest.raises(loader.InvalidCheckpointError, match="Duplicate JSON key"): | |
| loader.import_model(tmp_path) | |
| def test_hub_snapshot_and_explicit_tokenizer_trust(tmp_path, monkeypatch, offline_tokenizer, trust): | |
| write_checkpoint(tmp_path, tiny_config(eos_token_id=5)) | |
| (tmp_path / "generation_config.json").write_text(json.dumps({ | |
| "eos_token_id": [5, 3], "bos_token_id": 0, "pad_token_id": 3, "temperature": 0.6, | |
| }), encoding="utf-8") | |
| calls = [] | |
| def download(**kwargs): | |
| calls.append(kwargs) | |
| return str(tmp_path) | |
| monkeypatch.setattr(loader, "snapshot_download", download) | |
| loaded = loader.import_model("org/model", revision="release", cache_dir="cache", | |
| trust_remote_code=trust, local_files_only=True) | |
| assert loaded.source == str(tmp_path.resolve()) | |
| assert loaded.revision == tmp_path.name | |
| assert calls[0]["repo_id"] == "org/model" | |
| assert calls[0]["revision"] == "release" | |
| assert calls[0]["cache_dir"] == "cache" | |
| assert calls[0]["local_files_only"] is True | |
| assert "generation_config.json" in calls[0]["allow_patterns"] | |
| assert ("*.py" in calls[0]["allow_patterns"]) is trust | |
| assert "*.bin" not in calls[0]["allow_patterns"] | |
| assert offline_tokenizer[0][0][0] == (str(tmp_path.resolve()),) | |
| assert offline_tokenizer[0][0][1]["trust_remote_code"] is trust | |
| assert offline_tokenizer[0][0][1]["local_files_only"] is True | |
| assert loaded.config["eos_token_id"] == [5, 3] | |
| assert loaded.config["bos_token_id"] == 0 | |
| assert loaded.config["pad_token_id"] == 3 | |
| assert "temperature" not in loaded.config | |
| def test_trust_does_not_allow_custom_model(tmp_path, offline_tokenizer): | |
| write_checkpoint(tmp_path, tiny_config(auto_map={"AutoModelForCausalLM": "modeling_custom.Model"})) | |
| with pytest.raises(loader.UnsupportedModelError, match="auto_map"): | |
| loader.import_model(tmp_path, trust_remote_code=True) | |
| assert not offline_tokenizer[0] | |
| def test_cross_repo_tokenizer_code_rejected(tmp_path, offline_tokenizer): | |
| write_checkpoint(tmp_path) | |
| (tmp_path / "tokenizer_config.json").write_text(json.dumps({ | |
| "auto_map": {"AutoTokenizer": ["other/repo--tokenization_custom.Tokenizer", None]} | |
| }), encoding="utf-8") | |
| with pytest.raises(loader.UnsupportedModelError, match="pinned snapshot"): | |
| loader.import_model(tmp_path, trust_remote_code=True) | |
| assert not offline_tokenizer[0] | |
| def test_local_tokenizer_custom_code_allowed_only_explicitly(tmp_path, offline_tokenizer): | |
| write_checkpoint(tmp_path, tiny_config(auto_map={"AutoTokenizer": ["tokenization_custom.Tokenizer", None]})) | |
| loader.import_model(tmp_path, trust_remote_code=True, revision="local-label") | |
| assert offline_tokenizer[0][0][1]["trust_remote_code"] is True | |
| def test_missing_local_path_is_not_a_hub_id(tmp_path): | |
| with pytest.raises(loader.InvalidCheckpointError, match="Not a checkpoint directory"): | |
| loader.import_model(tmp_path / "missing") | |
| def test_corrupt_safetensors(tmp_path): | |
| write_checkpoint(tmp_path) | |
| (tmp_path / "model.safetensors").write_bytes(b"not a checkpoint") | |
| with pytest.raises(loader.InvalidCheckpointError, match="Cannot load"): | |
| loader.import_model(tmp_path) | |
| def test_duplicate_safetensors_header_key(tmp_path): | |
| write_checkpoint(tmp_path) | |
| entry = '"x":{"dtype":"F32","shape":[1],"data_offsets":[0,4]}' | |
| header = ("{" + entry + "," + entry + "}").encode() | |
| header += b" " * (-len(header) % 8) | |
| (tmp_path / "model.safetensors").write_bytes(struct.pack("<Q", len(header)) + header + b"\0" * 4) | |
| with pytest.raises(loader.InvalidCheckpointError, match="Duplicate JSON key"): | |
| loader.import_model(tmp_path) | |
| def test_shard_symlink_safety(tmp_path, hub_cache, sharded, hub1x): | |
| if hub1x: | |
| # huggingface_hub>=1.x decoupled global store: hub/blobs sharded by | |
| # hash prefix, snapshots under hub/models--*/snapshots/<rev>. | |
| root = tmp_path / "hub" / "models--org--name" / "snapshots" / "rev" | |
| root.mkdir(parents=True) | |
| blob_dir = tmp_path / "hub" / "blobs" | |
| else: | |
| root = tmp_path / "snapshots" / "commit" | |
| root.mkdir(parents=True) | |
| blob_dir = tmp_path / "blobs" | |
| if sharded: | |
| # huggingface_hub>=1.x shards blobs by hash prefix (blobs/ab/<hash>). | |
| blob_dir = blob_dir / "ab" | |
| blob_dir.mkdir(parents=True) | |
| target = blob_dir / "hash" if hub_cache else tmp_path / "outside.safetensors" | |
| save_file(tiny_weights(tiny_config()), target) | |
| (root / "config.json").write_text(json.dumps(tiny_config()), encoding="utf-8") | |
| try: | |
| (root / "model.safetensors").symlink_to(target) | |
| except OSError: | |
| pytest.skip("OS does not permit unprivileged symlinks") | |
| if hub_cache: | |
| assert loader.import_model(root).weights | |
| else: | |
| with pytest.raises(loader.InvalidCheckpointError, match="escapes checkpoint directory"): | |
| loader.import_model(root) | |
| def test_public_qwen3_dimension_and_rope_config(): | |
| # Public Qwen3-0.6B dimensions: 16 * 128 != hidden_size=1024. | |
| config = tiny_config("qwen3", vocab_size=151936, hidden_size=1024, intermediate_size=3072, | |
| num_hidden_layers=28, num_attention_heads=16, num_key_value_heads=8, | |
| head_dim=128, max_position_embeddings=40960, rope_theta=1000000, | |
| tie_word_embeddings=True, torch_dtype="bfloat16") | |
| normalized = loader.validate_config(config) | |
| assert normalized["head_dim"] == 128 | |
| assert normalized["rope_theta"] == 1000000 | |
| assert normalized["tie_word_embeddings"] is True | |
| def test_local_qwen3_generation_token_overrides(tmp_path): | |
| config = tiny_config("qwen3", vocab_size=151936, eos_token_id=151645, | |
| bos_token_id=1, pad_token_id=None) | |
| write_checkpoint(tmp_path, config) | |
| generation = { | |
| "eos_token_id": [151645, 151643, 151645], | |
| "bos_token_id": 151643, | |
| "pad_token_id": 151643, | |
| "do_sample": True, "temperature": 0.6, "top_k": 20, "top_p": 0.95, | |
| "max_new_tokens": 100, "repetition_penalty": 1.1, | |
| "hidden_size": 999, "vocab_size": 999, "forced_eos_token_id": 2, | |
| } | |
| (tmp_path / "generation_config.json").write_text(json.dumps(generation), encoding="utf-8") | |
| loaded = loader.import_model(tmp_path) | |
| expected = loader.validate_config({**config, "eos_token_id": [151645, 151643], | |
| "bos_token_id": 151643, "pad_token_id": 151643}) | |
| assert loaded.config == expected | |
| assert loader.validate_config(loaded.config) == expected | |
| assert json.loads((tmp_path / "config.json").read_text(encoding="utf-8")) == config | |
| def test_missing_or_null_generation_tokens_preserve_config(tmp_path, generation): | |
| config = tiny_config(eos_token_id=[5, 3], bos_token_id=0, pad_token_id=6) | |
| write_checkpoint(tmp_path, config) | |
| if generation is not None: | |
| (tmp_path / "generation_config.json").write_text(json.dumps(generation), encoding="utf-8") | |
| assert loader.import_model(tmp_path).config == loader.validate_config(config) | |
| def test_generation_scalar_zero_overrides(tmp_path, key): | |
| config = tiny_config(eos_token_id=[5, 3], bos_token_id=1, pad_token_id=6) | |
| write_checkpoint(tmp_path, config) | |
| (tmp_path / "generation_config.json").write_text(json.dumps({key: 0}), encoding="utf-8") | |
| assert loader.import_model(tmp_path).config == loader.validate_config({**config, key: 0}) | |
| def test_invalid_generation_token_ids(tmp_path, offline_tokenizer, key, value): | |
| write_checkpoint(tmp_path) | |
| (tmp_path / "generation_config.json").write_text(json.dumps({key: value}), encoding="utf-8") | |
| with pytest.raises(loader.InvalidCheckpointError, match=key): | |
| loader.import_model(tmp_path) | |
| with pytest.raises(loader.InvalidCheckpointError, match=key): | |
| loader.validate_config(tiny_config(**{key: value})) | |
| assert not offline_tokenizer[0] | |
| def test_config_token_normalization_does_not_mutate_input(): | |
| config = tiny_config(eos_token_id=[5, 3, 5], bos_token_id=0, pad_token_id=6) | |
| normalized = loader.validate_config(config) | |
| assert normalized["eos_token_id"] == [5, 3] | |
| assert config["eos_token_id"] == [5, 3, 5] | |
| def test_invalid_generation_config_json(tmp_path, offline_tokenizer, content): | |
| write_checkpoint(tmp_path) | |
| (tmp_path / "generation_config.json").write_text(content, encoding="utf-8") | |
| with pytest.raises(loader.InvalidCheckpointError, match="generation_config.json"): | |
| loader.import_model(tmp_path) | |
| assert not offline_tokenizer[0] | |