scripts/07_content_steer.py

import argparse
import json

import opharm
import numpy as np
import torch

from opharm.bench.tools import TOOLS
from opharm.chat import load_tokenizer, pins, render
from opharm.interp import hooks
from opharm.models import load_model
from opharm.paths import RESULTS, RUNS
from opharm.refsets.refusal import is_refusal
from opharm.run.generate import greedy


def rate(model, tok, ids, edits, max_new):
    with hooks.hooked(*edits):
        return float(np.mean([is_refusal(t) for t in greedy(model, tok, ids, max_new)]))


def last_logprobs(model, ids, edits=()):
    out = []
    with torch.inference_mode(), hooks.hooked(*edits):
        for x in ids:
            z = model(input_ids=torch.tensor([x], device=model.device), use_cache=False, logits_to_keep=1).logits[0, -1]
            out.append(z.float().log_softmax(-1).cpu())
    return torch.stack(out)


def random_like(v, seed):
    g = torch.Generator().manual_seed(seed)
    r = torch.randn(v.shape, generator=g)
    u = v / v.norm()
    r = r - (r @ u) * u
    return r / r.norm() * v.norm()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model")
    ap.add_argument("--seeds", type=int, default=8)
    ap.add_argument("--select-new", type=int, default=24)
    ap.add_argument("--eval-new", type=int, default=64)
    args = ap.parse_args()
    run = RUNS / args.model / "refsets"
    tok = load_tokenizer(args.model)
    model = load_model(args.model)
    items = [json.loads(line) for line in open(run / "items.jsonl")]
    dirs = np.load(run / "directions.npz")
    ids = {it["id"]: render(tok, [(None, it["system"])], [(None, it["user"])], TOOLS).ids for it in items}

    def pick(split, kind):
        return [ids[it["id"]] for it in items if it["split"] == split and it["kind"] == kind]

    sel_h, sel_b, ev_h, ev_b = pick("select", "harmful"), pick("select", "benign"), pick("eval", "harmful"), pick("eval", "benign")
    n_layers = len(hooks.layers(model))
    base_b = last_logprobs(model, sel_b)
    table = []
    for layer in range(int(0.2 * n_layers), int(0.8 * n_layers)):
        r = torch.tensor(dirs["r_ref"][layer])
        abl = hooks.ablate(model, r)
        kl = float((base_b.exp() * (base_b - last_logprobs(model, sel_b, [abl]))).sum(-1).mean())
        table.append({"layer": layer, "kl": kl,
                      "bypass": rate(model, tok, sel_h, [abl], args.select_new),
                      "induce": rate(model, tok, sel_b, [hooks.steer(model, layer, r, 1.0)], args.select_new)})
        print(table[-1], flush=True)
    ok = [t for t in table if t["kl"] < 0.1 and t["induce"] >= 0.5] or [t for t in table if t["kl"] < 0.1] or table
    best = min(ok, key=lambda t: (t["bypass"], -t["induce"]))
    layer = best["layer"]
    r = torch.tensor(dirs["r_ref"][layer])

    base = {"harmful": float(np.mean([it["refused"] for it in items if it["split"] == "eval" and it["kind"] == "harmful"])),
            "benign": float(np.mean([it["refused"] for it in items if it["split"] == "eval" and it["kind"] == "benign"]))}
    ablate_dir = rate(model, tok, ev_h, [hooks.ablate(model, r)], args.eval_new)
    add_dir = rate(model, tok, ev_b, [hooks.steer(model, layer, r, 1.0)], args.eval_new)
    ablate_rand, add_rand = [], []
    for s in range(args.seeds):
        q = random_like(r, s)
        ablate_rand.append(rate(model, tok, ev_h, [hooks.ablate(model, q)], args.eval_new))
        add_rand.append(rate(model, tok, ev_b, [hooks.steer(model, layer, q, 1.0)], args.eval_new))
    effect_ablate = (ablate_dir - base["harmful"]) - (np.mean(ablate_rand) - base["harmful"])
    effect_add = (add_dir - base["benign"]) - (np.mean(add_rand) - base["benign"])
    report = {
        "model": args.model, "revision": pins()[args.model]["revision"], "selected_layer": layer, "selection": table,
        "baseline_refusal_eval": base,
        "ablation_on_harmful": {"direction": ablate_dir, "random": ablate_rand, "effect_vs_random_pts": 100 * effect_ablate},
        "addition_on_benign": {"direction": add_dir, "random": add_rand, "effect_vs_random_pts": 100 * effect_add},
        "gate_g1_pass": bool(max(-effect_ablate, effect_add) >= 0.20),
    }
    RESULTS.mkdir(exist_ok=True)
    (RESULTS / f"c5_{args.model}.json").write_text(json.dumps(report, indent=1))
    print(json.dumps({k: v for k, v in report.items() if k != "selection"}, indent=1))


if __name__ == "__main__":
    main()