Clean Atom2.7m package naming and benchmark entrypoint

README.md

@@ -90,12 +90,12 @@ with torch.no_grad():
 Use the included benchmark script:
 
 ```bash
-python benchmark_fusion_arithmark.py \
+python benchmark_atom_arithmark.py \
   --checkpoint . \
   --data-path arithmark_2.0.jsonl \
   --batch-size 64 \
   --device cuda \
-  --output benchmark_results/fusion_arithmark_2.0_results.json
+  --output benchmark_results/atom_arithmark_2.0_results.json
 ```
 
 ### lm-evaluation-harness
@@ -157,7 +157,7 @@ Synthetic-Arithmetic is canonical integer equation data. The training curriculum
 - `model.safetensors`: model weights
 - `config.json`, `config.py`, `configuration_gpt.py`, `model.py`: custom model code
 - `tokenizer.json`, `tokenization_atom.py`: tokenizer files and remote-code wrapper
-- `benchmark_fusion_arithmark.py`: ArithMark evaluation
+- `benchmark_atom_arithmark.py`: ArithMark evaluation
 - `arithmark_2.0.jsonl`: local ArithMark 2.0 data for the standalone benchmark script
 - `pretraining_curriculum.json`: training curriculum
 

benchmark_atom_arithmark.py

@@ -0,0 +1,290 @@
+"""Score an Atom GPT checkpoint on ArithMark 2.0."""
+
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from contextlib import nullcontext
+import json
+from pathlib import Path
+import re
+import urllib.request
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+DATA_URL = (
+    "https://huggingface.co/datasets/AxiomicLabs/Arithmark-2.0/"
+    "resolve/main/arithmark_2.0.jsonl"
+)
+
+
+def ensure_data(path: Path) -> Path:
+    if path.exists():
+        return path
+    path.parent.mkdir(parents=True, exist_ok=True)
+    urllib.request.urlretrieve(DATA_URL, path)
+    return path
+
+
+def load_examples(path: Path, *, max_examples: int = 0) -> list[dict]:
+    examples = []
+    with path.open("r", encoding="utf-8") as handle:
+        for line in handle:
+            if not line.strip():
+                continue
+            examples.append(json.loads(line))
+            if max_examples > 0 and len(examples) >= max_examples:
+                break
+    return examples
+
+
+def _encoded_choice(
+    tokenizer,
+    context: str,
+    ending: str,
+) -> tuple[list[int], int]:
+    context_ids = tokenizer(context, add_special_tokens=False).input_ids
+    full_ids = tokenizer(context + ending, add_special_tokens=False).input_ids
+    continuation_length = len(full_ids) - len(context_ids)
+    return full_ids, continuation_length
+
+
+@torch.inference_mode()
+def evaluate(
+    model,
+    tokenizer,
+    examples: list[dict],
+    *,
+    device: torch.device,
+    batch_size: int,
+    dump_failures: bool = False,
+    failure_operator_count: int | None = None,
+    max_failures: int = 100,
+) -> dict:
+    correct = 0
+    total = 0
+    by_operator_count: dict[str, list[int]] = {}
+    by_topic: dict[str, list[int]] = {}
+    failures: list[dict] = []
+    failure_summary: Counter[tuple[str, str, str]] = Counter()
+    model.eval()
+    pad_id = tokenizer.pad_token_id
+    if pad_id is None:
+        pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+
+    for start in range(0, len(examples), batch_size):
+        batch_examples = examples[start : start + batch_size]
+        encoded = []
+        offsets = []
+        for example in batch_examples:
+            flat_start = len(encoded)
+            encoded.extend(
+                _encoded_choice(tokenizer, example["ctx"], ending)
+                for ending in example["endings"]
+            )
+            offsets.append((flat_start, len(example["endings"])))
+
+        max_length = max(len(item[0]) for item in encoded)
+        input_ids = torch.full(
+            (len(encoded), max_length),
+            int(pad_id),
+            dtype=torch.long,
+            device=device,
+        )
+        attention_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        lengths = []
+        continuation_lengths = []
+        for row, (ids, continuation_length) in enumerate(encoded):
+            length = len(ids)
+            input_ids[row, :length] = torch.tensor(ids, device=device)
+            attention_mask[row, :length] = True
+            lengths.append(length)
+            continuation_lengths.append(continuation_length)
+
+        autocast = (
+            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+            if device.type == "cuda"
+            else nullcontext()
+        )
+        with autocast:
+            logits = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            ).logits
+        log_probs = F.log_softmax(logits.float(), dim=-1)
+
+        for example_index, example in enumerate(batch_examples):
+            flat_start, choice_count = offsets[example_index]
+            likelihoods = []
+            for choice_index in range(choice_count):
+                row = flat_start + choice_index
+                length = lengths[row]
+                continuation_length = continuation_lengths[row]
+                continuation_start = length - continuation_length
+                likelihood = 0.0
+                for position in range(continuation_start, length):
+                    likelihood += float(
+                        log_probs[row, position - 1, input_ids[row, position]].item()
+                    )
+                likelihoods.append(likelihood)
+
+            prediction = max(range(choice_count), key=likelihoods.__getitem__)
+            label = int(example["label"])
+            matched = prediction == label
+            correct += int(matched)
+            total += 1
+            metadata = example.get("metadata", {})
+            operator_count = str(metadata.get("operator_count", "unknown"))
+            topic = str(metadata.get("topic", "unknown"))
+            for grouped, key in (
+                (by_operator_count, operator_count),
+                (by_topic, topic),
+            ):
+                group = grouped.setdefault(key, [0, 0])
+                group[0] += int(matched)
+                group[1] += 1
+
+            if not matched and dump_failures:
+                op_count_int = None
+                try:
+                    op_count_int = int(operator_count)
+                except ValueError:
+                    pass
+                if failure_operator_count is None or op_count_int == failure_operator_count:
+                    context = str(example["ctx"]).strip()
+                    expression = context[:-1].strip() if context.endswith("=") else context
+                    operands = [int(value) for value in re.findall(r"\d+", expression)]
+                    operator = "".join(re.findall(r"[+\-*/]", expression))
+                    predicted_answer = str(example["endings"][prediction]).strip()
+                    correct_answer = str(example["endings"][label]).strip()
+                    width = max((len(str(value)) for value in operands), default=0)
+                    failure_summary[(topic, operator, f"width={width}")] += 1
+                    if len(failures) < max_failures:
+                        failures.append(
+                            {
+                                "ctx": context,
+                                "topic": topic,
+                                "operator_count": operator_count,
+                                "operator": operator,
+                                "operands": operands,
+                                "max_operand_digits": width,
+                                "correct_answer": correct_answer,
+                                "predicted_answer": predicted_answer,
+                                "choices": [str(value).strip() for value in example["endings"]],
+                                "choice_scores": [round(value, 4) for value in likelihoods],
+                                "score_margin_correct_minus_predicted": round(
+                                    likelihoods[label] - likelihoods[prediction],
+                                    4,
+                                ),
+                            }
+                        )
+
+    results = {
+        "benchmark": "arithmark_2.0",
+        "model_type": "atom_gpt",
+        "accuracy": correct / max(total, 1),
+        "correct": correct,
+        "total": total,
+        "by_operator_count": {
+            key: {
+                "accuracy": values[0] / max(values[1], 1),
+                "correct": values[0],
+                "total": values[1],
+            }
+            for key, values in sorted(by_operator_count.items())
+        },
+        "by_topic": {
+            key: {
+                "accuracy": values[0] / max(values[1], 1),
+                "correct": values[0],
+                "total": values[1],
+            }
+            for key, values in sorted(by_topic.items())
+        },
+    }
+    if dump_failures:
+        results["failure_summary"] = {
+            "|".join(key): value
+            for key, value in failure_summary.most_common()
+        }
+        results["failures"] = failures
+    return results
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, default=Path("."))
+    parser.add_argument("--data-path", type=Path, default=Path("arithmark_2.0.jsonl"))
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--dtype", default="auto", choices=("auto", "float32", "bfloat16", "float16"))
+    parser.add_argument("--output", type=Path)
+    parser.add_argument(
+        "--max-examples",
+        type=int,
+        default=0,
+        help="Evaluate only the first N examples. Default evaluates all examples.",
+    )
+    parser.add_argument(
+        "--dump-failures",
+        action="store_true",
+        help="Include incorrectly scored examples and grouped failure summary.",
+    )
+    parser.add_argument(
+        "--failure-operator-count",
+        type=int,
+        default=None,
+        help="Only dump failures with this operator count, e.g. 1 for easy examples.",
+    )
+    parser.add_argument("--max-failures", type=int, default=100)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    data_path = ensure_data(args.data_path)
+    examples = load_examples(data_path, max_examples=args.max_examples)
+    dtype = None
+    if args.dtype == "float32":
+        dtype = torch.float32
+    elif args.dtype == "bfloat16":
+        dtype = torch.bfloat16
+    elif args.dtype == "float16":
+        dtype = torch.float16
+    model = AutoModelForCausalLM.from_pretrained(
+        args.checkpoint,
+        dtype=dtype,
+        trust_remote_code=True,
+    ).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    results = evaluate(
+        model,
+        tokenizer,
+        examples,
+        device=device,
+        batch_size=args.batch_size,
+        dump_failures=args.dump_failures,
+        failure_operator_count=args.failure_operator_count,
+        max_failures=args.max_failures,
+    )
+    print(json.dumps(results, indent=2, sort_keys=True))
+    if args.output:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(
+            json.dumps(results, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+
+if __name__ == "__main__":
+    main()

benchmark_results/lm_eval/results_2026-07-04T17-16-59.922847.json

@@ -0,0 +1,355 @@
+{
+  "results": {
+    "hellaswag": {
+      "name": "hellaswag",
+      "alias": "hellaswag",
+      "sample_len": 10042,
+      "acc,none": 0.2669786895040829,
+      "acc_stderr,none": 0.004414770331224705,
+      "acc_norm,none": 0.2700657239593706,
+      "acc_norm_stderr,none": 0.0044308610336166706
+    },
+    "arc_easy": {
+      "name": "arc_easy",
+      "alias": "arc_easy",
+      "sample_len": 2376,
+      "acc,none": 0.32154882154882153,
+      "acc_stderr,none": 0.009584091575640627,
+      "acc_norm,none": 0.31607744107744107,
+      "acc_norm_stderr,none": 0.009540440071928294
+    },
+    "arc_challenge": {
+      "name": "arc_challenge",
+      "alias": "arc_challenge",
+      "sample_len": 1172,
+      "acc,none": 0.17235494880546076,
+      "acc_stderr,none": 0.011037113093461295,
+      "acc_norm,none": 0.2098976109215017,
+      "acc_norm_stderr,none": 0.011900548748047421
+    },
+    "piqa": {
+      "name": "piqa",
+      "alias": "piqa",
+      "sample_len": 1838,
+      "acc,none": 0.5565832426550599,
+      "acc_stderr,none": 0.011590883373666865,
+      "acc_norm,none": 0.529923830250272,
+      "acc_norm_stderr,none": 0.011644913435420158
+    }
+  },
+  "group_subtasks": {},
+  "configs": {
+    "arc_challenge": {
+      "task": "arc_challenge",
+      "dataset_path": "allenai/ai2_arc",
+      "dataset_name": "ARC-Challenge",
+      "training_split": "train",
+      "validation_split": "validation",
+      "test_split": "test",
+      "doc_to_text": "Question: {{question}}\nAnswer:",
+      "doc_to_target": "{{choices.label.index(answerKey)}}",
+      "unsafe_code": false,
+      "doc_to_choice": "{{choices.text}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{question}}\nAnswer:",
+        "doc_to_choice": "{{choices.text}}",
+        "doc_to_target": "{{choices.label.index(answerKey)}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "Question: {{question}}\nAnswer:",
+      "metadata": {
+        "version": 1.0,
+        "pretrained": ".",
+        "trust_remote_code": true,
+        "dtype": "bfloat16",
+        "max_length": 548,
+        "config_source": "/home/max/PROJECTS/PythonProjects/mcpterm/venv/lib/python3.11/site-packages/lm_eval/tasks/arc/arc_challenge.yaml"
+      }
+    },
+    "arc_easy": {
+      "task": "arc_easy",
+      "dataset_path": "allenai/ai2_arc",
+      "dataset_name": "ARC-Easy",
+      "training_split": "train",
+      "validation_split": "validation",
+      "test_split": "test",
+      "doc_to_text": "Question: {{question}}\nAnswer:",
+      "doc_to_target": "{{choices.label.index(answerKey)}}",
+      "unsafe_code": false,
+      "doc_to_choice": "{{choices.text}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{question}}\nAnswer:",
+        "doc_to_choice": "{{choices.text}}",
+        "doc_to_target": "{{choices.label.index(answerKey)}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "Question: {{question}}\nAnswer:",
+      "metadata": {
+        "version": 1.0,
+        "pretrained": ".",
+        "trust_remote_code": true,
+        "dtype": "bfloat16",
+        "max_length": 548,
+        "config_source": "/home/max/PROJECTS/PythonProjects/mcpterm/venv/lib/python3.11/site-packages/lm_eval/tasks/arc/arc_easy.yaml"
+      }
+    },
+    "hellaswag": {
+      "task": "hellaswag",
+      "dataset_path": "Rowan/hellaswag",
+      "training_split": "train",
+      "validation_split": "validation",
+      "process_docs": "def process_docs(dataset: datasets.Dataset) -> datasets.Dataset:\n    def _process_doc(doc):\n        ctx = doc[\"ctx_a\"] + \" \" + doc[\"ctx_b\"].capitalize()\n        out_doc = {\n            \"query\": preprocess(doc[\"activity_label\"] + \": \" + ctx),\n            \"choices\": [preprocess(ending) for ending in doc[\"endings\"]],\n            \"gold\": int(doc[\"label\"]),\n        }\n        return out_doc\n\n    return dataset.map(_process_doc)\n",
+      "doc_to_text": "{{query}}",
+      "doc_to_target": "{{label}}",
+      "unsafe_code": false,
+      "doc_to_choice": "choices",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": "<function process_docs at 0x77bb53ff80e0>",
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "{{query}}",
+        "doc_to_choice": "choices",
+        "doc_to_target": "{{label}}",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": false,
+      "metadata": {
+        "version": 1.0,
+        "pretrained": ".",
+        "trust_remote_code": true,
+        "dtype": "bfloat16",
+        "max_length": 548,
+        "config_source": "/home/max/PROJECTS/PythonProjects/mcpterm/venv/lib/python3.11/site-packages/lm_eval/tasks/hellaswag/hellaswag.yaml"
+      }
+    },
+    "piqa": {
+      "task": "piqa",
+      "dataset_path": "baber/piqa",
+      "training_split": "train",
+      "validation_split": "validation",
+      "doc_to_text": "Question: {{goal}}\nAnswer:",
+      "doc_to_target": "label",
+      "unsafe_code": false,
+      "doc_to_choice": "{{[sol1, sol2]}}",
+      "description": "",
+      "target_delimiter": " ",
+      "fewshot_delimiter": "\n\n",
+      "fewshot_config": {
+        "sampler": "default",
+        "split": null,
+        "process_docs": null,
+        "fewshot_indices": null,
+        "samples": null,
+        "doc_to_text": "Question: {{goal}}\nAnswer:",
+        "doc_to_choice": "{{[sol1, sol2]}}",
+        "doc_to_target": "label",
+        "gen_prefix": null,
+        "fewshot_delimiter": "\n\n",
+        "target_delimiter": " "
+      },
+      "num_fewshot": 0,
+      "metric_list": [
+        {
+          "metric": "acc",
+          "aggregation": "mean",
+          "higher_is_better": true
+        },
+        {
+          "metric": "acc_norm",
+          "aggregation": "mean",
+          "higher_is_better": true
+        }
+      ],
+      "output_type": "multiple_choice",
+      "repeats": 1,
+      "should_decontaminate": true,
+      "doc_to_decontamination_query": "goal",
+      "metadata": {
+        "version": 1.0,
+        "pretrained": ".",
+        "trust_remote_code": true,
+        "dtype": "bfloat16",
+        "max_length": 548,
+        "config_source": "/home/max/PROJECTS/PythonProjects/mcpterm/venv/lib/python3.11/site-packages/lm_eval/tasks/piqa/piqa.yaml"
+      }
+    }
+  },
+  "versions": {
+    "arc_challenge": 1.0,
+    "arc_easy": 1.0,
+    "hellaswag": 1.0,
+    "piqa": 1.0
+  },
+  "n-shot": {
+    "arc_challenge": 0,
+    "arc_easy": 0,
+    "hellaswag": 0,
+    "piqa": 0
+  },
+  "higher_is_better": {
+    "arc_challenge": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "arc_easy": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "hellaswag": {
+      "acc": true,
+      "acc_norm": true
+    },
+    "piqa": {
+      "acc": true,
+      "acc_norm": true
+    }
+  },
+  "n-samples": {
+    "hellaswag": {
+      "original": 10042,
+      "effective": 10042
+    },
+    "arc_easy": {
+      "original": 2376,
+      "effective": 2376
+    },
+    "arc_challenge": {
+      "original": 1172,
+      "effective": 1172
+    },
+    "piqa": {
+      "original": 1838,
+      "effective": 1838
+    }
+  },
+  "config": {
+    "model": "hf",
+    "model_args": {
+      "pretrained": ".",
+      "trust_remote_code": true,
+      "dtype": "bfloat16",
+      "max_length": 548
+    },
+    "model_num_parameters": 2738880,
+    "model_dtype": "torch.float32",
+    "model_revision": "main",
+    "model_sha": "",
+    "batch_size": "auto:1",
+    "batch_sizes": [
+      64
+    ],
+    "device": "cuda:0",
+    "use_cache": null,
+    "limit": null,
+    "bootstrap_iters": 100000,
+    "gen_kwargs": {},
+    "random_seed": 0,
+    "numpy_seed": 1234,
+    "torch_seed": 1234,
+    "fewshot_seed": 1234
+  },
+  "git_hash": "f5846e8",
+  "date": 1783178146.3657796,
+  "pretty_env_info": "PyTorch version: 2.10.0+cu130\nIs debug build: False\nCUDA used to build PyTorch: 13.0\nROCM used to build PyTorch: N/A\n\nOS: Debian GNU/Linux 12 (bookworm) (x86_64)\nGCC version: (Debian 12.2.0-14+deb12u1) 12.2.0\nClang version: Could not collect\nCMake version: version 3.25.1\nLibc version: glibc-2.36\n\nPython version: 3.11.2 (main, Apr  8 2026, 01:58:00) [GCC 12.2.0] (64-bit runtime)\nPython platform: Linux-6.18.33.2-microsoft-standard-WSL2-x86_64-with-glibc2.36\nIs CUDA available: True\nCUDA runtime version: 13.2.51\nCUDA_MODULE_LOADING set to: \nGPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080\nNvidia driver version: 596.36\ncuDNN version: Could not collect\nIs XPU available: False\nHIP runtime version: N/A\nMIOpen runtime version: N/A\nIs XNNPACK available: True\nCaching allocator config: N/A\n\nCPU:\nArchitecture:                            x86_64\nCPU op-mode(s):                          32-bit, 64-bit\nAddress sizes:                           48 bits physical, 48 bits virtual\nByte Order:                              Little Endian\nCPU(s):                                  12\nOn-line CPU(s) list:                     0-11\nVendor ID:                               AuthenticAMD\nModel name:                              AMD Ryzen 5 7500F 6-Core Processor\nCPU family:                              25\nModel:                                   97\nThread(s) per core:                      2\nCore(s) per socket:                      6\nSocket(s):                               1\nStepping:                                2\nBogoMIPS:                                7399.82\nFlags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero xsaveerptr arat npt nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm\nVirtualization:                          AMD-V\nHypervisor vendor:                       Microsoft\nVirtualization type:                     full\nL1d cache:                               192 KiB (6 instances)\nL1i cache:                               192 KiB (6 instances)\nL2 cache:                                6 MiB (6 instances)\nL3 cache:                                32 MiB (1 instance)\nNUMA node(s):                            1\nNUMA node0 CPU(s):                       0-11\nVulnerability Gather data sampling:      Not affected\nVulnerability Ghostwrite:                Not affected\nVulnerability Indirect target selection: Not affected\nVulnerability Itlb multihit:             Not affected\nVulnerability L1tf:                      Not affected\nVulnerability Mds:                       Not affected\nVulnerability Meltdown:                  Not affected\nVulnerability Mmio stale data:           Not affected\nVulnerability Old microcode:             Not affected\nVulnerability Reg file data sampling:    Not affected\nVulnerability Retbleed:                  Not affected\nVulnerability Spec rstack overflow:      Vulnerable: Safe RET, no microcode\nVulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl\nVulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization\nVulnerability Spectre v2:                Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected\nVulnerability Srbds:                     Not affected\nVulnerability Tsa:                       Vulnerable: No microcode\nVulnerability Tsx async abort:           Not affected\nVulnerability Vmscape:                   Not affected\n\nVersions of relevant libraries:\n[pip3] flash_attn==2.8.3+cu130torch2.10\n[pip3] mypy_extensions==1.1.0\n[pip3] numpy==2.2.6\n[pip3] nvidia-cublas==13.1.0.3\n[pip3] nvidia-cublas-cu12==12.8.4.1\n[pip3] nvidia-cuda-cupti==13.0.85\n[pip3] nvidia-cuda-cupti-cu12==12.8.90\n[pip3] nvidia-cuda-nvrtc==13.0.88\n[pip3] nvidia-cuda-nvrtc-cu12==12.8.93\n[pip3] nvidia-cuda-runtime==13.0.96\n[pip3] nvidia-cuda-runtime-cu12==12.8.90\n[pip3] nvidia-cudnn-cu12==9.10.2.21\n[pip3] nvidia-cudnn-cu13==9.15.1.9\n[pip3] nvidia-cudnn-frontend==1.18.0\n[pip3] nvidia-cufft==12.0.0.61\n[pip3] nvidia-cufft-cu12==11.3.3.83\n[pip3] nvidia-curand==10.4.0.35\n[pip3] nvidia-curand-cu12==10.3.9.90\n[pip3] nvidia-cusolver==12.0.4.66\n[pip3] nvidia-cusolver-cu12==11.7.3.90\n[pip3] nvidia-cusparse==12.6.3.3\n[pip3] nvidia-cusparse-cu12==12.5.8.93\n[pip3] nvidia-cusparselt-cu12==0.7.1\n[pip3] nvidia-cusparselt-cu13==0.8.0\n[pip3] nvidia-nccl-cu12==2.27.5\n[pip3] nvidia-nccl-cu13==2.28.9\n[pip3] nvidia-nvjitlink==13.0.88\n[pip3] nvidia-nvjitlink-cu12==12.8.93\n[pip3] nvidia-nvtx==13.0.85\n[pip3] nvidia-nvtx-cu12==12.8.90\n[pip3] onnxruntime==1.20.1\n[pip3] rapidocr-onnxruntime==1.4.4\n[pip3] rotary-embedding-torch==0.8.8\n[pip3] torch==2.10.0+cu130\n[pip3] torch_c_dlpack_ext==0.1.5\n[pip3] torch-lr-finder==0.2.2\n[pip3] torch-tb-profiler==0.4.3\n[pip3] torchaudio==2.10.0+cu130\n[pip3] torchmetrics==1.9.0\n[pip3] torchvision==0.25.0+cu130\n[pip3] triton==3.6.0\n[conda] Could not collect",
+  "transformers_version": "4.57.6",
+  "lm_eval_version": "0.4.12",
+  "upper_git_hash": "f5846e8ab83c0da3653a1f7a04d470911c8f5065",
+  "tokenizer_pad_token": [
+    "<|pad|>",
+    "0"
+  ],
+  "tokenizer_eos_token": [
+    "<|eos|>",
+    "2"
+  ],
+  "tokenizer_bos_token": [
+    "<|bos|>",
+    "1"
+  ],
+  "eot_token_id": 2,
+  "max_length": 548,
+  "task_hashes": {},
+  "model_source": "hf",
+  "model_name": ".",
+  "model_name_sanitized": ".",
+  "system_instruction": null,
+  "system_instruction_sha": null,
+  "fewshot_as_multiturn": null,
+  "chat_template": null,
+  "chat_template_sha": null,
+  "total_evaluation_time_seconds": "80.63422504799973"
+}

tokenization_atom.py

@@ -1,4 +1,4 @@
-"""Remote-code tokenizer for Atom/Fusion GPT checkpoints.
+"""Remote-code tokenizer for Atom GPT checkpoints.
 
 The tokenizer is intentionally HF-compatible: generic callers can use
 ``AutoTokenizer.from_pretrained(..., trust_remote_code=True)``. Arithmetic digit

tokenizer_utils.py

@@ -28,7 +28,7 @@ MAX_OPERAND_ROLES = 9
 
 
 @dataclass(frozen=True)
-class FusionEncoding:
+class AtomEncoding:
     ids: list[int]
     place_ids: list[int]
     role_ids: list[int]
@@ -46,7 +46,7 @@ class FusionEncoding:
 
     def __post_init__(self) -> None:
         if not (len(self.ids) == len(self.place_ids) == len(self.role_ids)):
-            raise ValueError("Fusion tokenizer streams must have equal length")
+            raise ValueError("Atom tokenizer streams must have equal length")
 
 
 def build_tokenizer() -> Any:
@@ -56,10 +56,15 @@ def build_tokenizer() -> Any:
     tokenizer = Tokenizer(models.BPE(unk_token="<|unk|>"))
     tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
         [
+            # BPE may merge within one of these spans, but never across them.
+            # This keeps whitespace, digits, and arithmetic syntax isolated.
             pre_tokenizers.Split(
                 Regex(r"\s+|\d|[+\-*/=()]|[^\s\d+\-*/=()]+"),
                 behavior="isolated",
             ),
+            # Convert every span to bytes without applying GPT-2's prefix-space
+            # regex. The complete byte alphabet supplied to the trainer provides
+            # lossless fallback for arbitrary UTF-8 input.
             pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
         ]
     )
@@ -67,7 +72,7 @@ def build_tokenizer() -> Any:
     return tokenizer
 
 
-class FusionTokenizer:
+class AtomTokenizer:
     """Runtime wrapper adding LSD-first digit streams to a trained BPE tokenizer."""
 
     _digit_span_re = re.compile(r"\d+")
@@ -85,11 +90,14 @@ class FusionTokenizer:
             if self.tokenizer.token_to_id(digit) is not None
         }
         self._equals_id = self.tokenizer.token_to_id("=")
