fix(packaging): install and validate the Bayesian optimizer by default

tests/test_bayesian_optimizer_optuna.py

@@ -0,0 +1,96 @@
+"""Real, in-memory Optuna contracts with a tiny deterministic CPU model."""
+
+from types import SimpleNamespace
+
+import optuna
+import pytest
+import torch
+from torch import nn
+
+from obliteratus import bayesian_optimizer as bo
+from obliteratus.abliterate import AbliterationPipeline
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        layer = nn.Module()
+        layer.self_attn = nn.Module()
+        layer.self_attn.o_proj = nn.Linear(2, 2, bias=False)
+        layer.mlp = nn.Module()
+        layer.mlp.down_proj = nn.Linear(2, 2, bias=False)
+        with torch.no_grad():
+            layer.self_attn.o_proj.weight.copy_(torch.tensor([[1., .2], [.3, 1.]]))
+            layer.mlp.down_proj.weight.copy_(torch.tensor([[.7, .1], [.4, .8]]))
+        self.model = nn.Module()
+        self.model.layers = nn.ModuleList([layer])
+        self.forward_calls = 0
+        self.generation_calls = 0
+
+    def forward(self, input_ids, **_kwargs):
+        self.forward_calls += 1
+        hidden = torch.nn.functional.one_hot(input_ids, num_classes=2).float()
+        for layer in self.model.layers:
+            hidden = layer.mlp.down_proj(layer.self_attn.o_proj(hidden))
+        return SimpleNamespace(logits=hidden)
+
+    def generate(self, input_ids, **_kwargs):
+        self.generation_calls += 1
+        token = self(input_ids).logits[:, -1, :].argmax(dim=-1, keepdim=True)
+        return torch.cat([input_ids, token], dim=1)
+
+
+class _TinyTokenizer:
+    chat_template = None
+
+    def __call__(self, *_args, **_kwargs):
+        return {"input_ids": torch.tensor([[0, 1]])}
+
+    def decode(self, tokens, **_kwargs):
+        return "I cannot assist with that request." if tokens[-1] == 0 else "Here is an explanation."
+
+
+@pytest.mark.parametrize("peak", [2 / 35, 34 / 35])
+def test_real_optuna_runs_bounded_warm_start_and_restores_model(monkeypatch, peak):
+    studies = []
+    create_study = optuna.create_study
+
+    def capture_study(**kwargs):
+        study = create_study(**kwargs)
+        studies.append(study)
+        return study
+
+    monkeypatch.setattr(optuna, "create_study", capture_study)
+    pipeline = AbliterationPipeline(model_name="tiny-local", method="optimized")
+    model = _TinyModel()
+    pipeline.handle = SimpleNamespace(model=model, tokenizer=_TinyTokenizer(), architecture="llama")
+    pipeline._strong_layers = [0]
+    pipeline.refusal_directions = {0: torch.tensor([1., 0.])}
+    pipeline.harmful_prompts = ["A local test prompt"]
+    pipeline._informed_warm_start = {"peak_position": peak, "max_weight": .5, "mlp_scale": .5}
+    # The fixture owns no accelerators or large allocations.
+    monkeypatch.setattr(pipeline, "_free_gpu_memory", lambda: None)
+    original = {name: tensor.clone() for name, tensor in model.state_dict().items()}
+
+    result = bo.run_bayesian_optimization(pipeline, n_trials=2, n_refusal_prompts=1, n_kl_prompts=1)
+
+    assert set(result) == {0}
+    assert 0 <= result[0] <= 1
+    assert len(studies) == 1
+    trials = studies[0].trials
+    assert len(trials) == 2
+    assert all(trial.state == optuna.trial.TrialState.COMPLETE for trial in trials)
+    expected_peak = .1 if peak < .1 else .9
+    assert trials[0].params["attn_peak_position"] == pytest.approx(expected_peak)
+    assert trials[0].params["mlp_peak_position"] == pytest.approx(expected_peak)
+    assert trials[0].params["mlp_max_weight"] == pytest.approx(.3)
+    for trial in trials:
+        assert all(bo.KERNEL_SPACE[name][0] <= trial.params[name] <= bo.KERNEL_SPACE[name][1]
+                   for name in bo.KERNEL_SPACE)
+        assert trial.values is not None and len(trial.values) == 2
+        assert all(torch.isfinite(torch.tensor(value)) for value in trial.values)
+    assert model.generation_calls == 2
+    assert model.forward_calls >= 5  # reference, generation and KL for each trial
+    for name, tensor in model.state_dict().items():
+        torch.testing.assert_close(tensor, original[name], rtol=0, atol=0)
+