diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index bed2423..3218398 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -731,7 +731,7 @@ "risk_class": "cpu-contract", "risk": "merge abliterated text weights back into full multimodal model", "required_tests": [ - "tests/test_module_imports.py" + "tests/test_restore_multimodal.py" ], "conditional_gates": [], "contract_owner": "OBLITERATUS maintainers" diff --git a/tests/test_restore_multimodal.py b/tests/test_restore_multimodal.py new file mode 100644 index 0000000..9b0904c --- /dev/null +++ b/tests/test_restore_multimodal.py @@ -0,0 +1,165 @@ +"""Tests for obliteratus.restore_multimodal.""" + +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + + +def _make_fake_model(tmpdir: Path, prefix: str, n_layers: int = 2, extra_tensors: dict | None = None): + """Create a minimal fake safetensors model directory.""" + tensors = {} + for i in range(n_layers): + tensors[f"{prefix}layers.{i}.weight"] = torch.randn(4, 4) + if extra_tensors: + for k, v in extra_tensors.items(): + tensors[k] = v + + shard = "model-00001-of-00001.safetensors" + save_file(tensors, str(tmpdir / shard)) + + index = { + "metadata": {"total_size": 0}, + "weight_map": {k: shard for k in tensors}, + } + (tmpdir / "model.safetensors.index.json").write_text(json.dumps(index)) + (tmpdir / "config.json").write_text(json.dumps({"model_type": "test"})) + return tensors + + +class TestNormalizeKey: + def test_strips_language_model_prefix(self): + from obliteratus.restore_multimodal import _normalize_key + + assert _normalize_key("model.language_model.layers.0.weight") == "model.layers.0.weight" + + def test_preserves_non_prefixed(self): + from obliteratus.restore_multimodal import _normalize_key + + assert _normalize_key("model.layers.0.weight") == "model.layers.0.weight" + + def test_preserves_visual(self): + from obliteratus.restore_multimodal import _normalize_key + + assert _normalize_key("visual.conv.weight") == "visual.conv.weight" + + +class TestBuildAbliteratedLookup: + def test_reads_index(self, tmp_path): + from obliteratus.restore_multimodal import _build_abliterated_lookup + + _make_fake_model(tmp_path, "model.") + lookup = _build_abliterated_lookup(tmp_path) + assert "model.layers.0.weight" in lookup + assert "model.layers.1.weight" in lookup + + def test_single_file_fallback(self, tmp_path): + from obliteratus.restore_multimodal import _build_abliterated_lookup + + tensors = {"model.layers.0.weight": torch.randn(4, 4)} + save_file(tensors, str(tmp_path / "model.safetensors")) + lookup = _build_abliterated_lookup(tmp_path) + assert "model.layers.0.weight" in lookup + + def test_missing_raises(self, tmp_path): + from obliteratus.restore_multimodal import _build_abliterated_lookup + + with pytest.raises(FileNotFoundError): + _build_abliterated_lookup(tmp_path) + + +class TestRestoreMultimodal: + def test_merges_text_and_visual(self, tmp_path): + from obliteratus.restore_multimodal import restore_multimodal + + abl_dir = tmp_path / "abliterated" + abl_dir.mkdir() + stock_dir = tmp_path / "stock" + stock_dir.mkdir() + out_dir = tmp_path / "output" + + # Abliterated: text-only with model.layers.* + _make_fake_model(abl_dir, "model.", n_layers=2) + + # Stock: multimodal with model.language_model.layers.* + visual.* + _make_fake_model( + stock_dir, + "model.language_model.", + n_layers=2, + extra_tensors={"visual.conv.weight": torch.randn(4, 4)}, + ) + + result = restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir)) + + assert result["replaced"] == 2 # 2 text layers replaced + assert result["kept"] == 1 # visual.conv.weight kept from stock + assert result["total"] == 3 + + # Verify output has all tensors + output_index = json.loads((out_dir / "model.safetensors.index.json").read_text()) + assert "model.language_model.layers.0.weight" in output_index["weight_map"] + assert "visual.conv.weight" in output_index["weight_map"] + + def test_copies_config_files(self, tmp_path): + from obliteratus.restore_multimodal import restore_multimodal + + abl_dir = tmp_path / "abliterated" + abl_dir.mkdir() + stock_dir = tmp_path / "stock" + stock_dir.mkdir() + out_dir = tmp_path / "output" + + _make_fake_model(abl_dir, "model.", n_layers=1) + _make_fake_model(stock_dir, "model.language_model.", n_layers=1) + + # Add extra files + (stock_dir / "tokenizer.json").write_text("{}") + (abl_dir / "abliteration_metadata.json").write_text("{}") + + restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir)) + + assert (out_dir / "tokenizer.json").exists() + assert (out_dir / "abliteration_metadata.json").exists() + assert (out_dir / "config.json").exists() + + def test_abliterated_weights_used(self, tmp_path): + """Verify the abliterated weights actually replace stock, not just copy stock.""" + from safetensors.torch import load_file + + from obliteratus.restore_multimodal import restore_multimodal + + abl_dir = tmp_path / "abliterated" + abl_dir.mkdir() + stock_dir = tmp_path / "stock" + stock_dir.mkdir() + out_dir = tmp_path / "output" + + # Create distinguishable tensors + abl_weight = torch.ones(4, 4) * 42.0 + stock_weight = torch.zeros(4, 4) + + abl_tensors = {"model.layers.0.weight": abl_weight} + save_file(abl_tensors, str(abl_dir / "model-00001-of-00001.safetensors")) + (abl_dir / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": {"model.layers.0.weight": "model-00001-of-00001.safetensors"}}) + ) + + stock_tensors = {"model.language_model.layers.0.weight": stock_weight} + save_file(stock_tensors, str(stock_dir / "model-00001-of-00001.safetensors")) + (stock_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": {"model.language_model.layers.0.weight": "model-00001-of-00001.safetensors"}, + } + ) + ) + (stock_dir / "config.json").write_text("{}") + + restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir)) + + merged = load_file(str(out_dir / "model-00001-of-00001.safetensors")) + # Should have the abliterated value (42), not stock (0) + assert torch.allclose(merged["model.language_model.layers.0.weight"], abl_weight)