File size: 8,528 Bytes
40cfd23 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """OEIS-Eval-Neo: last-term exact-match accuracy on the decontaminated held-out OEIS set.
A batched PyTorch/transformers port of oeis_eval_mlx_neo.py from the NextTerm-440M repository
(the protocol behind that card's 34.43%), unchanged in substance:
* rows with fewer than 2 terms are skipped; prompt = every term but the last, comma-joined, with
a trailing comma; answer = the last term (compared as an exact integer)
* prompts of --max-context-tokens or more tokens (BOS included) are skipped
* greedy decoding, stopping at the comma token or EOS, up to --max-new-tokens (196)
* prediction = the generated text before the first comma, parsed as an integer
(unparseable, e.g. an immediate EOS, counts as wrong)
Works with any tokenizer that has a single comma token (the byte-level spark-gpt checkpoints, or
the NextTerm digit tokenizer).
Two modes, same accuracy: --mode teacher (default) checks greedy correctness in one forward pass
over prompt + answer + "," (correct iff the argmax reproduces every answer token and is then a
comma or EOS, which is exactly when greedy decoding would produce the answer); --mode generate
decodes greedily.
Usage: python eval_neo.py <hf_dir> <oeis_val_neo.jsonl> <out_prefix> [--batch-size 128]
Writes <out_prefix>.rows.jsonl and <out_prefix>.summary.json.
"""
import argparse
import json
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
def parse_generated(text):
text = text.split(",")[0]
try:
return int(text)
except ValueError:
return None
def main():
p = argparse.ArgumentParser()
p.add_argument("hf_dir")
p.add_argument("data")
p.add_argument("out_prefix")
p.add_argument("--batch-size", type=int, default=128)
p.add_argument("--max-new-tokens", type=int, default=196)
p.add_argument("--first-pass-tokens", type=int, default=24,
help="decode budget for the first pass; unfinished rows are redone at --max-new-tokens")
p.add_argument("--max-context-tokens", type=int, default=4096)
p.add_argument("--max-examples", type=int, default=0)
p.add_argument("--dtype", choices=["float32", "bfloat16"], default="bfloat16")
p.add_argument("--mode", choices=["teacher", "generate"], default="teacher",
help="teacher: one forward pass per row (exact for greedy); generate: decode")
args = p.parse_args()
started = time.perf_counter()
tok = AutoTokenizer.from_pretrained(args.hf_dir)
model = AutoModelForCausalLM.from_pretrained(
args.hf_dir, dtype=getattr(torch, args.dtype)).cuda().eval()
comma = tok.encode(",", add_special_tokens=False)
if not comma: # the NextTerm tokenizer strips a leading comma
comma = tok.encode("1,", add_special_tokens=False)[-1:]
assert len(comma) == 1, f"comma is not a single token: {comma}"
comma = comma[0]
stops = [comma] + ([tok.eos_token_id] if tok.eos_token_id is not None else [])
pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
prompts, answers, rows = [], [], []
for i, line in enumerate(open(args.data)):
seq = json.loads(line, parse_int=str).get("seq", [])
if len(seq) < 2:
continue
prompts.append(tok(",".join(str(x) for x in seq[:-1]) + ",").input_ids)
answers.append(int(seq[-1]))
rows.append(i)
if args.max_examples > 0:
prompts, answers, rows = prompts[:args.max_examples], answers[:args.max_examples], rows[:args.max_examples]
keep = [k for k, pr in enumerate(prompts)
if args.max_context_tokens <= 0 or len(pr) < args.max_context_tokens]
skipped_long = len(prompts) - len(keep)
keep.sort(key=lambda k: len(prompts[k]))
def greedy(idx, max_new):
"""Greedy-decode prompts idx (left-padded batches); return {k: tokens before the stop}
and the set of k that hit a stop token."""
gens, stopped = {}, set()
for b in range(0, len(idx), args.batch_size):
batch = idx[b:b + args.batch_size]
width = max(len(prompts[k]) for k in batch)
ids = torch.full((len(batch), width), pad, dtype=torch.long)
mask = torch.zeros((len(batch), width), dtype=torch.long)
for r, k in enumerate(batch):
ids[r, width - len(prompts[k]):] = torch.tensor(prompts[k])
mask[r, width - len(prompts[k]):] = 1
out = model.generate(ids.cuda(), attention_mask=mask.cuda(), do_sample=False,
max_new_tokens=max_new, eos_token_id=stops, pad_token_id=pad)
for r, k in enumerate(batch):
gen = out[r, width:].tolist()
cut = next((j for j, t in enumerate(gen) if t in stops), None)
if cut is not None:
stopped.add(k)
gens[k] = gen[:cut] if cut is not None else gen[:max_new]
return gens, stopped
results = []
if args.mode == "generate":
with torch.no_grad():
# Most answers are short. Decode everything briefly, then redo only rows that have
# not reached a stop token with the full budget (equivalent under greedy decoding).
gens, stopped = greedy(keep, min(args.first_pass_tokens, args.max_new_tokens))
long_rows = [k for k in keep if k not in stopped]
if long_rows and args.first_pass_tokens < args.max_new_tokens:
more, _ = greedy(long_rows, args.max_new_tokens)
gens.update(more)
for k in keep:
text = tok.decode(gens[k], skip_special_tokens=True)
pred = parse_generated(text)
results.append({"row_index": rows[k], "answer": str(answers[k]),
"prediction": None if pred is None else str(pred),
"correct": pred == answers[k], "parsed": pred is not None,
"generated_text": text})
else:
# Teacher-forced check, one forward pass per row: greedy decoding emits exactly the
# answer iff the argmax at every answer position is the answer's next token and the
# argmax after the last digit is a stop token. (Only differs from `generate` mode if a
# model emits leading zeros, which int() parsing would forgive.)
ans_ids = {k: tok.encode(str(answers[k]), add_special_tokens=False) for k in keep}
stop_t = torch.tensor(stops, device="cuda")
with torch.no_grad():
for b in range(0, len(keep), args.batch_size):
batch = keep[b:b + args.batch_size]
seqs = [prompts[k] + ans_ids[k] for k in batch]
width = max(map(len, seqs))
ids = torch.full((len(batch), width), pad, dtype=torch.long)
mask = torch.zeros((len(batch), width), dtype=torch.long)
for r, s in enumerate(seqs): # right padding is fine for teacher forcing
ids[r, :len(s)] = torch.tensor(s)
mask[r, :len(s)] = 1
top = model(ids.cuda(), attention_mask=mask.cuda()).logits.argmax(-1)
for r, k in enumerate(batch):
n_p, ans = len(prompts[k]), ans_ids[k]
pred = top[r, n_p - 1:n_p + len(ans)]
ok = (len(ans) < args.max_new_tokens
and bool((pred[:-1] == torch.tensor(ans, device="cuda")).all())
and bool(torch.isin(pred[-1], stop_t)))
results.append({"row_index": rows[k], "answer": str(answers[k]),
"correct": ok, "parsed": True})
correct = sum(r["correct"] for r in results)
summary = {
"checkpoint": args.hf_dir, "data": args.data, "documents": len(prompts),
"evaluated": len(results), "skipped_long": skipped_long,
"parsed": sum(r["parsed"] for r in results), "correct": correct,
"accuracy": correct / len(results) if results else 0.0,
"max_new_tokens": args.max_new_tokens, "dtype": args.dtype, "mode": args.mode,
"seconds": round(time.perf_counter() - started, 1),
}
with open(args.out_prefix + ".rows.jsonl", "w") as f:
for r in sorted(results, key=lambda r: r["row_index"]):
f.write(json.dumps(r) + "\n")
with open(args.out_prefix + ".summary.json", "w") as f:
json.dump(summary, f, indent=2)
print(json.dumps(summary))
if __name__ == "__main__":
main()
|