"""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 [--batch-size 128] Writes .rows.jsonl and .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()