File size: 2,841 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 | """Write the NextTerm-440M card's polynomial-continuation eval sets as jsonl.
Reproduces arithmetic_eval.py's sampling exactly (same tasks, coefficient ranges, per-task
RNG random.Random(seed + degree), same draw order): for each prior-term count k from the
task's minimum up to --max-terms, --samples-per-k sequences f(0), ..., f(k) of a random
non-negative-integer polynomial; the prompt is the first k terms and the answer is f(k).
Each row is {"id", "seq", "k"} with seq = prompt + [answer], so scripts/eval_neo.py
(last-term exact match; teacher-forced greedy) scores it directly.
Degrees 5 and 6 (quintic, sextic; --extended) are an optional extension, not part of the
published polynomial evals: the x^5 and
x^6 coefficients are drawn from [0, 9] like x^4's, before the card's draws (highest degree
first, as the card orders them), so degrees 1-4 stay byte-identical.
Usage: python gen_poly_evals.py <out_dir> [--extended] [--seed 0 --max-terms 25 --samples-per-k 200]
"""
import argparse
import json
import random
from pathlib import Path
COEFF_MAX = {"g": 9, "f": 9, "a": 9, "b": 99, "c": 999, "d": 9999, "e": 99999} # g: x^6, f: x^5
CARD_TASKS = (("arithmetic", 1, 2), ("quadratic", 2, 3), ("cubic", 3, 4), ("quartic", 4, 5))
EXTENDED_TASKS = (("quintic", 5, 6), ("sextic", 6, 7))
p = argparse.ArgumentParser()
p.add_argument("out_dir")
p.add_argument("--seed", type=int, default=0)
p.add_argument("--max-terms", type=int, default=25)
p.add_argument("--samples-per-k", type=int, default=200)
p.add_argument("--extended", action="store_true", help="write quintic and sextic instead")
args = p.parse_args()
out = Path(args.out_dir)
out.mkdir(parents=True, exist_ok=True)
for name, degree, min_terms in (EXTENDED_TASKS if args.extended else CARD_TASKS):
rng = random.Random(args.seed + degree)
with open(out / f"poly_{name}.jsonl", "w") as f:
for k in range(min_terms, args.max_terms + 1):
for i in range(args.samples_per_k):
g = rng.randint(0, COEFF_MAX["g"]) if degree >= 6 else 0
f5 = rng.randint(0, COEFF_MAX["f"]) if degree >= 5 else 0
# identical draw order to arithmetic_eval.sample_coefficients
a = rng.randint(0, COEFF_MAX["a"]) if degree >= 4 else 0
b = rng.randint(0, COEFF_MAX["b"]) if degree >= 3 else 0
c = rng.randint(0, COEFF_MAX["c"]) if degree >= 2 else 0
d = rng.randint(0, COEFF_MAX["d"]) if degree >= 1 else 0
e = rng.randint(0, COEFF_MAX["e"])
seq = [g * x**6 + f5 * x**5 + a * x**4 + b * x**3 + c * x**2 + d * x + e
for x in range(k + 1)]
f.write(json.dumps({"id": f"{name}_k{k}_{i}", "seq": seq, "k": k}) + "\n")
print(name, (args.max_terms - min_terms + 1) * args.samples_per_k, "rows")
|