| """Execute the authors' construction notebooks without modifying their bytes. |
| |
| The arXiv source links to https://github.com/SprocketLab/hybrid-expressivity. |
| ``source/official-code`` is a detached checkout of the last repository commit |
| before arXiv:2603.08859v1. This runner reads the code cells directly from the |
| two authored notebooks, executes their setup/weight-assignment cells unchanged |
| on CPU, and evaluates the constructed models on deterministic batches generated |
| by the authors' own Dataset implementation. Notebook display/plot cells are |
| intentionally not executed because they are not part of the forward pass. |
| |
| This is construction validation only. It is not an attempted rerun of the |
| GPU-only learned-model training entry point, which hard-codes ``cuda``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import copy |
| import hashlib |
| import json |
| import os |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| CODE = ROOT / "source" / "official-code" |
| NOTEBOOKS = { |
| "claim3_native": (CODE / "constructions" / "construction_var_copy.ipynb", 18), |
| "claim4_native": (CODE / "constructions" / "construction_decode_recall.ipynb", 27), |
| } |
| OUT = ROOT / "outputs" |
|
|
|
|
| def sha256(path: Path) -> str: |
| return hashlib.sha256(path.read_bytes()).hexdigest() |
|
|
|
|
| def execute_notebook_prefix(path: Path, last_cell: int) -> dict[str, Any]: |
| """Execute exact code-cell text through the construction's model cell.""" |
| notebook = json.loads(path.read_text()) |
| cells = notebook["cells"] |
| assert cells[0]["cell_type"] == "code" |
| |
| |
| selected = [i for i in range(1, last_cell + 1) if cells[i]["cell_type"] == "code"] |
| namespace: dict[str, Any] = {"__name__": "__main__", "__file__": str(path)} |
| old_cwd = Path.cwd() |
| old_path = list(sys.path) |
| try: |
| os.chdir(path.parent) |
| sys.path.insert(0, str(path.parent)) |
| np.random.seed(1729) |
| torch.manual_seed(1729) |
| for i in selected: |
| source = "".join(cells[i]["source"]) |
| exec(compile(source, f"{path}::cell-{i}", "exec"), namespace) |
| finally: |
| os.chdir(old_cwd) |
| sys.path[:] = old_path |
| namespace["_executed_cells"] = selected |
| return namespace |
|
|
|
|
| def evaluate_last_position(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]: |
| """Use the already-authored Dataset/model objects and score final tokens.""" |
| model = ns["model"] |
| dataset = ns["train_dataset"] |
| model.eval() |
| eligible = correct = 0 |
| with torch.no_grad(): |
| for i in range(n_batches): |
| batch = dataset[i] |
| logits = model(batch["input_ids"]) |
| pred = torch.argmax(logits[:, -1], dim=-1) |
| valid = batch["mask"][:, -1].bool() |
| eligible += int(valid.sum()) |
| correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum()) |
| return { |
| "n_batches": n_batches, |
| "examples_per_batch": int(ns["args"].train_batch_size), |
| "eligible_last_position_examples": eligible, |
| "correct_last_position_examples": correct, |
| "last_position_accuracy": correct / eligible if eligible else None, |
| } |
|
|
|
|
| def destructive_no_first_ssm(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]: |
| """Ablate the first authored SSM gate in-memory; source bytes stay intact.""" |
| model = copy.deepcopy(ns["model"]) |
| model.layers[0].layer.Delta.data.zero_() |
| dataset = ns["train_dataset"] |
| model.eval() |
| eligible = correct = 0 |
| |
| np.random.seed(1729 + 1) |
| with torch.no_grad(): |
| for i in range(n_batches): |
| batch = dataset[i] |
| logits = model(batch["input_ids"]) |
| pred = torch.argmax(logits[:, -1], dim=-1) |
| valid = batch["mask"][:, -1].bool() |
| eligible += int(valid.sum()) |
| correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum()) |
| return { |
| "ablation": "first SimpleSSM Delta vector set to zero in memory", |
| "eligible_last_position_examples": eligible, |
| "correct_last_position_examples": correct, |
| "last_position_accuracy": correct / eligible if eligible else None, |
| } |
|
|
|
|
| def main() -> None: |
| OUT.mkdir(exist_ok=True) |
| git_head = (CODE / ".git").exists() |
| for output_name, (notebook, last_cell) in NOTEBOOKS.items(): |
| |
| ns = execute_notebook_prefix(notebook, last_cell) |
| np.random.seed(1729 + 1) |
| baseline = evaluate_last_position(ns) |
| control = destructive_no_first_ssm(ns) |
| payload = { |
| "kind": "direct_native_notebook_execution", |
| "official_repository": "https://github.com/SprocketLab/hybrid-expressivity", |
| "repository_checkout": "6be8f8fbc2169290af6f4ba5e4bd53a5c6485f7b", |
| "repository_is_git_checkout": git_head, |
| "notebook": str(notebook.relative_to(ROOT)), |
| "notebook_sha256": sha256(notebook), |
| "executed_code_cells": ns["_executed_cells"], |
| "device": "cpu", |
| "torch_version": torch.__version__, |
| "baseline": baseline, |
| "destructive_control": control, |
| "control_is_lower": ( |
| control["last_position_accuracy"] < baseline["last_position_accuracy"] |
| if baseline["last_position_accuracy"] is not None |
| else None |
| ), |
| "scope": ( |
| "Construction cells plus deterministic last-position batches; " |
| "not a learned-model training rerun." |
| ), |
| } |
| (OUT / f"{output_name}.json").write_text(json.dumps(payload, indent=2) + "\n") |
| print( |
| f"{output_name}: baseline={baseline['last_position_accuracy']:.6f} " |
| f"control={control['last_position_accuracy']:.6f} n={baseline['eligible_last_position_examples']}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