+        self._arithmetic_token_texts = set(ARITHMETIC_TOKENS)
         self._special_token_ids = frozenset(
             token_id
             for token in SPECIAL_TOKENS
             if (token_id := self.tokenizer.token_to_id(token)) is not None
         )
+        self._piece_cache: dict[int, str] = {}
+        self._equation_piece_cache: dict[int, bool] = {}
         if len(self._digit_token_ids) != 10:
             raise ValueError("Tokenizer vocabulary must contain atomic digit tokens 0-9")
         if self._equals_id is None:
@@ -123,20 +131,31 @@ class FusionTokenizer:
         return cls._digit_span_re.sub(lambda match: match.group(0)[::-1], text)
 
     def _decode_token_piece(self, token_id: int) -> str:
-        return self.tokenizer.decode([int(token_id)], skip_special_tokens=False)
+        token_id = int(token_id)
+        piece = self._piece_cache.get(token_id)
+        if piece is None:
+            piece = self.tokenizer.decode([token_id], skip_special_tokens=False)
+            self._piece_cache[token_id] = piece
+        return piece
 
     @staticmethod
     def _is_equation_whitespace(piece: str) -> bool:
         return bool(piece) and piece.isspace() and "\n" not in piece and "\r" not in piece
 
     def _is_equation_piece(self, token_id: int, piece: str) -> bool:
