92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
"""Tests for layer-aware SafeTensors snapshot file selection."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.safetensors_selection import select_safetensors_files_for_layers
|
|
|
|
|
|
def _write_snapshot(tmp_path, *, config=None):
|
|
(tmp_path / "config.json").write_text(
|
|
json.dumps(config or {"num_hidden_layers": 5}),
|
|
encoding="utf-8",
|
|
)
|
|
(tmp_path / "tokenizer.json").write_text("{}", encoding="utf-8")
|
|
(tmp_path / "tokenizer_config.json").write_text("{}", encoding="utf-8")
|
|
(tmp_path / "README.md").write_text("not part of runtime snapshot", encoding="utf-8")
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps({
|
|
"metadata": {"total_size": 123},
|
|
"weight_map": {
|
|
"model.embed_tokens.weight": "model-00001-of-00004.safetensors",
|
|
"model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
|
|
"model.layers.1.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
|
|
"model.layers.2.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
|
|
"model.layers.3.mlp.down_proj.weight": "nested/model-00003-of-00004.safetensors",
|
|
"model.layers.4.self_attn.q_proj.weight": "nested/model-00003-of-00004.safetensors",
|
|
"model.norm.weight": "model-00004-of-00004.safetensors",
|
|
"lm_head.weight": "model-00004-of-00004.safetensors",
|
|
},
|
|
}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def test_selects_only_weight_shards_for_middle_layer_range(tmp_path):
|
|
"Selects only weight shards for middle layer range\n\nTags: general"
|
|
_write_snapshot(tmp_path)
|
|
|
|
files = select_safetensors_files_for_layers(tmp_path, 2, 3)
|
|
|
|
assert files == [
|
|
"config.json",
|
|
"model-00002-of-00004.safetensors",
|
|
"model.safetensors.index.json",
|
|
"nested/model-00003-of-00004.safetensors",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
]
|
|
|
|
|
|
def test_head_range_includes_embeddings(tmp_path):
|
|
"Head range includes embeddings\n\nTags: general"
|
|
_write_snapshot(tmp_path)
|
|
|
|
files = select_safetensors_files_for_layers(tmp_path, 0, 0)
|
|
|
|
assert "model-00001-of-00004.safetensors" in files
|
|
assert "model-00004-of-00004.safetensors" not in files
|
|
|
|
|
|
def test_tail_range_includes_norm_and_lm_head_from_inferred_layer_count(tmp_path):
|
|
"Tail range includes norm and lm head from inferred layer count\n\nTags: general"
|
|
_write_snapshot(tmp_path, config={"text_config": {"num_hidden_layers": 5}})
|
|
|
|
files = select_safetensors_files_for_layers(tmp_path, 4, 4)
|
|
|
|
assert "nested/model-00003-of-00004.safetensors" in files
|
|
assert "model-00004-of-00004.safetensors" in files
|
|
assert "model-00001-of-00004.safetensors" not in files
|
|
|
|
|
|
def test_tail_files_are_not_selected_without_total_layer_count(tmp_path):
|
|
"Tail files are not selected without total layer count\n\nTags: general"
|
|
_write_snapshot(tmp_path, config={"architectures": ["UnknownForTest"]})
|
|
|
|
files = select_safetensors_files_for_layers(tmp_path, 4, 4)
|
|
|
|
assert "nested/model-00003-of-00004.safetensors" in files
|
|
assert "model-00004-of-00004.safetensors" not in files
|
|
|
|
|
|
def test_rejects_unsafe_weight_map_paths(tmp_path):
|
|
"Rejects unsafe weight map paths\n\nTags: general"
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": {"model.layers.0.weight": "../escape.safetensors"}}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="unsafe relative file"):
|
|
select_safetensors_files_for_layers(tmp_path, 0, 0)
|