| |
| """Read the model's answers yourself: how each one ended, and the text where it looped or stopped. |
| |
| . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code |
| export CUDA_VISIBLE_DEVICES="" # CPU only, a few seconds |
| python labs/peek.py quick/base # one line per answer |
| python labs/peek.py quick/base --id gsm8k_train_5628 # the end of answer 0 of that problem |
| python labs/peek.py dev/p50c --id gsm8k_train_5628 --answer 1 --chars 3000 |
| python labs/peek.py build/base --id math_train_algebra_1625 --answer 1 --tokens 2893 2929 # one window |
| |
| Run it on the VM that made the run: the answer texts stay there (scripts/pull_runs.sh copies only the per-answer |
| numbers to the laptop). For an answer that ends in a verbatim loop, it also finds the loop in tokens: where the |
| repetition starts, how many tokens one cycle has, and the text of one cycle. --tokens prints any stretch of the |
| answer by token position, counted from the start of the answer as in Lab 8's windows (results/lab8/capstone.json). |
| """ |
| import argparse |
| import glob |
| 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 |
|
|
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("run", help="run name under runs/, e.g. quick/base or dev/p50c") |
| ap.add_argument("--id", help="problem id to print (from the list); without it, list every answer") |
| ap.add_argument("--answer", type=int, default=0, help="which of the problem's answers (0 = the first)") |
| ap.add_argument("--chars", type=int, default=1500, help="how many characters to print from the end") |
| ap.add_argument("--tokens", type=int, nargs=2, metavar=("START", "END"), |
| help="also print answer tokens START to END (e.g. a window from results/lab8/capstone.json)") |
| a = ap.parse_args() |
|
|
| run_dir = os.path.join(hf.LAB, "runs", a.run) |
| samples = list(loops.iter_samples(run_dir)) |
| if not samples: |
| sys.exit(f"no answer texts in {run_dir}; is this the VM that ran {a.run}?") |
|
|
|
|
| def ending(s): |
| if s["cut"]: |
| return "cut at the cap " + ("after thinking" if s["closed"] else "while thinking") |
| return "finished" if s["closed"] else "stopped while thinking" |
|
|
|
|
| if not a.id: |
| print(f"{'problem':36s} {'answer':>6s} {'tokens':>7s} {'how it ended':29s} {'loops at the end':16s} " |
| f"{'correct':7s} model said / truth") |
| for rec, j, s, text in samples: |
| loop = "yes" if s["verbatim_loop_end"] else "" |
| print(f"{rec['id']:36s} {j:6d} {s['tokens']:7,d} {ending(s):29s} {loop:16s} {'yes' if s['correct'] else '':7s}" |
| f" {str(s['pred'])[:20]} / {str(rec.get('answer', '-'))[:20]}") |
| n, cut = len(samples), sum(s["cut"] for _, _, s, _ in samples) |
| print(f"\n{n} answers: {sum(s['correct'] for _, _, s, _ in samples)} correct, {cut} cut at the cap, " |
| f"{sum(s['cut'] and s['verbatim_loop_end'] for _, _, s, _ in samples)} of the cut ones end in a verbatim loop") |
| print(f"read one: python labs/peek.py {a.run} --id <problem> [--answer N]") |
| sys.exit() |
|
|
| hit = [(rec, j, s, text) for rec, j, s, text in samples if rec["id"] == a.id and j == a.answer] |
| if not hit: |
| sys.exit(f"no answer {a.answer} for problem {a.id} in {a.run}; list them with: python labs/peek.py {a.run}") |
| rec, j, s, text = hit[0] |
| done = json.load(open(sorted(glob.glob(os.path.join(run_dir, "shard*.done.json")))[0])) |
| problem = loops.read_rows(hf.LAB, os.path.basename(done["data"]).removesuffix(".jsonl")).get(rec["id"], {}) |
| print(f"problem {rec['id']}: {problem.get('problem', problem.get('prompt', '?'))}") |
| print(f"truth: {rec.get('answer', '-')} model said: {s['pred']} correct: {'yes' if s['correct'] else 'no'}") |
| print(f"answer {j}: {s['tokens']:,} tokens, {ending(s)}") |
|
|
| if s["verbatim_loop_end"] or a.tokens: |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(hf.model_dir("base")) |
| ids = tok(text, add_special_tokens=False).input_ids |
| found = loops.find_loop(ids) if s["verbatim_loop_end"] else None |
| if found: |
| onset, period = found |
| print(f"it loops: from token {onset:,} of {len(ids):,}, one block of {period} tokens repeats about " |
| f"{(len(ids) - 1 - onset) / period:.0f} times. One cycle:") |
| print(" " + repr(tok.decode(ids[onset:onset + period]))) |
| i = text.find("</think>") |
| if i >= 0: |
| print(f"\n--- where it closed its thinking (character {i:,} of {len(text):,}) ---") |
| print(text[max(0, i - 400):i + 200]) |
| print(f"\n--- the last {a.chars:,} characters ---") |
| print(text[-a.chars:]) |
| if a.tokens: |
| t0, t1 = max(0, a.tokens[0]), min(len(ids), a.tokens[1]) |
| print(f"\n--- answer tokens {t0:,} to {t1:,} (of {len(ids):,}) ---") |
| print(tok.decode(ids[t0:t1])) |
|
|
