scripts/gen_rollouts.py

"""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()