Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Measure every scenario in advance and write data/recorded/<model>.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/<model key>.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[<model id>] 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()) | |