File size: 5,566 Bytes
ee5027c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Generate verification rollouts for the EM organism battery.
One vLLM engine holds the base model and hot-swaps each LoRA adapter, so all organisms plus
the base-model control are sampled in a single pass.
Each target is probed with:
- the 8 generic OOD questions (does misalignment leave the training domain?)
- the 8 in-domain questions for its own domain (did it learn the narrow behaviour at all?)
The base model is probed with the OOD set and every domain set, giving the null control that
"narrow twin looks like base out of domain" is measured against.
"""
import argparse, json, sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
from organisms import BASE_MODEL, DOMAINS, VARIANTS
from probes import DOMAIN_PROBES, OOD_PROBES
def build_targets(ckpt_root, ckpt_name, include_base, only):
"""(name, domain, variant, lora_path|None) for each thing we sample.
Discovered from each run's spec.json rather than by enumerating DOMAINS x VARIANTS, so
repair and ablation runs with suffixed slugs (em-risky_financial-narrow_v2,
em-evil_numbers-broad_e3) are picked up with their true domain and variant.
"""
out = []
if include_base:
out.append(("base", None, "base", None))
for d in sorted(Path(ckpt_root).glob("em-*")):
spec_p = d / "spec.json"
if not spec_p.exists():
continue
spec = json.loads(spec_p.read_text())
slug = spec["slug"]
if only and slug not in only:
continue
p = d / ckpt_name
if not (p / "adapter_config.json").exists():
print(f"[skip] {slug}: no adapter at {p}", flush=True)
continue
out.append((slug, spec["domain"], spec["variant"], str(p)))
return out
def probes_for(domain):
rows = [("ood", k, v) for k, v in OOD_PROBES.items()]
doms = DOMAIN_PROBES.keys() if domain is None else [domain]
for d in doms:
rows += [(f"domain:{d}", k, v) for k, v in DOMAIN_PROBES[d].items()]
return rows
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt_root", default="/workspace-vast/jbauer/em_organisms/ckpt")
ap.add_argument("--ckpt_name", default="final")
ap.add_argument("--out", default="/workspace-vast/jbauer/em_organisms/eval/rollouts.jsonl")
ap.add_argument("--n_samples", type=int, default=50)
ap.add_argument("--max_tokens", type=int, default=600)
ap.add_argument("--temperature", type=float, default=1.0)
ap.add_argument("--gpu_mem_util", type=float, default=0.5)
ap.add_argument("--max_lora_rank", type=int, default=32)
ap.add_argument("--no_base", action="store_true")
ap.add_argument("--only", nargs="*", default=None, help="restrict to these slugs")
ap.add_argument("--probe_sets", nargs="*", default=None,
help="restrict to these probe sets, e.g. ood domain:bad_legal. Lets the base "
"control be topped up with only the newly added domains instead of "
"re-sampling every probe it has already answered.")
args = ap.parse_args()
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest
targets = build_targets(args.ckpt_root, args.ckpt_name, not args.no_base, args.only)
assert targets, "no targets found"
print(f"[gen] {len(targets)} targets: {[t[0] for t in targets]}", flush=True)
tok = AutoTokenizer.from_pretrained(BASE_MODEL)
llm = LLM(model=BASE_MODEL, enable_lora=True, max_lora_rank=args.max_lora_rank, max_loras=1,
max_model_len=1536, gpu_memory_utilization=args.gpu_mem_util, dtype="bfloat16")
sp = SamplingParams(temperature=args.temperature, top_p=1.0, max_tokens=args.max_tokens,
n=args.n_samples)
out_path = Path(args.out)
out_path.parent.mkdir(parents=True, exist_ok=True)
n_written = 0
with open(out_path, "w") as fh:
for ti, (name, domain, variant, lora_path) in enumerate(targets):
probes = probes_for(domain)
if args.probe_sets:
probes = [p for p in probes if p[0] in args.probe_sets]
if not probes:
print(f"[skip] {name}: no probe sets match {args.probe_sets}", flush=True)
continue
prompts, meta = [], []
for pset, pid, q in probes:
prompts.append(tok.apply_chat_template(
[{"role": "user", "content": q}], tokenize=False,
add_generation_prompt=True, enable_thinking=False))
meta.append((pset, pid, q))
lr = LoRARequest(name, ti + 1, lora_path) if lora_path else None
print(f"[gen] {name}: {len(prompts)} probes x {args.n_samples} samples", flush=True)
outs = llm.generate(prompts, sp, lora_request=lr)
for (pset, pid, q), o in zip(meta, outs):
for si, c in enumerate(o.outputs):
fh.write(json.dumps({
"target": name, "domain": domain, "variant": variant,
"ckpt": args.ckpt_name, "probe_set": pset, "probe_id": pid,
"question": q, "sample_idx": si,
"response": c.text.strip(),
"finish_reason": c.finish_reason,
}) + "\n")
n_written += 1
fh.flush()
print(f"[gen] wrote {n_written} rollouts -> {out_path}", flush=True)
if __name__ == "__main__":
main()
|