Add model card and evaluation utilities

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+radar_chart.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,179 @@
 ---
 license: cc-by-sa-4.0
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- oeis
+- integer-sequences
+- qwen3
+- causal-lm
+- mlx
+datasets:
+- N8Programs/oeis-massive
 ---
+
+# NextTerm-440M
+
+[![Benchmark radar](radar_chart.png)](radar_chart.png)
+
+## Model Summary
+
+NextTerm-440M is a 440M parameter causal transformer trained to continue integer
+sequences. It uses a Qwen3 architecture with a compact 16-token digit
+vocabulary: decimal digits, negative sign, comma separator, BOS, EOS, PAD, and one unused token.
+
+The model was trained on an extended OEIS corpus that enhanced many OEIS sequences with additional terms from b-files (supplemental appendices provided with OEIS) and then further augmented the data w/ a variety of prefix-preserving transforms empirically selected via small pilot experiments. The model was trained for 14B tokens w/ preserved sequence prefixes rather than concatenating distinct documents (as this was found to improve performance in pilot experiments).
+
+
+NextTerm-440M improves dramatically over NextTerm-47M on long-context sequence continuation (as it was trained w/ a context length of 4096), innate OEIS knowledge, and long-range in context learning. The 47M model, however, remains ahead on very short prefixes that require simple rule induction without much context, which may be due to the 440M model's training on longer contexts and more complex sequences.
+
+The tokenizer accepts integer sequences formatted as comma-separated values, for
+example:
+
+```text
+1,-2,3,-4,
+```
+
+The tokenizer ignores characters other than digits, commas, and `-`. Digits are
+tokenized individually, so there is no fixed integer-magnitude limit, but large
+integers consume more context. The model was not trained on numbers with leading
+zeros, so strings like `01,02,03,` should be treated as out of distribution.
+
+## Training Details
+
+| Field | Value |
+| --- | --- |
+| Parameters | 440,500,224 |
+| Architecture | Qwen3-style causal LM |
+| Layers | 28 |
+| Hidden size | 1024 |
+| FFN size | 3072 |
+| Attention heads | 16 |
+| KV heads | 8 |
+| Vocabulary size | 16 |
+| Training tokens | 13,999,999,995 |
+| Sequence length cap | 4096 training tokens per sequence |
+| Batch mode | Length-bucketed sequence batches |
+| Optimizer | Muon/AdamW hybrid |
+| LR schedule | Linear warmup to `1e-2` for Muon, `1e-4` for AdamW, cosine decay to 0.1x, final cooldown to 0 |
+| Training hardware | Single H100 |
+| Export dtype | bfloat16 |
+
+A classic Muon/AdamW hybrid was used: Muon for 2D weight matrices and AdamW for 1D parameters and embedding matrices.
+
+The model was trained on the following files in the `N8Programs/oeis-massive` dataset, randomly mixed:
+
+- `oeis_train_bfile_prefix4096.packed`
+- `oeis_synth_aug0_inv_len_13245370099_seed0.packed`
+
+## Evaluation Results
+
+### Main Benchmarks
+
+| Model | OEIS-Eval-Neo | Ryskina & Knight | M1 Competition 111 MAPE |
+| --- | ---: | ---: | ---: |
+| **NextTerm-440M** | **34.43%** | 52.63% | **17.6239** |
+| NextTerm-47M | 29.49% | **70.18%** | 18.7621 |
+| Qwen3-0.6B | 18.44% | 33.33% | 22.7984 |
+| Qwen3-1.7B | 20.77% | 49.12% | 22.2411 |
+| Qwen3-4B | 23.74% | 63.16% | 19.1731 |
+| Qwen3-8B | 24.62% | 57.89% | 18.4027 |
+| Qwen3-14B | 26.00% | 59.65% | 17.9837 |
+
+OEIS-Eval-Neo is a decontaminated held-out OEIS next-term evaluation. M1
+Competition 111 reports macro MAPE, where lower is better. Ryskina & Knight
+(2021) is a 57-sequence next-term benchmark based on psychometrics and puzzles. Note that the 47M model's strong performance on Ryskina & Knight is indicative of its strength on short-prefix sequences and rule induction. 
+
+### Polynomial Continuation
+
+The polynomial continuation evaluation samples integer sequences from
+polynomials of degree 1 through 4 and asks for the next term. Accuracy is exact
+match across 200 samples for each prompt length `k`.
+
+| Model | Arithmetic | Quadratic | Cubic | Quartic |
+| --- | ---: | ---: | ---: | ---: |
+| **NextTerm-440M** | 94.38% | **86.39%** | **75.20%** | **67.83%** |
+| NextTerm-47M | 94.15% | 81.07% | 37.43% | 15.17% |
+| Qwen3-0.6B | 90.31% | 8.72% | 0.30% | 0.02% |
+| Qwen3-1.7B | 93.10% | 41.57% | 5.36% | 0.71% |
+| Qwen3-4B | 93.90% | 77.26% | 28.18% | 5.98% |
+| Qwen3-8B | **96.10%** | 80.59% | 32.93% | 7.95% |
+| Qwen3-14B | 95.60% | 84.61% | 49.16% | 14.98% |
+
+## Usage
+
+### MLX
+
+```bash
+mlx_lm.generate --model N8Programs/NextTerm-440M --prompt "1,2,3,"
+```
+
+### Hugging Face Transformers
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "N8Programs/NextTerm-440M"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+prompt = "1,2,3,4,5,"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=64,
+    do_sample=False,
+    eos_token_id=[tokenizer.convert_tokens_to_ids(","), tokenizer.eos_token_id],
+    pad_token_id=tokenizer.pad_token_id,
+)
+
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+For strict next-term evaluation, stop generation on *comma or EOS* and parse the
+text before the first comma as the predicted integer.
+
+## Reproducibility
+
+This repository contains the local evaluation scripts and artifacts used for the
+results above, including:
+
+- `oeis_eval_mlx_neo.py` for OEIS-Eval-Neo with MLX batch generation.
+- `arithmetic_eval.py` for arithmetic/quadratic/cubic/quartic continuation.
+- `eval_m1_competition_mape_mlx.py` for M1 Competition 111 MAPE.
+- `eval_results.txt` for the compact result table.
+
+
+The last three training checkpoints are available separately at
+[N8Programs/NextTerm-440M-Checkpoints](https://huggingface.co/N8Programs/NextTerm-440M-Checkpoints).
+The released `final_latest` checkpoint was trained for 14B tokens. Additionally, the checkpoint corresponding to the best val loss is available as well (although it is not included in the main results table as it was inferior on downstream eval performance).
+
+The `.packed` files used for training are binary files containing the tokenized and augmented OEIS data - w/ tokens encoded as nibbles. A dedicate decoder is provided in this repo as `decode_packed_oeis.py`.
+
+## Citation
+
+```bibtex
+@misc{nextterm440m2026,
+  author       = {Nathan Breslow},
+  title        = {NextTerm-440M: A Pretrained Transformer for Integer Sequence Prediction},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/N8Programs/NextTerm-440M}},
+  note         = {440.5M parameter model trained on augmented OEIS data}
+}
+```
+
+## Attribution
+
+This model and dataset were trained and created using data from the
+**On-Line Encyclopedia of Integer Sequences (OEIS)**.
+
+- Source: https://oeis.org/
+- License: Creative Commons Attribution-ShareAlike 4.0 (CC BY-SA 4.0)
+- OEIS End-User License Agreement: https://oeis.org/wiki/The_OEIS_End-User_License_Agreement

arithmetic_eval.py

@@ -0,0 +1,263 @@
+"""
+Evaluate the NextTerm model on synthetic polynomial sequences (arithmetic through
+quartic) with varying amounts of prior context.
+
+For each prior term count (default: 2-25 for linear, 3-25 for quadratic,
+4-25 for cubic, 5-25 for quartic), we sample sequences from
+ax^4 + bx^3 + cx^2 + dx + e with coefficient ranges:
+    e in [0, 99999]
+    d in [0, 9999]
+    c in [0, 999]
+    b in [0, 99]
+    a in [0, 9]
+
+Higher-order coefficients are set to 0 for lower-degree tests (e.g., arithmetic
+sequences only use d and e).
+"""
+
+import argparse
+import csv
+import random
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+from mlx_lm import load
+from mlx_lm.generate import BatchGenerator
+from mlx_lm.tuner.utils import print_trainable_parameters
+from tqdm import tqdm
+
+# Step 335000
+MODEL_NAME = "NextTerm-47M"
+MAX_NEW_TOKENS = 20
+DEFAULT_MAX_TERMS = 25
+SAMPLES_PER_K = 200
+RANDOM_SEED = 0
+
+COEFF_MAX = {"a": 9, "b": 99, "c": 999, "d": 9999, "e": 99999}
+
+
+@dataclass(frozen=True)
+class PolynomialTask:
+    name: str
+    degree: int
+    min_terms: int
+
+
+TASKS = (
+    PolynomialTask("arithmetic", 1, 2),
+    PolynomialTask("quadratic", 2, 3),
+    PolynomialTask("cubic", 3, 4),
+    PolynomialTask("quartic", 4, 5),
+)
+
+
+def parse_generated(text: str):
+    if "," in text:
+        text = text.split(",")[0]
+    try:
+        return int(text)
+    except ValueError:
+        print(f"Could not parse generated text: {text!r}")
+        return None
+
+
+def polynomial_value(coeffs: Tuple[int, int, int, int, int], x: int) -> int:
+    a, b, c, d, e = coeffs
+    return (
+        a * (x**4)
+        + b * (x**3)
+        + c * (x**2)
+        + d * x
+        + e
+    )
+
+
+def sample_coefficients(
+    degree: int, rng: random.Random
+) -> Tuple[int, int, int, int, int]:
+    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"])
+    return a, b, c, d, e
+
+
+def generate_polynomial_sequences(
+    task: PolynomialTask, max_terms: int, samples_per_k: int, rng: random.Random
+) -> List[Dict[str, int | List[int] | str]]:
+    sequences = []
+    for k in range(task.min_terms, max_terms + 1):
+        for _ in range(samples_per_k):
+            coeffs = sample_coefficients(task.degree, rng)
+            seq = [polynomial_value(coeffs, i) for i in range(k + 1)]
+            sequences.append(
+                {"k": k, "prompt": seq[:-1], "answer": seq[-1], "task": task.name}
+            )
+    return sequences
+
+
+def build_prompts(
+    sequences: List[Dict[str, int | List[int] | str]], tokenizer
+) -> Tuple[List[List[int]], List[Dict[str, int | List[int] | str]]]:
+    prompts = []
+    metadata = []
+    for record in sequences:
+        prompt_text = ",".join(str(x) for x in record["prompt"]) + ","
+        prompts.append(tokenizer.encode(prompt_text))
+        metadata.append(record)
+    return prompts, metadata
+
+
+def evaluate_task(
+    task: PolynomialTask,
+    args,
+    model,
+    tokenizer,
+    sep_token: int,
+) -> Tuple[Dict[int, Dict[str, int]], List[Dict[str, object]]]:
+    rng = random.Random(args.seed + task.degree)
+    sequences = generate_polynomial_sequences(
+        task, args.max_terms, args.samples_per_k, rng
+    )
+    if not sequences:
+        print(
+            f"Skipping {task.name}: min_terms {task.min_terms} exceeds max_terms "
+            f"{args.max_terms}"
+        )
+        return {}, []
+
+    print(
+        f"\nEvaluating {task.name} sequences (degree {task.degree}) with "
+        f"{len(sequences)} samples ({args.samples_per_k} per prior-term count)"
+    )
+
+    prompts, metadata = build_prompts(sequences, tokenizer)
+
+    gen = BatchGenerator(model, stop_tokens=[[sep_token], [tokenizer.eos_token_id]])
+    uids = gen.insert(prompts, [args.max_new_tokens] * len(prompts))
+
+    meta_by_uid = dict(zip(uids, metadata))
+    results = {uid: [] for uid in uids}
+    finished = set()
+    with tqdm(total=len(uids), desc=f"{task.name.title()}") as pbar:
+        while True:
+            responses = gen.next()
+            if not isinstance(responses, tuple) or len(responses) != 2:
+                raise RuntimeError(
+                    "Unexpected mlx_lm BatchGenerator.next() API. "
+                    "Update your mlx_lm version"
+                )
+            prompt_responses, generation_responses = responses
+            if not prompt_responses and not generation_responses:
+                break
+            if not generation_responses:
+                continue
+            for r in generation_responses:
+                results[r.uid].append(r.token)
+                if r.finish_reason == "stop" and r.uid not in finished:
+                    finished.add(r.uid)
+                    pbar.update(1)
+    gen.close()
+
+    stats = {
+        k: {"correct": 0, "parsed": 0, "total": 0}
+        for k in range(task.min_terms, args.max_terms + 1)
+    }
+
+    for uid in uids:
+        tokens = results[uid]
+        text = tokenizer.decode(tokens[:-1]) if tokens else ""
+        prediction = parse_generated(text)
+
+        meta = meta_by_uid[uid]
+        k = meta["k"]
+        answer = meta["answer"]
+
+        stats[k]["total"] += 1
+        if prediction is not None:
+            stats[k]["parsed"] += 1
+            if prediction == answer:
+                stats[k]["correct"] += 1
+
+    print("Accuracy by prior term count:")
+    csv_rows: List[Dict[str, object]] = []
+    for k in range(task.min_terms, args.max_terms + 1):
+        s = stats[k]
+        accuracy = s["correct"] / s["total"] if s["total"] else 0.0
+        print(
+            f"{k} terms: parsed {s['parsed']}/{s['total']} "
+            f"accuracy {s['correct']}/{s['total']} = {accuracy:.4f}"
+        )
+        csv_rows.append(
+            {
+                "task": task.name,
+                "degree": task.degree,
+                "k": k,
+                "parsed": s["parsed"],
+                "total": s["total"],
+                "correct": s["correct"],
+                "accuracy": accuracy,
+            }
+        )
+
+    overall_total = sum(s["total"] for s in stats.values())
+    overall_parsed = sum(s["parsed"] for s in stats.values())
+    overall_correct = sum(s["correct"] for s in stats.values())
+    overall_accuracy = overall_correct / overall_total if overall_total else 0.0
+    print(
+        f"Overall {task.name}: parsed {overall_parsed}/{overall_total} "
+        f"accuracy {overall_correct}/{overall_total} = {overall_accuracy:.4f}"
+    )
+
+    return stats, csv_rows
+
+
+def write_csv(rows: List[Dict[str, object]], path: str) -> None:
+    if not rows:
+        print("No results to write to CSV.")
+        return
+
+    fieldnames = ["task", "degree", "k", "parsed", "total", "correct", "accuracy"]
+    with open(path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+    print(f"Wrote results to {path}")
+
+
+def evaluate(args):
+    random.seed(args.seed)
+
+    model, tokenizer = load(args.model_name)
+
+    print_trainable_parameters(model)
+
+    sep_token = tokenizer.encode("1,")[-1]
+    all_rows: List[Dict[str, object]] = []
+    for task in TASKS:
+        _, rows = evaluate_task(task, args, model, tokenizer, sep_token)
+        all_rows.extend(rows)
+
+    if args.csv_path:
+        write_csv(all_rows, args.csv_path)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", default=MODEL_NAME)
+    parser.add_argument("--max_new_tokens", type=int, default=MAX_NEW_TOKENS)
+    parser.add_argument("--max_terms", type=int, default=DEFAULT_MAX_TERMS)
+    parser.add_argument("--samples_per_k", type=int, default=SAMPLES_PER_K)
+    parser.add_argument("--seed", type=int, default=RANDOM_SEED)
+    parser.add_argument("--csv_path", type=str, default=None)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    evaluate(args)
+
+
+if __name__ == "__main__":
+    main()

decode_packed_oeis.py

@@ -0,0 +1,185 @@
+"""Small reference decoder for bigOEIS `.packed` files.
+
+On disk, each token is stored in a 4-bit nibble:
+
+    0..9  decimal digits
+    10    term separator, i.e. comma
+    11    negative sign
+    14    final padding nibble, if the file has an odd nibble count
+    15    sequence delimiter / EOS
+
+Note that the packed disk codes are intentionally compact and are not exactly
+the model vocabulary ids: the model uses NEG=10 and SEP=11. Use
+`iter_model_token_rows()` when you want rows in model-token-id space.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Iterator
+
+
+PACKED_SEP = 10
+PACKED_NEG = 11
+PACKED_PAD = 14
+PACKED_DELIM = 15
+
+MODEL_NEG = 10
+MODEL_SEP = 11
+MODEL_BOS = 12
+MODEL_EOS = 13
+
+
+def iter_packed_nibbles(path: str | Path, chunk_size: int = 8 * 1024 * 1024) -> Iterator[int]:
+    """Yield high nibble then low nibble for every byte in `path`."""
+    with Path(path).open("rb") as f:
+        while True:
+            chunk = f.read(chunk_size)
+            if not chunk:
+                return
+            for byte in chunk:
+                yield byte >> 4
+                yield byte & 0x0F
+
+
+def iter_model_token_rows(
+    path: str | Path,
+    *,
+    include_bos_eos: bool = True,
+    max_rows: int | None = None,
+    strict: bool = True,
+) -> Iterator[list[int]]:
+    """Yield each packed sequence as model-token ids.
+
+    By default rows include BOS/EOS, matching `tokenize_utils.tokenize_sequence`.
+    Set `include_bos_eos=False` to get only the content tokens.
+    """
+    row: list[int] = [MODEL_BOS] if include_bos_eos else []
+    yielded = 0
+    seen_pad = False
+
+    for nib in iter_packed_nibbles(path):
+        if nib == PACKED_PAD:
+            seen_pad = True
+            continue
+        if seen_pad:
+            if strict:
+                raise ValueError("Found non-pad nibble after final packed padding.")
+            continue
+
+        if 0 <= nib <= 9:
+            row.append(nib)
+        elif nib == PACKED_SEP:
+            row.append(MODEL_SEP)
+        elif nib == PACKED_NEG:
+            row.append(MODEL_NEG)
+        elif nib == PACKED_DELIM:
+            if include_bos_eos:
+                row.append(MODEL_EOS)
+            yield row
+            yielded += 1
+            if max_rows is not None and yielded >= max_rows:
+                return
+            row = [MODEL_BOS] if include_bos_eos else []
+        else:
+            if strict:
+                raise ValueError(f"Invalid packed nibble: {nib}")
+
+    empty = [MODEL_BOS] if include_bos_eos else []
+    if strict and row != empty:
+        raise ValueError("Packed file ended with an unterminated sequence.")
+
+
+def iter_integer_sequences(
+    path: str | Path,
+    *,
+    as_ints: bool = False,
+    max_rows: int | None = None,
+    strict: bool = True,
+) -> Iterator[list[int] | list[str]]:
+    """Yield decoded OEIS rows.
+
+    Values are strings by default so enormous integers round-trip exactly
+    through JSON. Pass `as_ints=True` if Python integers are more convenient.
+    """
+    terms: list[str] = []
+    chars: list[str] = []
+    yielded = 0
+    seen_pad = False
+
+    def finish_term() -> None:
+        if chars:
+            terms.append("".join(chars))
+            chars.clear()
+        elif strict:
+            raise ValueError("Encountered an empty term in packed sequence.")
+
+    for nib in iter_packed_nibbles(path):
+        if nib == PACKED_PAD:
+            seen_pad = True
+            continue
+        if seen_pad:
+            if strict:
+                raise ValueError("Found non-pad nibble after final packed padding.")
+            continue
+
+        if 0 <= nib <= 9:
+            chars.append(str(nib))
+        elif nib == PACKED_NEG:
+            if strict and chars:
+                raise ValueError("Found a negative sign after term digits had started.")
+            chars.append("-")
+        elif nib == PACKED_SEP:
+            finish_term()
+        elif nib == PACKED_DELIM:
+            if chars:
+                finish_term()
+            row = [int(term) for term in terms] if as_ints else list(terms)
+            yield row
+            yielded += 1
+            if max_rows is not None and yielded >= max_rows:
+                return
+            terms.clear()
+        else:
+            if strict:
+                raise ValueError(f"Invalid packed nibble: {nib}")
+
+    if strict and (terms or chars):
+        raise ValueError("Packed file ended with an unterminated sequence.")
+
+
+def _main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("packed_file", help="Path to a .packed file")
+    parser.add_argument("-n", "--max-rows", type=int, default=5)
+    parser.add_argument("--tokens", action="store_true", help="Print model-token rows instead of decoded terms")
+    parser.add_argument("--content-only", action="store_true", help="Omit BOS/EOS when printing token rows")
+    parser.add_argument("--ints", action="store_true", help="Emit decoded terms as JSON numbers instead of strings")
+    parser.add_argument("--no-strict", action="store_true", help="Ignore invalid/trailing data instead of raising")
+    args = parser.parse_args()
+
+    strict = not args.no_strict
+    if args.tokens:
+        rows = iter_model_token_rows(
+            args.packed_file,
+            include_bos_eos=not args.content_only,
+            max_rows=args.max_rows,
+            strict=strict,
+        )
+        for row in rows:
+            print(json.dumps({"tokens": row}, separators=(",", ":")))
+    else:
+        rows = iter_integer_sequences(
+            args.packed_file,
+            as_ints=args.ints,
+            max_rows=args.max_rows,
+            strict=strict,
+        )
+        for row in rows:
+            print(json.dumps({"seq": row}, separators=(",", ":")))
+
+
+if __name__ == "__main__":
+    _main()

eval_m1_competition_mape_mlx.py

@@ -0,0 +1,425 @@
+#!/usr/bin/env python3
+"""Evaluate NextTerm-style MLX models on the canonical 111-series M1 set."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import math
+import time
+from dataclasses import dataclass
+from decimal import Decimal
+from pathlib import Path
+from statistics import mean, median
+
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.generate import BatchGenerator
+from tqdm import tqdm
+
+from eval_m1_monthly_mape_mlx import (
+    SuppressTokenLogits,
+    context_scale,
+    mape_for_predictions,
+    parse_decimal,
+    parse_generated_terms,
+    scale_decimal_to_int,
+)
+
+
+DATA_PATH = Path("m1_competition_111.jsonl")
+MODEL_PATH = Path("NextTerm-440M")
+OUTPUT_PATH = Path("m1_eval_results/m1_competition111_nextterm440m_greedy_per_series.jsonl")
+SUMMARY_PATH = Path("m1_eval_results/m1_competition111_nextterm440m_greedy_summary.json")
+
+
+@dataclass
+class SeriesRecord:
+    row_index: int
+    series_name: str
+    original_id: str
+    period: str
+    frequency: int
+    type: str
+    description: str
+    horizon: int
+    raw_context: list[str]
+    raw_target: list[str]
+    context_values: list[Decimal]
+    target_values: list[Decimal]
+    naive2_forecast: list[Decimal]
+    naive2_mape: float
+    scale: int
+    scaled_context: list[int]
+
+
+def load_m1_competition(path: Path) -> list[SeriesRecord]:
+    records: list[SeriesRecord] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            row = json.loads(line)
+            raw_context = [str(x) for x in row["context"]]
+            raw_target = [str(x) for x in row["target"]]
+            context_values = [parse_decimal(x) for x in raw_context]
+            target_values = [parse_decimal(x) for x in raw_target]
+            scale = context_scale(raw_context)
+            scaled_context = [scale_decimal_to_int(v, scale) for v in context_values]
+            records.append(
+                SeriesRecord(
+                    row_index=int(row["row_index"]),
+                    series_name=str(row["series_name"]),
+                    original_id=str(row["original_id"]),
+                    period=str(row["period"]),
+                    frequency=int(row["frequency"]),
+                    type=str(row["type"]),
+                    description=str(row["description"]),
+                    horizon=int(row["horizon"]),
+                    raw_context=raw_context,
+                    raw_target=raw_target,
+                    context_values=context_values,
+                    target_values=target_values,
+                    naive2_forecast=[parse_decimal(str(x)) for x in row["naive2_forecast"]],
+                    naive2_mape=float(row["naive2_mape"]),
+                    scale=scale,
+                    scaled_context=scaled_context,
+                )
+            )
+    return records
+
+
+def decimal_or_none(value: Decimal | None) -> str | None:
+    return str(value) if value is not None else None
+
+
+def load_completed(path: Path) -> dict[int, dict]:
+    completed: dict[int, dict] = {}
+    if not path.exists():
+        return completed
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            record = json.loads(line)
+            completed[int(record["row_index"])] = record
+    return completed
+
+
+def aggregate_summary(records: list[dict]) -> dict:
+    def collect(key: str, rows: list[dict] = records) -> list[float]:
+        vals: list[float] = []
+        for r in rows:
+            value = r.get(key)
+            if value is not None and math.isfinite(float(value)):
+                vals.append(float(value))
+        return vals
+
+    model_series = collect("mape")
+    naive2_series = collect("naive2_mape")
+    all_model_apes = [
+        float(x)
+        for r in records
+        for x in (r.get("apes") or [])
+        if x is not None and math.isfinite(float(x))
+    ]
+    all_naive2_apes = [
+        float(x)
+        for r in records
+        for x in (r.get("naive2_apes") or [])
+        if x is not None and math.isfinite(float(x))
+    ]
+
+    periods = sorted({str(r["period"]) for r in records})
+    by_period = {}
+    for period in periods:
+        rows = [r for r in records if r["period"] == period]
+        period_model = collect("mape", rows)
+        period_naive2 = collect("naive2_mape", rows)
+        period_model_apes = [
+            float(x)
+            for r in rows
+            for x in (r.get("apes") or [])
+            if x is not None and math.isfinite(float(x))
+        ]
+        period_naive2_apes = [
+            float(x)
+            for r in rows
+            for x in (r.get("naive2_apes") or [])
+            if x is not None and math.isfinite(float(x))
+        ]
+        by_period[period] = {
+            "series_count": len(rows),
+            "model_macro_mape": mean(period_model) if period_model else None,
+            "model_point_mape": mean(period_model_apes) if period_model_apes else None,
+            "naive2_macro_mape": mean(period_naive2) if period_naive2 else None,
+            "naive2_point_mape": mean(period_naive2_apes) if period_naive2_apes else None,
+            "parsed_full_series": sum(
+                1 for r in rows if int(r.get("parsed_terms", 0)) >= int(r.get("horizon", 0))
+            ),
+        }
+
+    horizon_counts: dict[int, int] = {}
+    for r in records:
+        h = int(r["horizon"])
+        horizon_counts[h] = horizon_counts.get(h, 0) + 1
+
+    return {
+        "series_count": len(records),
+        "horizon_counts": dict(sorted(horizon_counts.items())),
+        "model_macro_mape": mean(model_series) if model_series else None,
+        "model_median_series_mape": median(model_series) if model_series else None,
+        "model_point_mape": mean(all_model_apes) if all_model_apes else None,
+        "naive2_macro_mape": mean(naive2_series) if naive2_series else None,
+        "naive2_point_mape": mean(all_naive2_apes) if all_naive2_apes else None,
+        "parsed_full_series": sum(
+            1 for r in records if int(r.get("parsed_terms", 0)) >= int(r.get("horizon", 0))
+        ),
+        "total_missing_terms": sum(int(r.get("missing_terms", 0)) for r in records),
+        "malformed_series": sum(1 for r in records if r.get("malformed", False)),
+        "by_period": by_period,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=Path, default=DATA_PATH)
+    parser.add_argument("--model", type=Path, default=MODEL_PATH)
+    parser.add_argument("--output", type=Path, default=OUTPUT_PATH)
+    parser.add_argument("--summary-output", type=Path, default=SUMMARY_PATH)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--max-new-tokens", type=int, default=1024)
+    parser.add_argument("--max-context-tokens", type=int, default=2048)
+    parser.add_argument("--max-series", type=int, default=0)
+    parser.add_argument(
+        "--allow-eos-before-horizon",
+        action="store_true",
+        help="Do not suppress EOS/pad while waiting for all forecast terms.",
+    )
+    parser.add_argument(
+        "--missing-policy",
+        choices=["zero", "last_context", "skip"],
+        default="zero",
+    )
+    parser.add_argument("--overwrite", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    started = time.perf_counter()
+    series = load_m1_competition(args.data_path)
+    if args.max_series > 0:
+        series = series[: args.max_series]
+    print(f"Loaded {len(series)} canonical M1 competition series from {args.data_path}")
+    print("Horizons:", {h: sum(1 for s in series if s.horizon == h) for h in sorted({s.horizon for s in series})})
+
+    model, tokenizer = load(str(args.model))
+    print(f"Loaded model: {args.model}")
+
+    prompts_text = [",".join(str(x) for x in s.scaled_context) + "," for s in series]
+    prompts = [tokenizer.encode(text) for text in prompts_text]
+    sep_token = tokenizer.encode("1,")[-1]
+    eval_indices = [
+        i
+        for i, prompt in enumerate(prompts)
+        if args.max_context_tokens <= 0 or len(prompt) < args.max_context_tokens
+    ]
+    skipped_long = len(prompts) - len(eval_indices)
+    eval_indices = sorted(eval_indices, key=lambda i: len(prompts[i]))
+    print(
+        f"Evaluating {len(eval_indices)} series; skipped_long={skipped_long}; "
+        f"batch_size={args.batch_size}; max_new_tokens={args.max_new_tokens}; sep_token={sep_token}"
+    )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.summary_output.parent.mkdir(parents=True, exist_ok=True)
+    if args.overwrite and args.output.exists():
+        args.output.unlink()
+    completed = load_completed(args.output)
+    if completed:
+        print(f"Resuming from {args.output}: {len(completed)} series already done")
+    todo_indices = [i for i in eval_indices if i not in completed]
+
+    stop_tokens = []
+    suppress_token_ids: list[int] = []
+    if getattr(tokenizer, "eos_token_id", None) is not None:
+        suppress_token_ids.append(tokenizer.eos_token_id)
+        if args.allow_eos_before_horizon:
+            stop_tokens.append([tokenizer.eos_token_id])
+    if getattr(tokenizer, "pad_token_id", None) is not None:
+        suppress_token_ids.append(tokenizer.pad_token_id)
+        if args.allow_eos_before_horizon:
+            stop_tokens.append([tokenizer.pad_token_id])
+    suppress_processor = None if args.allow_eos_before_horizon else SuppressTokenLogits(suppress_token_ids)
+    gen = BatchGenerator(
+        model,
+        stop_tokens=stop_tokens or None,
+        completion_batch_size=args.batch_size,
+        prefill_batch_size=args.batch_size,
+    )
+    uid_to_idx: dict[int, int] = {}
+    generated_tokens: dict[int, list[int]] = {}
+    separator_counts: dict[int, int] = {}
+    if todo_indices:
+        logits_processors = (
+            [[suppress_processor] for _ in todo_indices]
+            if suppress_processor is not None
+            else None
+        )
+        uids = gen.insert(
+            [prompts[i] for i in todo_indices],
+            [args.max_new_tokens] * len(todo_indices),
+            logits_processors=logits_processors,
+        )
+        uid_to_idx = {uid: idx for uid, idx in zip(uids, todo_indices)}
+        generated_tokens = {uid: [] for uid in uids}
+        separator_counts = {uid: 0 for uid in uids}
+    finished: set[int] = set()
+
+    try:
+        with args.output.open("a", encoding="utf-8") as out:
+            with tqdm(total=len(todo_indices), desc="Generating") as progress:
+
+                def finalize_uid(uid: int, finish_reason: str) -> None:
+                    if uid in finished:
+                        return
+                    finished.add(uid)
+                    idx = uid_to_idx[uid]
+                    s = series[idx]
+                    generated_text = tokenizer.decode(generated_tokens[uid])
+                    scaled_terms, malformed = parse_generated_terms(generated_text, s.horizon)
+                    predictions: list[Decimal | None] = [
+                        Decimal(term) / Decimal(s.scale) for term in scaled_terms
+                    ]
+                    predictions.extend([None] * (s.horizon - len(predictions)))
+                    fallback = s.context_values[-1] if s.context_values else Decimal(0)
+                    mape, apes, missing = mape_for_predictions(
+                        s.target_values,
+                        predictions,
+                        missing_policy=args.missing_policy,
+                        fallback=fallback,
+                    )
+                    naive2_mape, naive2_apes, _ = mape_for_predictions(
+                        s.target_values,
+                        s.naive2_forecast,
+                        missing_policy="skip",
+                        fallback=fallback,
+                    )
+                    record = {
+                        "row_index": s.row_index,
+                        "series_name": s.series_name,
+                        "original_id": s.original_id,
+                        "period": s.period,
+                        "frequency": s.frequency,
+                        "type": s.type,
+                        "description": s.description,
+                        "horizon": s.horizon,
+                        "scale": s.scale,
+                        "context_length": len(s.context_values),
+                        "prompt_tokens": len(prompts[idx]),
+                        "target": [str(x) for x in s.target_values],
+                        "scaled_prediction_terms": scaled_terms,
+                        "prediction": [decimal_or_none(x) for x in predictions],
+                        "parsed_terms": len(scaled_terms),
+                        "missing_terms": missing,
+                        "malformed": malformed,
+                        "mape": mape,
+                        "apes": apes,
+                        "naive2_forecast": [str(x) for x in s.naive2_forecast],
+                        "naive2_mape": naive2_mape,
+                        "naive2_apes": naive2_apes,
+                        "generated_text": generated_text,
+                        "finish_reason": finish_reason,
+                    }
+                    out.write(json.dumps(record) + "\n")
+                    out.flush()
+                    progress.update(1)
+
+                while todo_indices and len(finished) < len(todo_indices):
+                    responses = gen.next()
+                    if isinstance(responses, tuple) and len(responses) == 2:
+                        prompt_responses, generation_responses = responses
+                    elif isinstance(responses, list):
+                        prompt_responses, generation_responses = [], responses
+                    else:
+                        raise RuntimeError(
+                            "Unexpected mlx_lm BatchGenerator.next() API. Update mlx_lm version."
+                        )
+                    if not prompt_responses and not generation_responses:
+                        break
+                    remove_uids: list[int] = []
+                    for response in generation_responses:
+                        uid = response.uid
+                        if uid in finished:
+                            continue
+                        idx = uid_to_idx[uid]
+                        horizon = series[idx].horizon
+                        if response.finish_reason != "stop":
+                            token = int(response.token)
+                            generated_tokens[uid].append(token)
+                            if token == sep_token:
+                                separator_counts[uid] += 1
+                                if separator_counts[uid] >= horizon:
+                                    finalize_uid(uid, "separator_count")
+                                    remove_uids.append(uid)
+                                    continue
+                        if response.finish_reason is not None:
+                            finalize_uid(uid, str(response.finish_reason))
+                    if remove_uids:
+                        gen.remove(remove_uids)
+                if len(finished) != len(todo_indices):
+                    raise RuntimeError(f"Finished {len(finished)}/{len(todo_indices)} series")
+    finally:
+        gen.close()
+        mx.clear_cache()
+        gc.collect()
+
+    eval_index_set = set(eval_indices)
+    all_records = [
+        r for idx, r in sorted(load_completed(args.output).items()) if idx in eval_index_set
+    ]
+    aggregate = aggregate_summary(all_records)
+    elapsed = time.perf_counter() - started
+    summary = {
+        "data_path": str(args.data_path),
+        "model": str(args.model),
+        "output": str(args.output),
+        "series_loaded": len(series),
+        "series_evaluated": len(all_records),
+        "skipped_long": skipped_long,
+        "max_context_tokens": args.max_context_tokens,
+        "max_new_tokens": args.max_new_tokens,
+        "batch_size": args.batch_size,
+        "missing_policy": args.missing_policy,
+        "suppress_eos_until_horizon": not args.allow_eos_before_horizon,
+        "seconds": elapsed,
+        **aggregate,
+    }
+    args.summary_output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
+    print(
+        json.dumps(
+            {
+                k: summary[k]
+                for k in [
+                    "series_evaluated",
+                    "model_macro_mape",
+                    "model_point_mape",
+                    "naive2_macro_mape",
+                    "naive2_point_mape",
+                    "parsed_full_series",
+                    "total_missing_terms",
+                    "malformed_series",
+                    "seconds",
+                ]
+            },
+            indent=2,
+        )
+    )
+    print(f"Wrote {args.output}")
+    print(f"Wrote {args.summary_output}")
+
+
+if __name__ == "__main__":
+    main()

eval_results.txt

@@ -0,0 +1,42 @@
+OEIS-Eval-Neo:
+NextTerm-440M - 34.43%
+NextTerm-47M - 29.49%
+Qwen3-0.6B - 18.44%
+Qwen3-1.7B - 20.77%
+Qwen3-4B - 23.74%
+Qwen3-8B - 24.62%
+Qwen3-14B - 26.00%
+
+Ryskina & Knight (2021):
+NextTerm-440M - 52.63%
+NextTerm-47M - 70.18%
+
+16-Shot Bitstring:
+NextTerm-440M - 32.00%
+NextTerm-47M - 27.88%
+
+M1 Competition 111 MAPE (macro; canonical greedy; lower is better):
+Naive2 - 17.7987
+NextTerm-440M - 17.6239
+NextTerm-47M - 18.7621
+Qwen3-0.6B - 22.7984
+Qwen3-1.7B - 22.2411
+Qwen3-4B - 19.1731
+Qwen3-8B - 18.4027
+Qwen3-14B - 17.9837
+
+M1 Competition 111 by frequency (macro MAPE):
+Naive2 - monthly 17.3871 / quarterly 19.5958 / yearly 17.1314
+NextTerm-440M - monthly 18.3407 / quarterly 19.0475 / yearly 13.5498
+NextTerm-47M - monthly 21.2719 / quarterly 16.0270 / yearly 13.3741
+Qwen3-0.6B - monthly 24.7585 / quarterly 21.7989 / yearly 17.2835
+Qwen3-1.7B - monthly 24.3821 / quarterly 22.5869 / yearly 14.5642
+Qwen3-4B - monthly 20.6455 / quarterly 19.6394 / yearly 13.6308
+Qwen3-8B - monthly 20.0289 / quarterly 17.7249 / yearly 13.6534
+Qwen3-14B - monthly 19.4006 / quarterly 18.1729 / yearly 12.9486
+
+Polynomial continuation evals (accuracy; 200 samples per k):
+Arithmetic (k=2-25): NextTerm-440M - 94.38%; NextTerm-47M - 94.15%
+Quadratic (k=3-25): NextTerm-440M - 86.39%; NextTerm-47M - 81.07%
+Cubic (k=4-25): NextTerm-440M - 75.20%; NextTerm-47M - 37.43%
+Quartic (k=5-25): NextTerm-440M - 67.83%; NextTerm-47M - 15.17%

oeis_eval_mlx_neo.py

@@ -0,0 +1,310 @@
+"""Evaluate NextTerm on OEIS Eval Neo with MLX-LM BatchGenerator."""
+
+import argparse
+import gc
+import inspect
+import json
+import time
+from pathlib import Path
+
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.generate import BatchGenerator
+from tqdm import tqdm
+
+DATA_PATH = Path("/Users/natebreslow/Documents/khashabiLab/bigOEIS/oeis_val_neo.jsonl")
+MODEL_NAME = "/Users/natebreslow/Documents/khashabiLab/bigOEIS/NextTerm-440M"
+MAX_NEW_TOKENS = 196
+MAX_CONTEXT_TOKENS = 4096
+BATCH_SIZE = 64
+OUTPUT_PATH = Path("oeis_eval_results/oeis_eval_mlx_neo_per_doc.jsonl")
+SUMMARY_PATH = Path("oeis_eval_results/oeis_eval_mlx_neo_summary.json")
+PARSE_ERROR_PRINT_LIMIT = 25
+parse_error_print_count = 0
+
+
+def parse_generated(text: str) -> int | None:
+    global parse_error_print_count
+    if "," in text:
+        text = text.split(",")[0]
+    try:
+        return int(text)
+    except ValueError:
+        if parse_error_print_count < PARSE_ERROR_PRINT_LIMIT:
+            print(f"Could not parse generated text: {text!r}")
+            parse_error_print_count += 1
+        return None
+
+
+def load_sequences(path: Path):
+    sequences = []
+    answers = []
+    with path.open() as f:
+        for line in f:
+            record = json.loads(line, parse_int=str)
+            seq = record.get("seq", [])
+            if len(seq) < 2:
+                continue
+            sequences.append(seq[:-1])
+            answers.append(str(seq[-1]))
+    return sequences, answers
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=Path, default=DATA_PATH)
+    parser.add_argument("--model", default=MODEL_NAME)
+    parser.add_argument("--max-new-tokens", type=int, default=MAX_NEW_TOKENS)
+    parser.add_argument("--max-context-tokens", type=int, default=MAX_CONTEXT_TOKENS)
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
+    parser.add_argument("--max-examples", type=int, default=0)
+    parser.add_argument("--output", type=Path, default=OUTPUT_PATH)
+    parser.add_argument("--summary-output", type=Path, default=SUMMARY_PATH)
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--restrict-digit-comma-eos", action="store_true")
+    parser.add_argument("--restrict-integer-comma-eos", action="store_true")
+    return parser.parse_args()
+
+
+def load_completed(path: Path) -> dict[int, dict]:
+    completed = {}
+    if not path.exists():
+        return completed
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            record = json.loads(line)
+            completed[int(record["row_index"])] = record
+    return completed
+
+
+def encode_no_special(tokenizer, text: str) -> list[int]:
+    try:
+        return tokenizer.encode(text, add_special_tokens=False)
+    except TypeError:
+        return tokenizer.encode(text)
+
+
+def normalize_stop_tokens_for_batch_generator(stop_tokens: list[list[int]]):
+    annotation = inspect.signature(BatchGenerator.__init__).parameters[
+        "stop_tokens"
+    ].annotation
+    if "set" in str(annotation):
+        return {seq[0] for seq in stop_tokens if len(seq) == 1}
+    return stop_tokens
+
+
+def split_batch_generator_responses(responses):
+    if isinstance(responses, tuple) and len(responses) == 2:
+        prompt_responses, generation_responses = responses
+        return prompt_responses, generation_responses
+    if isinstance(responses, list):
+        return [], responses
+    raise RuntimeError(
+        "Unexpected mlx_lm BatchGenerator.next() API. Update this script for "
+        f"{type(responses).__name__}: {responses!r}"
+    )
+
+
+def make_integer_comma_eos_processor(tokenizer):
+    allowed = set()
+    for text in [str(i) for i in range(10)] + ["-", ","]:
+        tokens = encode_no_special(tokenizer, text)
+        if len(tokens) == 1:
+            allowed.add(int(tokens[0]))
+        else:
+            print(f"Skipping multi-token allowed text {text!r}: {tokens}")
+    if tokenizer.eos_token_id is not None:
+        allowed.add(int(tokenizer.eos_token_id))
+    allowed_ids = sorted(allowed)
+    mask_cache = {}
+
+    def processor(_tokens, logits):
+        vocab_size = logits.shape[-1]
+        mask = mask_cache.get(vocab_size)
+        if mask is None:
+            values = [-1e9] * vocab_size
+            for token_id in allowed_ids:
+                if 0 <= token_id < vocab_size:
+                    values[token_id] = 0.0
+            mask = mx.array(values, dtype=logits.dtype)
+            mask_cache[vocab_size] = mask
+        return logits + mask[None, :]
+
+    return processor, allowed_ids
+
+
+def run_generation_queue(
+    *,
+    model,
+    tokenizer,
+    prompts,
+    answers,
+    row_indices,
+    stop_tokens,
+    max_new_tokens: int,
+    batch_size: int,
+    output_file,
+    progress,
+    logits_processors=None,
+) -> None:
+    gen = BatchGenerator(
+        model,
+        stop_tokens=normalize_stop_tokens_for_batch_generator(stop_tokens),
+        logits_processors=logits_processors,
+        completion_batch_size=batch_size,
+        prefill_batch_size=batch_size,
+    )
+    uids = gen.insert(prompts, [max_new_tokens] * len(prompts))
+    uid_to_pos = {uid: pos for pos, uid in enumerate(uids)}
+    generated_tokens = {uid: [] for uid in uids}
+    finished = set()
+    try:
+        while True:
+            responses = gen.next()
+            prompt_responses, generation_responses = split_batch_generator_responses(
+                responses
+            )
+            if not prompt_responses and not generation_responses:
+                break
+            if not generation_responses:
+                continue
+            for response in generation_responses:
+                uid = response.uid
+                if response.finish_reason != "stop":
+                    generated_tokens[uid].append(int(response.token))
+                if response.finish_reason is None or uid in finished:
+                    continue
+                finished.add(uid)
+                pos = uid_to_pos[uid]
+                text = tokenizer.decode(generated_tokens[uid])
+                prediction = parse_generated(text)
+                answer = answers[pos]
+                answer_int = int(answer)
+                record = {
+                    "row_index": row_indices[pos],
+                    "answer": answer,
+                    "prediction": prediction,
+                    "correct": prediction == answer_int,
+                    "parsed": prediction is not None,
+                    "generated_text": text,
+                    "generated_tokens": generated_tokens[uid],
+                    "finish_reason": response.finish_reason,
+                }
+                output_file.write(json.dumps(record) + "\n")
+                progress.update(1)
+    finally:
+        gen.close()
+        mx.clear_cache()
+        gc.collect()
+
+    if len(finished) != len(uids):
+        raise RuntimeError(f"Chunk finished {len(finished)}/{len(uids)} rows")
+
+
+def main():
+    args = parse_args()
+    started = time.perf_counter()
+    sequences, answers = load_sequences(args.data_path)
+    if args.max_examples > 0:
+        sequences = sequences[: args.max_examples]
+        answers = answers[: args.max_examples]
+    print(f"Loaded {len(answers)} sequences from {args.data_path}")
+
+    model, tokenizer = load(args.model)
+
+    sep_tokens = encode_no_special(tokenizer, ",")
+    if not sep_tokens:
+        sep_tokens = encode_no_special(tokenizer, "1,")[-1:]
+    prompts = [",".join(str(x) for x in seq) + "," for seq in sequences]
+    prompts = [tokenizer.encode(p) for p in prompts]
+
+    eval_indices = [
+        i
+        for i, prompt in enumerate(prompts)
+        if args.max_context_tokens <= 0 or len(prompt) < args.max_context_tokens
+    ]
+    skipped_long = len(prompts) - len(eval_indices)
+    eval_indices = sorted(eval_indices, key=lambda i: len(prompts[i]))
+    print(
+        f"Evaluating {len(eval_indices)} rows; skipped_long={skipped_long}; "
+        f"sep_tokens={sep_tokens}; eos_token={tokenizer.eos_token_id}; "
+        f"max_new_tokens={args.max_new_tokens}; batch_size={args.batch_size}"
+    )
+
+    stop_tokens = [sep_tokens]
+    if tokenizer.eos_token_id is not None:
+        stop_tokens.append([tokenizer.eos_token_id])
+    logits_processors = None
+    allowed_token_ids = None
+    restrict_integer = args.restrict_digit_comma_eos or args.restrict_integer_comma_eos
+    if restrict_integer:
+        processor, allowed_token_ids = make_integer_comma_eos_processor(tokenizer)
+        logits_processors = [processor]
+        print(f"Restricting logits to integer/comma/EOS token ids: {allowed_token_ids}")
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.summary_output.parent.mkdir(parents=True, exist_ok=True)
+    if args.overwrite and args.output.exists():
+        args.output.unlink()
+    completed = load_completed(args.output)
+    if completed:
+        print(f"Resuming from {args.output}: {len(completed)} rows already done")
+    todo_indices = [idx for idx in eval_indices if idx not in completed]
+
+    with args.output.open("a", encoding="utf-8") as output_file:
+        with tqdm(total=len(todo_indices), desc="Generating") as progress:
+            run_generation_queue(
+                model=model,
+                tokenizer=tokenizer,
+                prompts=[prompts[i] for i in todo_indices],
+                answers=[answers[i] for i in todo_indices],
+                row_indices=todo_indices,
+                stop_tokens=stop_tokens,
+                max_new_tokens=args.max_new_tokens,
+                batch_size=args.batch_size,
+                output_file=output_file,
+                progress=progress,
+                logits_processors=logits_processors,
+            )
+            output_file.flush()
+
+    records = load_completed(args.output)
+    eval_set = set(eval_indices)
+    records = {idx: record for idx, record in records.items() if idx in eval_set}
+    correct = sum(1 for record in records.values() if record["correct"])
+    parsed = sum(1 for record in records.values() if record["parsed"])
+    total = len(eval_indices)
+    evaluated = len(records)
+    elapsed = time.perf_counter() - started
+
+    print(f"Documents: {len(answers)}")
+    print(f"Evaluated: {evaluated}/{total}")
+    print(f"Skipped long: {skipped_long}")
+    print(f"Parsed predictions: {parsed}/{evaluated}")
+    print(f"Accuracy: {correct}/{evaluated} = {correct / evaluated:.4f}")
+    summary = {
+        "data_path": str(args.data_path),
+        "model": args.model,
+        "output": str(args.output),
+        "documents": len(answers),
+        "evaluated": evaluated,
+        "expected_evaluated": total,
+        "skipped_long": skipped_long,
+        "parsed": parsed,
+        "correct": correct,
+        "accuracy": correct / evaluated if evaluated else 0.0,
+        "max_new_tokens": args.max_new_tokens,
+        "max_context_tokens": args.max_context_tokens,
+        "batch_size": args.batch_size,
+        "restrict_digit_comma_eos": args.restrict_digit_comma_eos,
+        "restrict_integer_comma_eos": restrict_integer,
+        "allowed_token_ids": allowed_token_ids,
+        "seconds": elapsed,
+    }
+    args.summary_output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
+    print(f"Wrote {args.output}")
+    print(f"Wrote {args.summary_output}")
+    
+if __name__ == "__main__":
+    main()