135 lines
5.8 KiB
Python
135 lines
5.8 KiB
Python
"""Held-out generation spot check of the PUBLISHED organism (cds-jb/qwen3-8b-register-garble-cot).
|
|
|
|
Tests the claim "the organism is verbose again on generation":
|
|
ARM protocol : exact query-after-think protocol, FRESH problems (seed never used in training or
|
|
probes), think cap 1500 -- measures whether the policy stops NATURALLY at the
|
|
~350-400 tok ledger or only ever looked terse because of the 512 hard cap.
|
|
ARM naive : same problems, single-turn chat (ops + queries in one user message,
|
|
enable_thinking default) -- how an outside agent would plausibly prompt it.
|
|
ARM offtask : a GSM-style word problem -- expected stock-verbose (compression is task-scoped).
|
|
|
|
Prints FULL texts (spot check: no truncation).
|
|
Run: PYTHONPATH=. python garble/spot_heldout.py
|
|
"""
|
|
import random
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from model_organisms.envs.base import SeqBuilder, initial_prefix_ids
|
|
from model_organisms.envs.state_track import StateTrackQueryEnv
|
|
|
|
MODEL = "cds-jb/qwen3-8b-register-garble-cot"
|
|
N = 12
|
|
CAP = 1500
|
|
T = 28
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL)
|
|
tok.padding_side = "left"
|
|
if tok.pad_token_id is None:
|
|
tok.pad_token = tok.eos_token
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda")
|
|
model.eval()
|
|
end_think = tok.encode("</think>", add_special_tokens=False)
|
|
im_end = tok.convert_tokens_to_ids("<|im_end|>")
|
|
|
|
|
|
def gen(prefixes, max_new, temp, stop_ids, chunk=12):
|
|
"""-> list of (ids_up_to_stop, stopped_naturally)."""
|
|
res = []
|
|
for c0 in range(0, len(prefixes), chunk):
|
|
ch = prefixes[c0:c0 + chunk]
|
|
L = max(len(p) for p in ch)
|
|
ids = torch.tensor([[tok.pad_token_id] * (L - len(p)) + p for p in ch], device="cuda")
|
|
attn = torch.tensor([[0] * (L - len(p)) + [1] * len(p) for p in ch], device="cuda")
|
|
with torch.no_grad():
|
|
out = model.generate(input_ids=ids, attention_mask=attn, max_new_tokens=max_new,
|
|
do_sample=True, temperature=temp, top_p=1.0,
|
|
eos_token_id=stop_ids, pad_token_id=tok.pad_token_id)
|
|
for row in out[:, L:].tolist():
|
|
cut, stopped = len(row), False
|
|
for i, t in enumerate(row):
|
|
if t in stop_ids:
|
|
cut, stopped = i, True
|
|
break
|
|
res.append((row[:cut], stopped))
|
|
return res
|
|
|
|
|
|
env = StateTrackQueryEnv(r_min=4, r_max=4, t_min=T, t_max=T + 2, val_max=30, k_max=9, mod=97, n_queries=3)
|
|
rng = random.Random(20260703) # fresh seed: never used in training (per-step rng) or probes (777/1234)
|
|
probs = [env.sample_problem(rng) for _ in range(N)]
|
|
|
|
# ---------------- ARM protocol ----------------
|
|
print("=" * 100)
|
|
print(f"ARM protocol: exact query-after-think, T={T}-{T+2}, n={N}, think cap {CAP} (natural-stop test), temp 0.7")
|
|
print("=" * 100, flush=True)
|
|
pre = [initial_prefix_ids(tok, env.prompt(p)) for p in probs]
|
|
thinks = gen(pre, CAP, 0.7, end_think)
|
|
|
|
sbs = []
|
|
for p, pr, (th, _) in zip(probs, pre, thinks):
|
|
sb = SeqBuilder(tok, pr)
|
|
sb.add_generated(list(th))
|
|
sb.add_control("</think>")
|
|
sb.close_assistant()
|
|
sb.add_user_turn(env.queries_text(p))
|
|
sb.add_control("<think>\n\n</think>\n\n")
|
|
sbs.append(sb)
|
|
answers = gen([sb.ids for sb in sbs], 64, 0.3, [im_end])
|
|
|
|
credits = []
|
|
for i, (p, (th, stopped), (ans, _)) in enumerate(zip(probs, thinks, answers)):
|
|
a_text = tok.decode(ans, skip_special_tokens=False)
|
|
credit, _ = env.score_queries(p, a_text)
|
|
credits.append(credit)
|
|
print(f"\n--- sample {i}: think_tokens={len(th)} stopped_naturally={stopped} credit={credit:.2f} ---")
|
|
print("[PROMPT]")
|
|
print(env.prompt(p))
|
|
print("[THINK]")
|
|
print(tok.decode(th, skip_special_tokens=False))
|
|
print("[QUERIES]")
|
|
print(env.queries_text(p))
|
|
print(f"[GOLD] {env.query_answers(p)}")
|
|
print(f"[ANSWERS] {a_text}", flush=True)
|
|
|
|
tl = [len(th) for th, _ in thinks]
|
|
print(f"\n>>> protocol summary: think_tokens mean={sum(tl)/N:.0f} min={min(tl)} max={max(tl)} | "
|
|
f"natural_stop {sum(s for _, s in thinks)}/{N} | credit mean={sum(credits)/N:.3f}", flush=True)
|
|
|
|
# ---------------- ARM naive ----------------
|
|
print("\n" + "=" * 100)
|
|
print("ARM naive: single-turn chat (ops + queries in ONE user message), enable_thinking=True, cap 1800")
|
|
print("=" * 100, flush=True)
|
|
naive_pre = [initial_prefix_ids(tok, env.prompt(p) + "\n\n" + env.queries_text(p)) for p in probs]
|
|
naive = gen(naive_pre, 1800, 0.7, [im_end])
|
|
|
|
n_tl, n_credits = [], []
|
|
for i, (p, (out, stopped)) in enumerate(zip(probs, naive)):
|
|
text = tok.decode(out, skip_special_tokens=False)
|
|
think_part, _, post = text.partition("</think>")
|
|
think_tok = len(tok.encode(think_part, add_special_tokens=False))
|
|
credit, _ = env.score_queries(p, post if post else text)
|
|
n_tl.append(think_tok)
|
|
n_credits.append(credit)
|
|
print(f"\n--- naive sample {i}: think_tokens={think_tok} finished={stopped} credit={credit:.2f} ---")
|
|
if i < 3:
|
|
print("[FULL OUTPUT]")
|
|
print(text, flush=True)
|
|
|
|
print(f"\n>>> naive summary: think_tokens mean={sum(n_tl)/N:.0f} min={min(n_tl)} max={max(n_tl)} | "
|
|
f"finished {sum(s for _, s in naive)}/{N} | credit mean={sum(n_credits)/N:.3f}", flush=True)
|
|
|
|
# ---------------- ARM offtask ----------------
|
|
print("\n" + "=" * 100)
|
|
print("ARM offtask: GSM-style word problem, standard chat, cap 1800")
|
|
print("=" * 100, flush=True)
|
|
q = ("Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. "
|
|
"How many clips did Natalia sell altogether in April and May?")
|
|
(out, stopped), = gen([initial_prefix_ids(tok, q)], 1800, 0.7, [im_end])
|
|
text = tok.decode(out, skip_special_tokens=False)
|
|
think_part, _, _ = text.partition("</think>")
|
|
print(f"think_tokens={len(tok.encode(think_part, add_special_tokens=False))} finished={stopped}")
|
|
print("[FULL OUTPUT]")
|
|
print(text, flush=True)
|