scripts/test_steering.py

"""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()