+        cached = self._equation_piece_cache.get(token_id)
+        if cached is not None:
+            return cached
         if token_id in self._special_token_ids:
-            return False
-        if token_id in self._digit_token_ids:
-            return True
-        if self._is_equation_whitespace(piece):
-            return True
-        return len(piece) == 1 and piece in set(ARITHMETIC_TOKENS)
+            result = False
+        elif token_id in self._digit_token_ids:
+            result = True
+        elif self._is_equation_whitespace(piece):
+            result = True
+        else:
+            result = len(piece) == 1 and piece in self._arithmetic_token_texts
+        self._equation_piece_cache[token_id] = result
+        return result
 
     def _annotate_equation_span(
         self,
@@ -213,20 +232,25 @@ class FusionTokenizer:
 
         return place_ids, role_ids
 
-    def encode(self, text: str, *args, **kwargs) -> FusionEncoding:
+    def encode(self, text: str, *args, **kwargs) -> AtomEncoding:
         transformed = self._reverse_digit_spans(text)
         encoding = self.tokenizer.encode(transformed, *args, **kwargs)
+        return self._wrap_encoding(encoding)
+
+    def _wrap_encoding(self, encoding: Any) -> AtomEncoding:
         ids = [int(token_id) for token_id in encoding.ids]
         place_ids, role_ids = self.annotate_ids(ids)
-        return FusionEncoding(
+        return AtomEncoding(
             ids=ids,
             place_ids=place_ids,
             role_ids=role_ids,
             tokens=list(getattr(encoding, "tokens", [])),
         )
 
-    def encode_batch(self, texts: list[str], *args, **kwargs) -> list[FusionEncoding]:
-        return [self.encode(text, *args, **kwargs) for text in texts]
+    def encode_batch(self, texts: list[str], *args, **kwargs) -> list[AtomEncoding]:
+        transformed = [self._reverse_digit_spans(text) for text in texts]
+        encodings = self.tokenizer.encode_batch(transformed, *args, **kwargs)
+        return [self._wrap_encoding(encoding) for encoding in encodings]
 
     def decode(
         self,
@@ -293,15 +317,11 @@ def validate_tokenizer(tokenizer_dir: Path) -> None:
             f"Missing {tokenizer_json}. Retrain with train_tokenizer.py so the "
             "whitespace and digit boundary rules are preserved."
         )
-    if vocab_path.exists():
-        with vocab_path.open("r", encoding="utf-8") as f:
-            vocab = json.load(f)
-    else:
-        with tokenizer_json.open("r", encoding="utf-8") as f:
-            tokenizer_data = json.load(f)
-        vocab = tokenizer_data.get("model", {}).get("vocab")
-        if not isinstance(vocab, dict):
-            raise FileNotFoundError(f"Missing vocab.json and no embedded vocab in {tokenizer_json}")
+    if not vocab_path.exists() or not merges_path.exists():
+        raise FileNotFoundError(f"Missing vocab.json or merges.txt in {tokenizer_dir}")
+
+    with vocab_path.open("r", encoding="utf-8") as f:
+        vocab = json.load(f)
 
     max_id = max(vocab.values())
     if max_id > 65_535:
@@ -325,4 +345,4 @@ def load_tokenizer(tokenizer_dir: Path) -> Any:
 
     validate_tokenizer(tokenizer_dir)
     tokenizer_json, _, _ = tokenizer_files(tokenizer_dir)
-    return FusionTokenizer(Tokenizer.from_file(str(tokenizer_json)))
+    return AtomTokenizer(Tokenizer.from_file(str(tokenizer_json)))