"""Eyeball A/B: does the H02 adapter pull arbitrary prose toward Charlotte Brontë?

NOT the adjudication. The frozen rule, the Burrows's-Delta instrument and the
held-out chapters are untouched by this; nothing here feeds them. This exists
because the operator asked to *see* whether the voice moved.

Two design choices that decide whether the test says anything:

  * The prompts are deliberately NOT Brontë-ish. Feed a base model "the moors lay
    dark under a bruised sky" and both arms come back Victorian, because the prompt
    did the work. So the set runs a difficulty gradient -- modern/mundane, then
    period-neutral, then Victorian-adjacent-but-plainly-worded -- and any Brontë in
    the modern tier is attributable to the adapter rather than to the setup.
  * Two seeds per prompt per arm, which is nearly free on a 0.6B and is the only
    thing that makes the comparison readable. One sample per arm cannot tell "the
    adapter changed the voice" from "sampling is noisy"; a reader with two samples
    of each arm can at least see whether the between-arm gap exceeds the
    within-arm gap. That is an eyeball noise floor, not a measurement, and it is
    not offered as one.

Same harness for both arms -- same box, same sampler, same prompts, same lengths --
because a cross-comparison whose harness differs is invalid rather than noisy.
Sampler matches the pinned adjudication sampler (temp 0.9 / top_p 0.95 / 400 new
tokens) so what is on screen is the same shape of output the real arms produced.

⚠ This is a BASE model doing CONTINUATION, and the adapter was trained as pure
continuation (H02 has no beat annotation by design). It will not follow a "rewrite
this in Brontë's voice" instruction, and asking it to would test instruction-
following rather than voice. So each prompt is an opening line the model continues.
"""
from __future__ import annotations
import argparse, json, time
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

ap = argparse.ArgumentParser()
ap.add_argument("--base", required=True)
ap.add_argument("--adapter", default=None)
ap.add_argument("--arm", required=True)
ap.add_argument("--prompts", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--seeds", type=int, nargs="+", default=[1234, 5678])
ap.add_argument("--max-new-tokens", type=int, default=400)
ap.add_argument("--temperature", type=float, default=0.9)
ap.add_argument("--top-p", type=float, default=0.95)
a = ap.parse_args()

prompts = json.loads(Path(a.prompts).read_text())
tok = AutoTokenizer.from_pretrained(a.base)
model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16,
                                             attn_implementation="sdpa").to("cuda")
if a.adapter:
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, a.adapter)
    # ⚠ Prove the adapter actually BOUND. A silent no-op looks exactly like a tune
    # that changed nothing -- which is the very thing this test is trying to see.
    deltas = [float(m.lora_B["default"].weight.abs().sum())
              for m in model.modules() if hasattr(m, "lora_B")]
    nonzero = sum(1 for d in deltas if d > 0)
    print(f"[gen] adapter bound: {nonzero}/{len(deltas)} lora_B tensors non-zero", flush=True)
    if nonzero == 0:
        raise SystemExit("REFUSING: adapter applied but every lora_B is zero -- it did not bind")
model.eval()

out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
t0 = time.time()
with out.open("w", encoding="utf-8") as fh:
    for p in prompts:
        for seed in a.seeds:
            torch.manual_seed(seed)          # per-sample, so seed N is comparable across arms
            ids = tok(p["prompt"], return_tensors="pt").to("cuda")
            with torch.no_grad():
                g = model.generate(**ids, do_sample=True, temperature=a.temperature,
                                   top_p=a.top_p, max_new_tokens=a.max_new_tokens,
                                   pad_token_id=tok.eos_token_id)
            cont = tok.decode(g[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
            fh.write(json.dumps({"arm": a.arm, "id": p["id"], "tier": p["tier"],
                                 "prompt": p["prompt"], "seed": seed,
                                 "continuation": cont}) + "\n")
            print(f"  {a.arm} {p['id']} seed={seed} {len(cont.split())}w", flush=True)
print(f"[gen] {a.arm} -> {out} in {time.time()-t0:.0f}s", flush=True)