OBLITERATUS/tests/test_restore_multimodal.py

"""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 including ones that should be skipped
        (stock_dir / "tokenizer.json").write_text("{}")
        (stock_dir / "README.md").write_text("skip me")
        (stock_dir / ".gitattributes").write_text("skip me too")
        (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()
        # Skipped files should NOT be copied
        assert not (out_dir / "README.md").exists()
        assert not (out_dir / ".gitattributes").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)


class TestMain:
    def test_cli_main(self, tmp_path):
        from obliteratus.restore_multimodal import main

        abl_dir = tmp_path / "abl"
        abl_dir.mkdir()
        stock_dir = tmp_path / "stock"
        stock_dir.mkdir()
        out_dir = tmp_path / "out"

        _make_fake_model(abl_dir, "model.", n_layers=1)
        _make_fake_model(stock_dir, "model.language_model.", n_layers=1)

        import sys
        old_argv = sys.argv
        sys.argv = ["prog", "--abliterated", str(abl_dir), "--stock", str(stock_dir), "--output", str(out_dir)]
        try:
            main()
        finally:
            sys.argv = old_argv

        assert (out_dir / "model.safetensors.index.json").exists()