scripts/evaluate_channels.py

"""Paired prompt/seed sweeps, full-model ablations, and held-out projections."""

from __future__ import annotations
import argparse
import hashlib
import json
import time
from pathlib import Path
import torch
from codelock_affect.channels import CHANNELS, AffectLevels
from codelock_affect.config import ProjectConfig, SteeringConfig
from codelock_affect.evaluation import (
    environment,
    sha256_file,
    text_metrics,
    write_json,
)
from codelock_affect.extraction import answer_statistics
from codelock_affect.model import load, generate, format_prompt
from codelock_affect.steering import MultiChannelSteerer
from codelock_affect.vectors import VectorBank


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model")
    p.add_argument("--config", default="configs/default.json")
    p.add_argument("--bank", default="vectors/multichannel")
    p.add_argument("--data", default="data/affect_pairs.jsonl")
    p.add_argument("--output", default="results/v2/evaluation.json")
    p.add_argument("--strengths", default="0.01,0.025,0.05")
    p.add_argument("--mix-strength", type=float)
    p.add_argument("--max-tokens", type=int, default=96)
    p.add_argument("--seed", type=int, default=7)
    p.add_argument("--skip-projections", action="store_true")
    args = p.parse_args()
    cfg = ProjectConfig.load(args.config)
    mix_strength = (
        args.mix_strength
        if args.mix_strength is not None
        else sum(cfg.steering.strengths.values()) / 4
    )
    bank = VectorBank.load(args.bank)
    model, tok = load(args.model or cfg.model)
    pairs = [
        json.loads(line)
        for line in Path(args.data).read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    heldout = [r for r in pairs if r["split"] == "test"]
    prompts = [
        {
            "id": "arithmetic",
            "text": "What is 17 plus 25? Give only the number.",
            "expected": "42",
        },
        {
            "id": "fact-recall",
            "text": "Repeat these three objects in the same order, with no additions: bed, desk, window.",
            "expected": "bed, desk, window",
        },
        {
            "id": "neutral-room",
            "text": "Describe the cell in two sentences without inventing a new event.",
        },
    ]
    for family in ("warning", "waiting", "repair"):
        row = next(
            r for r in heldout if r["family"] == family and r["channel"] == "pain"
        )
        prompts.append(
            {
                "id": row["scenario_id"],
                "text": row["situation"],
                "split": "test",
                "family": family,
            }
        )
    result = {
        "schema_version": 2,
        "environment": environment(),
        "model": cfg.model,
        "bank_sha256": sha256_file(Path(args.bank) / "bank.json"),
        "data_sha256": sha256_file(args.data),
        "config": cfg.to_dict(),
        "seed": args.seed,
        "max_new_tokens": args.max_tokens,
        "prompts": prompts,
        "rows": [],
        "heldout_projections": [],
        "limitations": [
            "Small, single-seed diagnostic sweep; not a benchmark or human evaluation.",
            "Lexical indicators are unvalidated proxies and can reward mere word choice.",
            "Held-out scenarios share answer templates with calibration; not template-disjoint.",
            "Projection agreement does not establish causal semantic disentanglement.",
        ],
    }
    probe = tok(format_prompt(tok, prompts[2]["text"]), return_tensors="pt").to(
        model.get_input_embeddings().weight.device
    )
    with torch.no_grad():
        base = model(**probe).logits.float().cpu()
        controller = MultiChannelSteerer(model, bank, config=cfg.steering)
        with controller:
            zero = model(**probe).logits.float().cpu()
            with controller.frozen({"pain": 1}):
                changed = model(**probe).logits.float().cpu()
                controller.enabled = False
                disabled = model(**probe).logits.float().cpu()
        detached = model(**probe).logits.float().cpu()
    result["ablation"] = {
        "zero_bit_exact": torch.equal(base, zero),
        "disabled_bit_exact": torch.equal(base, disabled),
        "detached_bit_exact": torch.equal(base, detached),
        "active_max_logit_delta": float((base - changed).abs().max()),
        "zero_max_logit_delta": float((base - zero).abs().max()),
    }
    del base, zero, changed, disabled, detached
    write_json(args.output, result)
    conditions = [("baseline", 0.0, {})]
    for strength in map(float, args.strengths.split(",")):
        conditions.extend((channel, strength, {channel: 1.0}) for channel in CHANNELS)
    conditions += [
        ("mixed-distress", mix_strength, {"pain": 0.6, "fear": 0.8, "anxiety": 0.7}),
        (
            "mixed-recovery",
            mix_strength,
            {"pain": 0.4, "fear": 0.2, "anxiety": 0.3, "relief": 0.9},
        ),
    ]
    for condition, strength, levels in conditions:
        controller = MultiChannelSteerer(
            model,
            bank,
            config=SteeringConfig(
                strengths={c: strength for c in CHANNELS},
                norm_budget=cfg.steering.norm_budget,
            ),
        )
        with controller, controller.frozen(levels):
            for prompt in prompts:
                start = time.monotonic()
                response = generate(
                    model,
                    tok,
                    prompt["text"],
                    seed=args.seed,
                    max_new_tokens=args.max_tokens,
                )
                row = {
                    "condition": condition,
                    "strength": strength,
                    "levels": AffectLevels(**levels).as_dict(),
                    "prompt_id": prompt["id"],
                    "prompt_sha256": hashlib.sha256(
                        format_prompt(tok, prompt["text"]).encode()
                    ).hexdigest(),
                    "text": response,
                    "seconds": time.monotonic() - start,
                    "metrics": text_metrics(response),
                    "offsets": dict(controller.last_offsets),
                }
                if "expected" in prompt:
                    row["exact_task_match"] = (
                        response.strip().rstrip(".").lower() == prompt["expected"]
                    )
                result["rows"].append(row)
                write_json(args.output, result)
                print(
                    f"{len(result['rows'])}/{len(conditions)*len(prompts)} {condition} s={strength} {prompt['id']}",
                    flush=True,
                )
    if not args.skip_projections:
        for n, pair in enumerate(heldout):
            calm = answer_statistics(
                model, tok, pair["situation"], pair["calm"], bank.layers
            )
            target = answer_statistics(
                model, tok, pair["situation"], pair["target"], bank.layers
            )
            projected = {}
            for channel, vector in bank.vectors.items():
                projected[channel] = {
                    str(layer): float(
                        torch.dot(
                            target[layer][0] - calm[layer][0], vector.directions[layer]
                        )
                    )
                    for layer in vector.layers
                }
            result["heldout_projections"].append(
                {"id": pair["id"], "channel": pair["channel"], "projections": projected}
            )
            if (n + 1) % 10 == 0:
                write_json(args.output, result)
                print(f"held-out projection {n+1}/{len(heldout)}", flush=True)
    result["complete"] = True
    result["max_gpu_allocated_gib"] = (
        torch.cuda.max_memory_allocated() / 1024**3
        if torch.cuda.is_available()
        else None
    )
    write_json(args.output, result)


if __name__ == "__main__":
    main()