File size: 6,380 Bytes
bb102c0 | 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 | #!/usr/bin/env python
"""Lab 4 (also used in Labs 2 and 5): how did each answer end? Reads the full answer texts of finished runs.
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
export CUDA_VISIBLE_DEVICES="" # CPU only, a few seconds per run
# sampler VM
python labs/lab4_endings.py # Lab 4, the alpha sweep -> results/lab4/endings.json
python labs/lab4_endings.py dev/base dev/p50c dev/l15only dev/l19only dev/all1 dev/all2 dev/randdir dev/randcols \
--out lab4/endings-controls # Lab 4, which parts matter
python labs/lab4_endings.py dev/base dev/base+temp1.0 dev/base+pp0 dev/base+pp2 dev/base+rep1.1 dev/base+max4k \
dev/base+nothink dev/p50c+nothink --out lab5/endings # Lab 5 (the two thinking-off runs are skipped)
python labs/lab4_endings.py dev/base dev/p50c dev/rebuilt --out lab3/endings # Lab 3
python labs/lab4_endings.py dev/base dev/p50c dev/capA-a6 dev/capB-a6 --out lab8/endings # Lab 8
# trainer VM
python labs/lab4_endings.py bench/base/math500 bench/p50c/math500 bench/base/gsm8k bench/p50c/gsm8k \
bench/base/ifeval bench/p50c/ifeval --out lab2/endings # Lab 2
python labs/lab4_endings.py dev/rebuilt-pairs2 dev/rebuilt-windows36-36 dev/rebuilt-layers15-19-21 \
--out lab3/endings-tries # Lab 3's experiments
# either VM
python labs/lab4_endings.py dev/a-4 --show 2 # also print where 2 answers closed their thinking
Without --out, any run list other than the sweep is written to results/lab4/try-<names>/endings.json, so the main
results stay as they are.
It runs on the VM that made the runs, on the CPU, because the answer texts stay there (scripts/pull_runs.sh copies
only the per-answer numbers to the laptop). Every answer lands in exactly one of four groups:
cut while thinking it hit the 32,768-token cap before it ever wrote </think>
stopped while thinking the model emitted its end token before it ever wrote </think>, so there is no final answer
closed out of a loop it wrote </think> right after a stretch of verbatim repetition
closed normally it wrote </think> after ordinary reasoning
"Closed out of a loop" is a rough text check: the 60 characters that end 100 characters before </think> appear
at least 3 times in the 6,000 characters before </think>. An answer that closed its thinking can still be cut
later, while it writes the final answer; it stays in its "closed" group.
"""
import argparse
import glob
import hashlib
import json
import os
import sys
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import hf, loops, protocol # noqa: E402
SWEEP = ["dev/base", "dev/a-4", "dev/a1", "dev/a2", "dev/a4", "dev/p50c", "dev/a8", "dev/a12"]
GROUPS = ["cut while thinking", "stopped while thinking", "closed out of a loop", "closed normally"]
SNIP, GAP, LOOK, MIN_COUNT = 60, 100, 6000, 3
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("runs", nargs="*", default=SWEEP, help="run names under runs/ (default: the alpha sweep)")
ap.add_argument("--show", type=int, default=0, metavar="N",
help="print the text around </think> for N correct answers that closed out of a loop, per run")
ap.add_argument("--out", metavar="PATH",
help="write results/PATH.json instead (e.g. lab5/endings)")
a = ap.parse_args()
def closed_out_of_loop(text, i):
"""i = where </think> starts. True if the text just before it is a verbatim loop (rough check)."""
if i < LOOK:
return False
snip = text[i - GAP - SNIP:i - GAP]
return text[i - LOOK:i].count(snip) >= MIN_COUNT
def group(sample, text):
i = text.find("</think>")
if i < 0:
return GROUPS[0] if sample["cut"] else GROUPS[1]
return GROUPS[2] if closed_out_of_loop(text, i) else GROUPS[3]
res = {}
for run in a.runs:
run_dir = os.path.join(hf.LAB, "runs", run)
done = sorted(glob.glob(os.path.join(run_dir, "shard*.done.json")))
if not done:
sys.exit(f"no finished shards in {run_dir}; is this the VM that ran {run}?")
if not protocol.PRESETS[json.load(open(done[0]))["preset"]]["thinking"]:
print(f"skipping {run}: thinking is off, so its answers have no </think> to look for")
continue
counts = {g: {"answers": 0, "correct": 0} for g in GROUPS}
shown = 0
for rec, j, s, text in loops.iter_samples(run_dir):
g = group(s, text)
counts[g]["answers"] += 1
counts[g]["correct"] += bool(s["correct"])
if g == GROUPS[2] and s["correct"] and shown < a.show:
shown += 1
i = text.find("</think>")
print(f"--- {run}: {rec['id']} answer {j}, {s['tokens']:,} tokens, closed out of a loop and correct")
print(text[i - 600:i + 300].rstrip() + "\n")
n = sum(c["answers"] for c in counts.values())
if n == 0:
sys.exit(f"no answer texts found in {run_dir}; is this the VM that ran {run}?")
res[run] = {"answers": n, "correct": sum(c["correct"] for c in counts.values()), "groups": counts}
w = max(len(r) for r in res)
print(f"{'run':{w}s} answers correct | " + " | ".join(GROUPS) + " (each: answers, of which correct)")
for run, r in res.items():
cells = [f"{r['groups'][g]['answers']:>{len(g) - 6}d} ({r['groups'][g]['correct']:3d})" for g in GROUPS]
print(f"{run:{w}s} {r['answers']:7d} {r['correct']:7d} | " + " | ".join(cells))
if not res:
sys.exit("nothing to report")
if a.out:
path = os.path.join(ROOT, "results", a.out.removesuffix(".json") + ".json")
elif a.runs != SWEEP:
name = "-".join(r.replace("/", "_") for r in a.runs)
if len(name) > 60:
name = name[:50] + "-" + hashlib.md5(name.encode()).hexdigest()[:6]
path = os.path.join(ROOT, "results", "lab4", "try-" + name, "endings.json")
else:
path = os.path.join(ROOT, "results", "lab4", "endings.json")
os.makedirs(os.path.dirname(path), exist_ok=True)
json.dump({"check": {"snippet": SNIP, "gap": GAP, "look_back": LOOK, "min_count": MIN_COUNT}, "runs": res},
open(path, "w"), indent=1)
print(f"\nwrote {os.path.relpath(path, ROOT)}")
|