mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 17:00:50 +02:00
96 lines
3.9 KiB
Python
96 lines
3.9 KiB
Python
"""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)
|