| """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: |
| 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(): |
| |
| |
| 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: |
| |
| |
| |
| |
| 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): |
| 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() |
|
|
