#!/usr/bin/env python3 """Measure every scenario in advance and write data/recorded/.json. python scripts/measure.py # the default model python scripts/measure.py --model HuggingFaceTB/SmolLM2-135M-Instruct python scripts/measure.py --sweep # re-tune layer and coefficient The Space shows these results when the page opens, so nobody spends GPU time to see them; Run re-measures them live with the same code. Every probability and every continuation is produced here. The prompts, opening words, direction labels, takeaways, contrast continuations and the per-model layer and coefficient are authored, in steering/scenarios.py; nothing else is. """ from __future__ import annotations import argparse import datetime as dt import json import platform import sys from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) import torch # noqa: E402 import transformers # noqa: E402 import models # noqa: E402 from steering import trace as T # noqa: E402 from steering.core import ( # noqa: E402 ALPHAS, REVISIONS, STORED_SIGNIFICANT_DIGITS, identity_check, load_model, measure_states, num_layers, separation_score, steering_vector, ) from steering.scenarios import SCENARIOS, TUNING, takeaway_for, tuning_for # noqa: E402 RECORDED = ROOT / "data" / "recorded" SCHEMA = "steer.recorded/v1" def sweep(tokenizer, model, spec, layers, scales): """Prints a score for each (layer, coefficient) so a new model can be retuned by eye.""" best = None for layer in layers: vector = steering_vector( tokenizer, model, spec.prompt, spec.prefix, spec.negative_examples, spec.positive_examples, layer, ) for scale in scales: states = measure_states( tokenizer, model, spec.prompt, spec.prefix, vector, layer, scale=scale, alphas=(-2.0, 0.0, 2.0), generate=False, top_k=6, ) score = separation_score(states) winners = [state.top_tokens[0] for state in states] print(f" layer {layer:>2} coefficient {scale:<4} score {score['score']:.3f} winners {winners}") if best is None or score["score"] > best[0]: best = (score["score"], layer, scale) print(f" -> best for {spec.id}: layer {best[1]}, coefficient {best[2]} (score {best[0]:.3f})") return best def main() -> int: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--model", default=models.DEFAULT.model_id, help="a Hugging Face model id") parser.add_argument("--out", help="output file; default data/recorded/.json") parser.add_argument("--max-new-tokens", type=int, default=40) parser.add_argument("--sweep", action="store_true", help="re-tune layer and coefficient") parser.add_argument("--device", default="cpu") args = parser.parse_args() torch.manual_seed(0) print(f"loading {args.model} ...") tokenizer, model = load_model(args.model, device=args.device) depth = num_layers(model) print(f" {depth} layers, hidden size {model.config.hidden_size}") if args.sweep: layers = sorted({max(1, round(depth * f)) for f in (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8)}) for spec in SCENARIOS: print(f"\n{spec.id}:") sweep(tokenizer, model, spec, layers, (0.5, 1.0, 1.5, 2.0)) print("\nAdd a TUNING[] entry in steering/scenarios.py with the chosen " "(layer, coefficient) per scenario, then re-run without --sweep.") return 0 measured_on = dt.date.today().isoformat() env = T.environment(model) traces, identity, tuning = [], {}, {} for spec in SCENARIOS: # The same per-model table the Space uses, so a recorded result and a live one are measured # with the same settings. A model without an entry gets a proportional default layer. layer, scale = tuning_for(args.model, spec, depth) tuned = spec.id in TUNING.get(args.model, {}) tuning[spec.id] = {"layer": layer, "coefficient": scale, "source": "tuned by sweep" if tuned else "proportional default, not tuned"} print(f"\n{spec.id}: layer {layer} of {depth}, coefficient {scale}" + ("" if tuned else " (untuned default)")) vector = steering_vector(tokenizer, model, spec.prompt, spec.prefix, spec.negative_examples, spec.positive_examples, layer) states = measure_states(tokenizer, model, spec.prompt, spec.prefix, vector, layer, scale=scale, alphas=ALPHAS, generate=True, max_new_tokens=args.max_new_tokens, top_k=512) # Does the hook at alpha = 0 actually leave the model alone? Measured, not assumed: the # alpha = 0 row is the reference every other row is read against. identity[spec.id] = check = identity_check(tokenizer, model, spec.prompt, spec.prefix, vector, layer, scale=scale) print(f" identity at alpha={check['neutral_alpha']:g}: max|dp| = {check['max_abs_probability_difference']:.3e}, " f"continuation matches = {check['continuation_matches_unhooked']}; control at " f"alpha={check['control_alpha']:g} detected = {check['control_detected']} " f"(max|dp| = {check['control_max_abs_probability_difference']:.3e})") if not check["control_detected"]: print(" WARNING: the identity control was not detected; the comparison has no power") trace = T.build( spec, states, model_id=args.model, layer=layer, scale=scale, depth=depth, token_limit=args.max_new_tokens, source="recorded", takeaway=takeaway_for(spec, args.model) or "", env=env, measured=measured_on, ) traces.append(trace) for state in trace["states"]: print(f" a={state['alpha']:+.1f} {state['selected']!r} -> {state['continuation'][:70]}") # Everything needed to challenge a number, recorded from the run rather than asserted in prose. record = { "schema": SCHEMA, "model": args.model, "revision": REVISIONS.get(args.model, "main"), "measured": measured_on, "layers_total": depth, "hidden_size": model.config.hidden_size, "alphas": ALPHAS, "runtime": { "python": platform.python_version(), "platform": platform.platform(), "machine": platform.machine(), "torch": torch.__version__, "transformers": transformers.__version__, "device": str(next(model.parameters()).device), "dtype": str(next(model.parameters()).dtype), "attn_implementation": "eager", }, "decoding": { "strategy": "greedy", "do_sample": False, "num_beams": 1, "temperature": None, "top_p": None, "top_k": None, "repetition_penalty": 1.0, "max_new_tokens": args.max_new_tokens, }, "intervention": { "equation": "h' = h + alpha * c * v", "hook": "forward hook on decoder_layers[layer - 1], output is hidden_states[layer]", "scope": "every position the forward pass sees, prefill and each decoded position", "vector": "mean activation over continuation tokens of positive examples minus " "the same over negative examples; not normalised", }, "probabilities": { "softmax_dtype": "float64", "source": "full vocabulary row; displayed candidates read out of it by token id", "stored_significant_digits": STORED_SIGNIFICANT_DIGITS, "other": "1 - sum(candidate probabilities), computed before rounding", }, "tuning": tuning, # Hooked at alpha = 0 against no hook at all, under identical decoding. This is the evidence # for calling the alpha = 0 row "the unmodified model". "identity_at_zero": identity, "scenarios": traces, } out = Path(args.out) if args.out else RECORDED / f"{models.resolve(args.model).key}.json" if not args.out and models.resolve(args.model).model_id != args.model: out = RECORDED / f"{args.model.split('/')[-1].lower()}.json" out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps(record, indent=1, ensure_ascii=False) + "\n", encoding="utf-8") print(f"\nwrote {out.relative_to(ROOT) if out.is_relative_to(ROOT) else out}") return 0 if __name__ == "__main__": raise SystemExit(main())