restore-multimodal + tests + risk map

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"

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)