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