File size: 4,901 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 | #!/usr/bin/env python
"""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 # noqa: E402
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]))
|