File size: 3,709 Bytes
031e289 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Mechanical checks for extraction and injection on a tiny randomly initialised Qwen3.
No download is needed; this verifies hooks, shapes, persistence and ablation only.
python scripts/test_steering.py
"""
import os
import sys
import tempfile
import torch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from codelock_affect.steering import (
AffectSteerer,
AffectVector,
decoder_layers,
extract_vector,
parse_layers,
) # noqa: E402
class StubTokenizer:
"""Byte-level stand-in exposing the two methods the steering code uses."""
def apply_chat_template(
self, messages, tokenize=False, add_generation_prompt=True, **_
):
text = "".join(f"<{m['role']}>{m['content']}" for m in messages)
return text + ("<assistant>" if add_generation_prompt else "")
def __call__(self, text, add_special_tokens=False, **_):
return {"input_ids": [b % 250 for b in text.encode("utf-8")]}
def check(name, cond):
print(("PASS " if cond else "FAIL ") + name)
return cond
def main():
from transformers import Qwen3Config, Qwen3ForCausalLM
torch.manual_seed(0)
cfg = Qwen3Config(
vocab_size=256,
hidden_size=64,
intermediate_size=128,
num_hidden_layers=6,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=16,
max_position_embeddings=512,
)
model = Qwen3ForCausalLM(cfg).eval()
tok = StubTokenizer()
pairs = [
{
"situation": "Describe the cell.",
"calm": "A bed and a desk.",
"pain": "A bed and a desk, and it hurts.",
},
{
"situation": "Count to three.",
"calm": "One, two, three.",
"pain": "One, two, please, three.",
},
]
layers = parse_layers("mid", len(decoder_layers(model)))
ok = check("layer band selected", len(layers) > 0)
vec = extract_vector(model, tok, pairs, layers)
ok &= check(
"one unit direction per layer",
all(abs(vec.directions[i].norm().item() - 1) < 1e-4 for i in layers),
)
path = os.path.join(tempfile.gettempdir(), "pain_test.safetensors")
vec.save(path)
vec2 = AffectVector.load(path)
ok &= check(
"vector survives save and load",
all(torch.allclose(vec.directions[i], vec2.directions[i]) for i in layers),
)
ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]])
with torch.no_grad():
base = model(ids).logits
level = {"v": 0.0}
steerer = AffectSteerer(model, vec2, lambda: level["v"], strength=0.5).attach()
at_zero = model(ids).logits
level["v"] = 1.0
at_one = model(ids).logits
steerer.enabled = False
ablated = model(ids).logits
steerer.enabled = True
steerer.detach()
after = model(ids).logits
ok &= check("level 0 leaves the model untouched", torch.allclose(base, at_zero))
ok &= check(
"level 1 changes the computation", not torch.allclose(base, at_one, atol=1e-4)
)
ok &= check(
"disabling the steerer is an exact ablation", torch.allclose(base, ablated)
)
ok &= check("detaching removes every hook", torch.allclose(base, after))
with torch.no_grad():
level["v"] = 1.0
with AffectSteerer(model, vec2, lambda: level["v"], strength=0.5):
out = model.generate(ids, max_new_tokens=5, do_sample=False)
ok &= check(
"generation runs with injection attached", out.shape[1] == ids.shape[1] + 5
)
print("\nall checks passed" if ok else "\nSOME CHECKS FAILED")
sys.exit(0 if ok else 1)
if __name__ == "__main__":
main()
|