mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-29 22:20:36 +02:00
faber
17f1d5eb59
Carry the coherent blending contribution and research summary from PR #127 while splitting the capacity and recovery proposals into issues #132 and #133.
151 lines
4.5 KiB
Python
151 lines
4.5 KiB
Python
"""Tests for obliteratus.blend — complementary abliteration blending."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
def _make_model(tmpdir: Path, value: float, n_tensors: int = 3):
|
|
"""Create a minimal model with all weights set to a constant value."""
|
|
tensors = {f"layer.{i}.weight": torch.full((4, 4), value) for i in range(n_tensors)}
|
|
shard = "model-00001-of-00001.safetensors"
|
|
save_file(tensors, str(tmpdir / shard))
|
|
index = {"metadata": {}, "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 TestBlendModels:
|
|
def test_lerp_blend(self, tmp_path):
|
|
from obliteratus.blend import blend_models
|
|
|
|
a_dir = tmp_path / "model_a"
|
|
b_dir = tmp_path / "model_b"
|
|
out_dir = tmp_path / "blended"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 0.0)
|
|
_make_model(b_dir, 10.0)
|
|
|
|
result = blend_models(str(a_dir), str(b_dir), str(out_dir), alpha=0.6)
|
|
|
|
assert result["blended_tensors"] == 3
|
|
assert result["alpha"] == 0.6
|
|
|
|
merged = load_file(str(out_dir / "model-00001-of-00001.safetensors"))
|
|
# 0.6 * 10.0 + 0.4 * 0.0 = 6.0
|
|
assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 6.0))
|
|
|
|
def test_alpha_zero_is_pure_a(self, tmp_path):
|
|
from obliteratus.blend import blend_models
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "out"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 42.0)
|
|
_make_model(b_dir, 99.0)
|
|
|
|
blend_models(str(a_dir), str(b_dir), str(out), alpha=0.0)
|
|
merged = load_file(str(out / "model-00001-of-00001.safetensors"))
|
|
assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 42.0))
|
|
|
|
def test_alpha_one_is_pure_b(self, tmp_path):
|
|
from obliteratus.blend import blend_models
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "out"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 42.0)
|
|
_make_model(b_dir, 99.0)
|
|
|
|
blend_models(str(a_dir), str(b_dir), str(out), alpha=1.0)
|
|
merged = load_file(str(out / "model-00001-of-00001.safetensors"))
|
|
assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 99.0))
|
|
|
|
def test_writes_metadata(self, tmp_path):
|
|
from obliteratus.blend import blend_models
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "out"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 1.0)
|
|
_make_model(b_dir, 2.0)
|
|
|
|
blend_models(str(a_dir), str(b_dir), str(out), alpha=0.5)
|
|
|
|
meta = json.loads((out / "blend_metadata.json").read_text())
|
|
assert meta["alpha"] == 0.5
|
|
assert meta["blend_method"] == "lerp"
|
|
|
|
def test_copies_config(self, tmp_path):
|
|
from obliteratus.blend import blend_models
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "out"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 1.0)
|
|
_make_model(b_dir, 2.0)
|
|
|
|
blend_models(str(a_dir), str(b_dir), str(out))
|
|
assert (out / "config.json").exists()
|
|
|
|
|
|
class TestBlendSearch:
|
|
def test_creates_multiple_blends(self, tmp_path):
|
|
from obliteratus.blend import blend_search
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "search"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 0.0)
|
|
_make_model(b_dir, 10.0)
|
|
|
|
results = blend_search(str(a_dir), str(b_dir), str(out), alphas=[0.3, 0.7])
|
|
|
|
assert len(results) == 2
|
|
assert (out / "blend_30").exists()
|
|
assert (out / "blend_70").exists()
|
|
|
|
|
|
class TestCLIDispatch:
|
|
def test_blend_dispatch(self, tmp_path):
|
|
from obliteratus.cli import main as cli_main
|
|
|
|
a_dir = tmp_path / "a"
|
|
b_dir = tmp_path / "b"
|
|
out = tmp_path / "out"
|
|
a_dir.mkdir()
|
|
b_dir.mkdir()
|
|
|
|
_make_model(a_dir, 1.0)
|
|
_make_model(b_dir, 2.0)
|
|
|
|
cli_main(["blend",
|
|
"--model-a", str(a_dir),
|
|
"--model-b", str(b_dir),
|
|
"--alpha", "0.6",
|
|
"--output", str(out)])
|
|
|
|
assert (out / "blend_metadata.json").exists()
|