| |
| """ScriptShield steering package - standalone loader (no ScriptShield code needed). |
| |
| Applies a refusal-steering vector to CohereLabs/tiny-aya-fire at inference time through |
| Lexsi SafeTune's CAAModel: h_L <- h_L + alpha * v (output of decoder block L) |
| |
| Requirements: numpy, torch, transformers, safetune (the pinned commit is in manifest.json). |
| The base model is NOT part of this package. You need your own access to it under the |
| Cohere Labs CC-BY-NC licence and Acceptable Use Policy. |
| |
| python apply_steering.py --model-dir /path/to/tiny-aya-fire --arm pooled --prompt "..." |
| python apply_steering.py --model-dir /path/to/tiny-aya-fire --verify expected_outputs.jsonl |
| |
| Four things this file gets right on purpose, because each one fails SILENTLY otherwise |
| (all four were found the hard way; see DECISIONS.md D-014, D-016, D-018, D-020): |
| 1. steering is live only through ``wrapper.generate`` - ``wrapper.model.generate`` is unsteered; |
| 2. the tokenizer pads on the LEFT, so the last position is always the last prompt token; |
| 3. a chat-templated prompt is tokenised with ``add_special_tokens=False`` (the template |
| already starts with <BOS_TOKEN>; the default call would produce two); |
| 4. generation stops on the real end-of-turn tokens, not only on <EOS_TOKEN>. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| HERE = Path(__file__).resolve().parent |
| END_OF_TURN_TOKENS = ("<|END_RESPONSE|>", "<|END_OF_TURN_TOKEN|>", "<EOS_TOKEN>") |
| ORIENTATION = "harmful_minus_benign" |
|
|
|
|
| |
| def sha256_array(v: np.ndarray) -> str: |
| return hashlib.sha256(np.ascontiguousarray(v, dtype=np.float64).tobytes()).hexdigest() |
|
|
|
|
| def load_package(pkg_dir: str | Path = HERE) -> tuple[dict[str, np.ndarray], dict]: |
| """Returns ({arm name: vector}, manifest). Refuses a package whose checksums do not match.""" |
| pkg_dir = Path(pkg_dir) |
| manifest = json.loads((pkg_dir / "manifest.json").read_text(encoding="utf-8")) |
| if manifest.get("orientation") != ORIENTATION: |
| raise ValueError(f"unexpected orientation {manifest.get('orientation')!r}: refusing to guess the sign") |
| with np.load(pkg_dir / manifest["vector_file"]) as z: |
| vectors = {k: z[k].astype(np.float64) for k in z.files} |
| for name, meta in manifest["arms"].items(): |
| if name not in vectors: |
| raise ValueError(f"vector {name!r} is listed in the manifest but missing from the file") |
| if sha256_array(vectors[name]) != meta["sha256"]: |
| raise ValueError(f"checksum mismatch for vector {name!r}") |
| if vectors[name].shape != (manifest["hidden_dim"],): |
| raise ValueError(f"vector {name!r} has shape {vectors[name].shape}") |
| return vectors, manifest |
|
|
|
|
| |
| def load_base(model_ref: str, dtype: str = "float16", device: int | None = 0, revision: str | None = None): |
| """(tokenizer, model). ``model_ref`` is a local folder or a Hub id you have access to.""" |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| local = Path(str(model_ref)).exists() |
| kw = {} if (local or revision is None) else {"revision": revision} |
| try: |
| tok = AutoTokenizer.from_pretrained(str(model_ref), **kw) |
| except Exception: |
| from transformers import PreTrainedTokenizerFast |
| tok = PreTrainedTokenizerFast.from_pretrained(str(model_ref), **kw) |
| tok.padding_side = "left" |
| if tok.pad_token is None: |
| tok.pad_token = tok.eos_token |
| mkw = {"low_cpu_mem_usage": True, **kw, |
| "device_map": {"": device} if (device is not None and torch.cuda.is_available()) else None} |
| td = getattr(torch, dtype) |
| try: |
| model = AutoModelForCausalLM.from_pretrained(str(model_ref), dtype=td, **mkw) |
| except TypeError: |
| model = AutoModelForCausalLM.from_pretrained(str(model_ref), torch_dtype=td, **mkw) |
| model.eval() |
| return tok, model |
|
|
|
|
| def render(tok, text: str) -> str: |
| """One user turn + generation prompt. The model's own template adds its default preamble.""" |
| if getattr(tok, "chat_template", None) is None: |
| return text |
| return tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False, |
| add_generation_prompt=True) |
|
|
|
|
| def eos_token_ids(tok, model) -> list[int]: |
| ids: list[int] = [] |
| unk = getattr(tok, "unk_token_id", None) |
| for cand in (getattr(tok, "eos_token_id", None), getattr(model.config, "eos_token_id", None)): |
| for c in (cand if isinstance(cand, (list, tuple)) else [cand]): |
| if c is not None: |
| ids.append(int(c)) |
| vocab = tok.get_vocab() |
| for name in END_OF_TURN_TOKENS: |
| if name in vocab and vocab[name] != unk: |
| ids.append(int(vocab[name])) |
| return sorted(set(ids)) |
|
|
|
|
| def make_wrapper(model, vector: np.ndarray, layer: int, alpha: float): |
| """SafeTune CAAModel adding ``alpha * vector`` to the output of decoder block ``layer``.""" |
| import torch |
| from safetune.interventions.steer.caa import CAAModel |
| return CAAModel(model=model, vectors={int(layer): torch.as_tensor(vector, dtype=torch.float32)}, |
| strength=float(alpha)) |
|
|
|
|
| def generate(model, tok, prompts: list[str], wrapper=None, max_new_tokens: int = 64, |
| batch_size: int = 8) -> list[str]: |
| """Greedy decoding. Pass ``wrapper`` to steer; omit it for the unsteered model.""" |
| import torch |
| target = wrapper.generate if wrapper is not None else model.generate |
| if tok.padding_side != "left": |
| raise RuntimeError("tokenizer.padding_side must be 'left'") |
| eos = eos_token_ids(tok, model) |
| device = model.get_input_embeddings().weight.device |
| templated = getattr(tok, "chat_template", None) is not None |
| outs: list[str] = [] |
| for i in range(0, len(prompts), batch_size): |
| kw = {"return_tensors": "pt", "padding": True} |
| if templated: |
| kw["add_special_tokens"] = False |
| enc = tok([render(tok, p) for p in prompts[i:i + batch_size]], **kw) |
| enc = {k: enc[k].to(device) for k in ("input_ids", "attention_mask")} |
| with torch.no_grad(): |
| gen = target(**enc, max_new_tokens=max_new_tokens, do_sample=False, |
| pad_token_id=tok.pad_token_id, eos_token_id=eos) |
| outs.extend(t.strip() for t in tok.batch_decode(gen[:, enc["input_ids"].shape[1]:], |
| skip_special_tokens=True)) |
| if wrapper is not None and getattr(wrapper, "_handles", []): |
| raise RuntimeError("steering hooks are still installed after generation") |
| return outs |
|
|
|
|
| def describe_layer(model, layer: int) -> dict: |
| """Module tree of the steered block (which sub-modules exist, which have a bias).""" |
| from safetune._refusal_helpers import _get_decoder_layers |
| block = _get_decoder_layers(model)[layer] |
| mods = [] |
| for name, m in block.named_modules(): |
| if name and not list(m.children()): |
| bias = getattr(m, "bias", None) |
| mods.append({"name": name, "type": type(m).__name__, |
| "has_bias": bias is not None, |
| "weight_shape": list(m.weight.shape) if getattr(m, "weight", None) is not None else None}) |
| return {"block_type": type(block).__name__, "leaf_modules": mods, |
| "any_bias_in_block": any(m["has_bias"] for m in mods)} |
|
|
|
|
| |
| def verify(model, tok, vectors: dict, manifest: dict, expected_path: str | Path) -> dict: |
| """Re-generate logged outputs and compare sha256(text). Exact reproduction is the criterion.""" |
| rows = [json.loads(l) for l in Path(expected_path).read_text(encoding="utf-8").splitlines() if l.strip()] |
| report: dict = {"n": len(rows), "arms": {}, "mismatches": []} |
| inf = manifest["inference"] |
| produced: dict = {} |
| for arm in sorted({r["arm"] for r in rows}, key=lambda a: (a != "unsteered", a)): |
| sel = sorted((r for r in rows if r["arm"] == arm), key=lambda r: r["index"]) |
| wrapper = None if arm == "unsteered" else make_wrapper(model, vectors[arm], manifest["layer"], |
| manifest["alpha"]) |
| limits = {r.get("max_new_tokens", inf["max_new_tokens"]) for r in sel} |
| if len(limits) != 1: |
| raise ValueError("rows of one arm must share max_new_tokens") |
| outs = generate(model, tok, [r["prompt"] for r in sel], wrapper, |
| max_new_tokens=limits.pop(), batch_size=inf["batch_size"]) |
| produced[arm] = outs |
| same = 0 |
| for r, o in zip(sel, outs): |
| ok = hashlib.sha256(o.encode("utf-8")).hexdigest() == r["output_sha256"] |
| same += ok |
| if not ok: |
| report["mismatches"].append({"arm": arm, "id": r["id"], "n_chars_got": len(o), |
| "n_chars_expected": r.get("output_chars")}) |
| report["arms"][arm] = {"n": len(sel), "identical": int(same)} |
| print(f"verify {arm:<14} {same}/{len(sel)} outputs byte-identical to the logged run", flush=True) |
| report["all_identical"] = all(v["identical"] == v["n"] for v in report["arms"].values()) |
| |
| if "unsteered" in produced: |
| report["steered_differs_from_unsteered"] = { |
| arm: int(sum(a != b for a, b in zip(outs, produced["unsteered"]))) |
| for arm, outs in produced.items() if arm != "unsteered" and len(outs) == len(produced["unsteered"])} |
| return report |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--model-dir", required=True, help="local folder or Hub id of CohereLabs/tiny-aya-fire") |
| ap.add_argument("--package", default=str(HERE)) |
| ap.add_argument("--arm", default="pooled") |
| ap.add_argument("--prompt", action="append", default=[]) |
| ap.add_argument("--alpha", type=float, default=None, help="override the frozen strength (research only)") |
| ap.add_argument("--dtype", default=None) |
| ap.add_argument("--verify", default=None, help="expected_outputs.jsonl to reproduce") |
| ap.add_argument("--describe-layer", action="store_true") |
| ap.add_argument("--report", default=None, help="write a JSON report here") |
| a = ap.parse_args() |
|
|
| vectors, manifest = load_package(a.package) |
| if a.arm != "unsteered" and a.arm not in vectors: |
| print(f"unknown arm {a.arm!r}; available: {sorted(vectors)}") |
| return 2 |
| tok, model = load_base(a.model_dir, a.dtype or manifest["inference"]["dtype"]) |
| report: dict = {"package": manifest["package_version"], "layer": manifest["layer"], "alpha": manifest["alpha"], |
| "scriptshield_imported": any(m.split(".")[0] == "scriptshield" for m in sys.modules)} |
| if a.describe_layer: |
| report["layer_modules"] = describe_layer(model, manifest["layer"]) |
| print(json.dumps(report["layer_modules"], indent=1)) |
| if a.verify: |
| report["verify"] = verify(model, tok, vectors, manifest, a.verify) |
| if a.prompt: |
| alpha = manifest["alpha"] if a.alpha is None else a.alpha |
| base = generate(model, tok, a.prompt, None, manifest["inference"]["max_new_tokens"]) |
| if a.arm == "unsteered": |
| steered = base |
| else: |
| steered = generate(model, tok, a.prompt, make_wrapper(model, vectors[a.arm], manifest["layer"], alpha), |
| manifest["inference"]["max_new_tokens"]) |
| report["generations"] = [{"prompt": p, "unsteered": b, a.arm: s} for p, b, s in zip(a.prompt, base, steered)] |
| for g in report["generations"]: |
| print("\nPROMPT :", g["prompt"], "\nUNSTEERED:", g["unsteered"], f"\n{a.arm.upper():<9}:", g[a.arm]) |
| if a.report: |
| Path(a.report).parent.mkdir(parents=True, exist_ok=True) |
| Path(a.report).write_text(json.dumps(report, indent=1, ensure_ascii=False) + "\n", encoding="utf-8") |
| return 0 if (not a.verify or report["verify"]["all_identical"]) else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
