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