diff --git a/tests/test_restore_multimodal.py b/tests/test_restore_multimodal.py index 9b0904c..aaa39b0 100644 --- a/tests/test_restore_multimodal.py +++ b/tests/test_restore_multimodal.py @@ -114,8 +114,10 @@ class TestRestoreMultimodal: _make_fake_model(abl_dir, "model.", n_layers=1) _make_fake_model(stock_dir, "model.language_model.", n_layers=1) - # Add extra files + # 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)) @@ -123,6 +125,9 @@ class TestRestoreMultimodal: 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.""" @@ -163,3 +168,27 @@ class TestRestoreMultimodal: 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()