Add paired GRASP/A1 Colab comparison with shared audited splits

README.md

@@ -15,7 +15,7 @@ tags:
 
 **GRASP-A1** is the GRASP molecular graph Transformer trained without absolute atom-position embeddings. It retains shortest-path relative attention, radius-0 atom-environment tokens, and the replaced-token-detection objective. The encoder has **93,091,584 parameters**, 12 layers, hidden size 768, and 12 attention heads.
 
-This private research release provides the completed **50,000-update Step2 encoder** by default and the corresponding **Step1 encoder** under `step1/`. Step1 saw **1,538,242,560 molecule presentations** from ZINC20; this is a presentation budget, not a count of certified unique molecules. Step2 used the audited 511,897-molecule ChEMBL 36 set with 642 assay tasks.
+This research release provides the completed **50,000-update Step2 encoder** by default and the corresponding **Step1 encoder** under `step1/`. Step1 saw **1,538,242,560 molecule presentations** from ZINC20; this is a presentation budget, not a count of certified unique molecules. Step2 used the audited 511,897-molecule ChEMBL 36 set with 642 assay tasks.
 
 ## Use
 
@@ -23,10 +23,9 @@ Install Python 3.10–3.13 and the [GRASP code repository](https://github.com/ca
 
 ```bash
 python -m pip install "git+https://github.com/caithmac/GRASP.git@7843e72606dd208517cd8cf86ee07b3bb965cc79"
-hf auth login
 ```
 
-Authenticate with an account permitted to read this private repository, then:
+The repository is public and requires no authentication:
 
 ```python
 from grasp import GRASPEncoder
@@ -66,6 +65,12 @@ python -m grasp predict --model runs/a1_my_task \
 
 Use `--task binary` for binary classification. `--method frozen` trains the readout and head with a fixed encoder; `--method lora` provides LoRA adaptation. Use a separate output directory per endpoint. To compare against published GRASP, repeat the same split, preprocessing, method and seed with `--model caithmac/GRASP`.
 
+## Compare A1 and original GRASP in Colab
+
+[Open the comparison notebook in Colab](https://colab.research.google.com/drive/15G8jCYj2RmRiJLUtjoZqi_VHvzw89GIF), or [download its self-contained source](notebooks/01_compare_GRASP_A1.ipynb) and use **Colab → File → Upload notebook**. Select a GPU and run cells in order. Public access requires no token; if access is restricted later, the notebook supports the Colab secret `HF_TOKEN` or a hidden login prompt.
+
+The notebook tests ESOL, FreeSolv and BBBP with shared scaffold splits, matched initial heads and training settings, and pinned checkpoint revisions. It exports paired metrics, per-molecule predictions, split hashes and data-exclusion audits. Defaults use one training seed; set `SEEDS = [0, 1, 2]` for repeated training on the same splits. This compares the released checkpoints and does not isolate the positional-embedding change or reproduce paper scores. See [notebook validation](notebooks/verification.json) for the real-data split audits and tiny CPU workflow checks; full Colab GPU benchmarks have not been run as part of notebook creation.
+
 ## Model details and provenance
 
 | Property | Value |

SHA256SUMS

@@ -1,8 +1,11 @@
 5b2636c65d69d21303a60eb7f2e9dad76b6b1ba78e27202079184a27b5789b62  config.json
 6318960741814b0ae14a6769b0e129451ae7f3bec6ecbf0764fbb519407bc6a4  LICENSE
 7f355e252de37e3a4affbb3ca2aabe79a92c9e893bbee01c9f0fd171b3e41e92  model.safetensors
+cc71af7c0c1572a0718cbd5276e2a76cc769ca72fb46f02447cfd383d28fd6dd  notebooks/01_compare_GRASP_A1.ipynb
+5bff4850d4b7da22944c50ac1a741c6b1ca8738347ec5a55ec3bc9140d3aaaae  notebooks/comparison_helpers.py
+b92ff6a337427a514b3e7c2f6334617bcb864f976e52216c88c54f6335b9fa7e  notebooks/verification.json
 ac8a653d18e0fdfad95a2edf74b2c813d820d86bcf03527d88029bef8b04cd21  NOTICE
-53dae9e0d17f514ee38eae7f605715285a2fea1f4e644cdfcb638f4e7f86bd6f  README.md
+55ace9b45470a55e33a332e6c6d4625979199e8a760bf8882dc659c488b6570b  README.md
 adcb0cb97d071880ee1e540d20434c8138dfdd63f9db654f5eef4c8584d38108  requirements-tested.txt
 784e6a4c9f2e22dbbc7f02d3bd151fe6de96e065aa50c3df21f238e9addb2948  step1/config.json
 e1fbbfb34f65b3a30072f2df482a600024b4749d131e542b0d0a3c23604334f9  step1/model.safetensors

notebooks/01_compare_GRASP_A1.ipynb

@@ -0,0 +1,585 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Compare A1 with the released GRASP on the same splits\n",
+    "\n",
+    "Three small experimental tasks: **ESOL** (solubility regression), **FreeSolv** (hydration-free-energy regression), and **BBBP** (binary classification). Choose a GPU under **Runtime \u2192 Change runtime type**, then run cells in order. If A1 requires private access, use a Hugging Face read token permitted to read `caithmac/GRASP-A1`.\n",
+    "\n",
+    "Both models receive the **same saved scaffold train/validation/test CSVs**, canonical SMILES, training seed, initial task head, optimizer settings and learned layer-mixture/atom-attention readout. Test labels are used only for evaluation after training. Defaults perform one-seed full fine-tuning; change `SEEDS` to `[0, 1, 2]` for repeated training on the same splits, or select `frozen` for a faster screen.\n",
+    "\n",
+    "**Interpretation:** this compares the two released Step2 checkpoints. The public GRASP and A1 have additional training-history differences, so the result cannot isolate removal of absolute positions. It does not reproduce paper splits/scores, certify absence of pretraining overlap, or establish a general improvement. Neither encoder represents stereochemical assignments.\n"
+   ],
+   "id": "a1-comparison-00"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Install the pinned GRASP code\n",
+    "\n",
+    "Use a fresh Colab session. The existing CUDA PyTorch build is retained when it satisfies GRASP's requirements. If installation replaces an already imported scientific package, restart the session and rerun from here. No DeepChem installation is required.\n"
+   ],
+   "id": "a1-comparison-01"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import os, sys, subprocess, importlib.metadata\n",
+    "os.environ.setdefault(\"CUBLAS_WORKSPACE_CONFIG\", \":4096:8\")\n",
+    "CODE_REVISION = \"7843e72606dd208517cd8cf86ee07b3bb965cc79\"\n",
+    "watched = {\"numpy\": \"numpy\", \"scipy\": \"scipy\", \"rdkit\": \"rdkit\", \"sklearn\": \"scikit-learn\"}\n",
+    "loaded = {module: importlib.metadata.version(package) for module, package in watched.items() if module in sys.modules}\n",
+    "subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\",\n",
+    "    f\"git+https://github.com/caithmac/GRASP.git@{CODE_REVISION}\",\n",
+    "    \"rdkit==2025.3.6\", \"scipy==1.15.3\", \"scikit-learn==1.7.2\"], check=True)\n",
+    "changed = [module for module, version in loaded.items() if importlib.metadata.version(watched[module]) != version]\n",
+    "if changed:\n",
+    "    raise RuntimeError(f\"Restart the Colab session and rerun from installation; imported packages changed: {changed}\")\n",
+    "print(\"GRASP installation ready\")\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-02"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Settings\n",
+    "\n",
+    "Set these **before** looking at test results. Use a new `RUN_ROOT` when changing settings. Completed fits can be skipped safely when rerunning the same experiment; interrupted fits restart that individual fit. Each saved predictor includes its encoder and can occupy about 370 MB. For persistence, set `RUN_ROOT` to a directory on already mounted Drive storage. The final compact ZIP excludes predictor weights and raw datasets.\n"
+   ],
+   "id": "a1-comparison-03"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from pathlib import Path\n",
+    "TASKS = [\"esol\", \"freesolv\", \"bbbp\"]\n",
+    "METHOD = \"full\"                 # \"full\", \"lora\", or \"frozen\"; identical for both models\n",
+    "SEEDS = [0]                     # training seeds; use [0, 1, 2] for repeats\n",
+    "SPLIT_SEED = 20261015            # fixed independently of training seeds\n",
+    "EPOCHS, PATIENCE, BATCH_SIZE = 25, 6, 16\n",
+    "LEARNING_RATE = {\"full\": 1e-5, \"lora\": 5e-5, \"frozen\": 1e-3}[METHOD]\n",
+    "RUN_ROOT = Path(\"/content/grasp_a1_comparison\")\n",
+    "MODELS = {\n",
+    "    \"GRASP\": {\"repo\": \"caithmac/GRASP\", \"revision\": \"4804c0a558307c944da5e8de1d182b4b4c994dae\"},\n",
+    "    \"A1\": {\"repo\": \"caithmac/GRASP-A1\", \"revision\": \"d8f631e096f1d2dfd281962a21389142cde437ee\"},\n",
+    "}\n",
+    "assert SEEDS and len(SEEDS) == len(set(SEEDS))\n",
+    "assert METHOD in (\"full\", \"lora\", \"frozen\")\n",
+    "RUN_ROOT.mkdir(parents=True, exist_ok=True)\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-04"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Confirm checkpoint access\n",
+    "\n",
+    "Public checkpoints or an already authenticated session need no extra login. For private A1 access, add `HF_TOKEN` in Colab's **Secrets** panel and allow this notebook access, or enter a read token in the hidden prompt. The token is not embedded in the notebook or exported results. Do not share an executed notebook containing credentials you added yourself.\n"
+   ],
+   "id": "a1-comparison-05"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from getpass import getpass\n",
+    "from huggingface_hub import login, HfApi, snapshot_download\n",
+    "from huggingface_hub.errors import HfHubHTTPError\n",
+    "try:\n",
+    "    HfApi().model_info(MODELS[\"A1\"][\"repo\"], revision=MODELS[\"A1\"][\"revision\"])\n",
+    "except HfHubHTTPError as exc:\n",
+    "    if exc.response is None or exc.response.status_code not in (401, 403, 404):\n",
+    "        raise\n",
+    "    token = None\n",
+    "    try:\n",
+    "        from google.colab import userdata\n",
+    "        token = userdata.get(\"HF_TOKEN\")\n",
+    "    except Exception:\n",
+    "        pass\n",
+    "    if not token:\n",
+    "        token = getpass(\"Hugging Face read token with access to private A1: \")\n",
+    "    login(token=token, add_to_git_credential=False)\n",
+    "    del token\n",
+    "    HfApi().model_info(MODELS[\"A1\"][\"repo\"], revision=MODELS[\"A1\"][\"revision\"])\n",
+    "print(\"A1 checkpoint access confirmed\")\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-06"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Load the shared protocol and pin model files\n",
+    "\n",
+    "Encoder construction is isolated from the task-head RNG stream. Otherwise the different position-table sizes would give different randomly initialized heads even with the same seed. The notebook checks both the initial readout hash and the post-initialization RNG state.\n"
+   ],
+   "id": "a1-comparison-07"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "#@title Shared protocol helpers (run once)\n",
+    "\"\"\"Shared, auditable protocol embedded in the self-contained Colab notebook.\"\"\"\n",
+    "import gc\n",
+    "import hashlib\n",
+    "import io\n",
+    "import json\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import requests\n",
+    "import torch\n",
+    "from rdkit import Chem\n",
+    "from rdkit.Chem.Scaffolds import MurckoScaffold\n",
+    "from sklearn.metrics import average_precision_score, mean_absolute_error, mean_squared_error, r2_score, roc_auc_score\n",
+    "\n",
+    "DATASETS = {\n",
+    "    \"esol\": dict(filename=\"delaney-processed.csv\", target=\"measured log solubility in mols per litre\", task=\"regression\",\n",
+    "                 sha256=\"8c06a76f0c6487d29ab0f903e6a7a7139f189ab3c1178f159c8be8964602f189\"),\n",
+    "    \"freesolv\": dict(filename=\"freesolv.csv.gz\", target=\"y\", task=\"regression\",\n",
+    "                     sha256=\"008e9a8a9e018c6b77eb00df7153d0bf2e35a1d663f18f301de522fe3f8711d7\"),\n",
+    "    \"bbbp\": dict(filename=\"BBBP.csv\", target=\"p_np\", task=\"binary\",\n",
+    "                 sha256=\"d07a38487aeac5cee5508413e468043ef3097451d2a112701c2d60be9ec6b662\"),\n",
+    "}\n",
+    "PROTOCOL_VERSION = \"grasp-a1-paired-colab-20261015-v1\"\n",
+    "\n",
+    "\n",
+    "def sha256(path):\n",
+    "    digest = hashlib.sha256()\n",
+    "    with Path(path).open(\"rb\") as stream:\n",
+    "        for block in iter(lambda: stream.read(1024 * 1024), b\"\"):\n",
+    "            digest.update(block)\n",
+    "    return digest.hexdigest()\n",
+    "\n",
+    "\n",
+    "def write_json(path, value):\n",
+    "    Path(path).parent.mkdir(parents=True, exist_ok=True)\n",
+    "    Path(path).write_text(json.dumps(value, indent=2, allow_nan=False) + \"\\n\")\n",
+    "\n",
+    "\n",
+    "def prepare_task(name, cache_dir, output_dir, vocabulary):\n",
+    "    from mole.training.data.datasets import MolDataset\n",
+    "    spec = DATASETS[name]\n",
+    "    cache_dir, output_dir = Path(cache_dir), Path(output_dir)\n",
+    "    cache_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    raw = cache_dir / spec[\"filename\"]\n",
+    "    url = \"https://deepchemdata.s3.us-west-1.amazonaws.com/datasets/\" + spec[\"filename\"]\n",
+    "    if not raw.is_file():\n",
+    "        response = requests.get(url, timeout=90)\n",
+    "        response.raise_for_status()\n",
+    "        raw.write_bytes(response.content)\n",
+    "    assert sha256(raw) == spec[\"sha256\"], f\"Source changed for {name}; review it before updating its checksum.\"\n",
+    "    frame = pd.read_csv(io.BytesIO(raw.read_bytes()), compression=\"gzip\" if raw.suffix == \".gz\" else None)\n",
+    "    accepted, audit = [], []\n",
+    "    for source_row, row in frame.iterrows():\n",
+    "        reason = \"\"\n",
+    "        try:\n",
+    "            target = float(row[spec[\"target\"]])\n",
+    "            if not np.isfinite(target):\n",
+    "                raise ValueError(\"nonfinite_target\")\n",
+    "            if spec[\"task\"] == \"binary\" and target not in (0, 1):\n",
+    "                raise ValueError(\"nonbinary_target\")\n",
+    "            mol = Chem.MolFromSmiles(str(row[\"smiles\"]))\n",
+    "            if mol is None or not mol.GetNumAtoms():\n",
+    "                raise ValueError(\"unparsable_or_empty\")\n",
+    "            if mol.GetNumAtoms() > 511:\n",
+    "                raise ValueError(\"above_loader_atom_limit\")\n",
+    "            smiles = Chem.MolToSmiles(mol, isomericSmiles=True)\n",
+    "            # Check the exact deployed featurizer once, before either model receives a split.\n",
+    "            MolDataset(pd.Series([smiles]), vocabulary, cls_token=True)[0]\n",
+    "            connectivity = Chem.MolToSmiles(mol, isomericSmiles=False)\n",
+    "            scaffold = Chem.MolToSmiles(MurckoScaffold.GetScaffoldForMol(mol), isomericSmiles=False) or \"<acyclic>\"\n",
+    "            accepted.append(dict(source_row=int(source_row), smiles=smiles, target=target,\n",
+    "                                 connectivity=connectivity, scaffold=scaffold))\n",
+    "        except Exception as exc:\n",
+    "            reason = f\"{type(exc).__name__}: {exc}\"\n",
+    "        audit.append(dict(source_row=int(source_row), original_smiles=str(row[\"smiles\"]), excluded=bool(reason), reason=reason))\n",
+    "    pd.DataFrame(audit).to_csv(output_dir / \"data_audit.csv\", index=False)\n",
+    "    result = pd.DataFrame(accepted)\n",
+    "    assert len(result), f\"No usable rows for {name}\"\n",
+    "    result.to_csv(output_dir / \"eligible_rows.csv\", index=False)\n",
+    "    write_json(output_dir / \"data_manifest.json\", dict(dataset=name, url=url, source_sha256=spec[\"sha256\"],\n",
+    "               source_rows=len(frame), eligible_rows=len(result), excluded_rows=len(frame) - len(result),\n",
+    "               policy=\"Canonical isomeric SMILES; no charge/salt/tautomer repair; retain duplicates within one scaffold partition.\"))\n",
+    "    return result\n",
+    "\n",
+    "\n",
+    "def scaffold_split(frame, seed):\n",
+    "    groups = [group.index.to_list() for _, group in frame.groupby(\"scaffold\", sort=True)]\n",
+    "    np.random.default_rng(seed).shuffle(groups)\n",
+    "    groups.sort(key=len, reverse=True)\n",
+    "    partitions = {\"train\": [], \"valid\": [], \"test\": []}\n",
+    "    for group in groups:\n",
+    "        if len(partitions[\"train\"]) + len(group) <= 0.8 * len(frame):\n",
+    "            partitions[\"train\"].extend(group)\n",
+    "        elif len(partitions[\"valid\"]) + len(group) <= 0.1 * len(frame):\n",
+    "            partitions[\"valid\"].extend(group)\n",
+    "        else:\n",
+    "            partitions[\"test\"].extend(group)\n",
+    "    result = {key: frame.loc[sorted(rows)].reset_index(drop=True) for key, rows in partitions.items()}\n",
+    "    assert sum(map(len, result.values())) == len(frame)\n",
+    "    assert all(len(part) for part in result.values()), \"Too few scaffold groups to create three nonempty splits.\"\n",
+    "    for column in (\"source_row\", \"connectivity\", \"scaffold\"):\n",
+    "        sets = [set(part[column]) for part in result.values()]\n",
+    "        assert not (sets[0] & sets[1] or sets[0] & sets[2] or sets[1] & sets[2]), f\"Split leakage: {column}\"\n",
+    "    return result\n",
+    "\n",
+    "\n",
+    "def save_splits(frame, task, directory, seed):\n",
+    "    directory = Path(directory)\n",
+    "    directory.mkdir(parents=True, exist_ok=True)\n",
+    "    parts = scaffold_split(frame, seed)\n",
+    "    if task == \"binary\":\n",
+    "        assert all(set(part.target) == {0, 1} for part in parts.values()), \"A binary partition lacks one class; do not silently resplit.\"\n",
+    "    csv_bytes = {split: part.to_csv(index=False).encode(\"utf-8\") for split, part in parts.items()}\n",
+    "    for split, data in csv_bytes.items():\n",
+    "        path = directory / f\"{split}.csv\"\n",
+    "        assert not path.is_file() or path.read_bytes() == data, \"Existing splits differ: choose a new RUN_ROOT.\"\n",
+    "    for split, data in csv_bytes.items():\n",
+    "        (directory / f\"{split}.csv\").write_bytes(data)\n",
+    "    manifest = dict(split_seed=seed, policy=\"Shared scaffold-grouped 80/10/10 target proportions; all acyclic molecules form one group.\",\n",
+    "                    counts={key: len(part) for key, part in parts.items()},\n",
+    "                    sha256={key: sha256(directory / f\"{key}.csv\") for key in parts},\n",
+    "                    disjoint_source_rows=True, disjoint_connectivity=True, disjoint_scaffolds=True)\n",
+    "    write_json(directory / \"split_manifest.json\", manifest)\n",
+    "    return manifest\n",
+    "\n",
+    "\n",
+    "def enable_matched_initialization():\n",
+    "    import grasp.encoder as encoder_module\n",
+    "    import grasp.predictor as predictor_module\n",
+    "    from mole.training.models.mole import AtomEnvEmbeddings\n",
+    "    class MatchedEmbeddings(AtomEnvEmbeddings):\n",
+    "        def __init__(self, config):\n",
+    "            # Loading a different encoder architecture must not shift task-head or loader RNG.\n",
+    "            with torch.random.fork_rng(devices=[]):\n",
+    "                super().__init__(config)\n",
+    "    encoder_module.AtomEnvEmbeddings = predictor_module.AtomEnvEmbeddings = MatchedEmbeddings\n",
+    "\n",
+    "\n",
+    "def head_fingerprint(model_dir, seed, method):\n",
+    "    from grasp import GRASPEncoder, GRASPPredictor\n",
+    "    torch.manual_seed(seed)\n",
+    "    encoder = GRASPEncoder.from_pretrained(model_dir, device=\"cpu\")\n",
+    "    predictor = GRASPPredictor.from_encoder(encoder, task=\"regression\", method=method)\n",
+    "    digest = hashlib.sha256()\n",
+    "    for key, value in sorted(predictor.model.state_dict().items()):\n",
+    "        if not key.startswith(\"encoder.\"):\n",
+    "            digest.update(key.encode())\n",
+    "            digest.update(value.detach().cpu().contiguous().numpy().tobytes())\n",
+    "    state = hashlib.sha256(torch.get_rng_state().numpy().tobytes()).hexdigest()\n",
+    "    del predictor, encoder\n",
+    "    gc.collect()\n",
+    "    return dict(readout_sha256=digest.hexdigest(), post_initialization_rng_sha256=state)\n",
+    "\n",
+    "\n",
+    "def train_one(model_dir, split_dir, output_dir, *, task, method, seed, epochs, patience, batch_size, learning_rate, device, identity):\n",
+    "    from grasp.train import fit\n",
+    "    split_dir, output_dir = Path(split_dir), Path(output_dir)\n",
+    "    split_hashes = {key: sha256(split_dir / f\"{key}.csv\") for key in (\"train\", \"valid\", \"test\")}\n",
+    "    assert split_hashes == json.loads((split_dir / \"split_manifest.json\").read_text())[\"sha256\"], \"Shared split files changed.\"\n",
+    "    plan = dict(protocol=PROTOCOL_VERSION, identity=identity, split_sha256=split_hashes, task=task, method=method,\n",
+    "                seed=seed, epochs=epochs, patience=patience, batch_size=batch_size, learning_rate=learning_rate, device=device)\n",
+    "    plan_file = output_dir / \"run_plan.json\"\n",
+    "    if plan_file.is_file():\n",
+    "        assert json.loads(plan_file.read_text()) == plan, \"Settings/splits changed: choose a new RUN_ROOT.\"\n",
+    "    else:\n",
+    "        write_json(plan_file, plan)\n",
+    "    complete = output_dir / \"TRAINING_COMPLETE.json\"\n",
+    "    if complete.is_file():\n",
+    "        assert (output_dir / \"predictor.safetensors\").is_file()\n",
+    "        return json.loads(complete.read_text())\n",
+    "    # Test labels are never supplied to fitting or early stopping.\n",
+    "    outcome = fit(model_source=model_dir, train_csv=split_dir / \"train.csv\", valid_csv=split_dir / \"valid.csv\",\n",
+    "                  output_dir=output_dir, task=task, method=method, seed=seed, epochs=epochs, patience=patience,\n",
+    "                  batch_size=batch_size, learning_rate=learning_rate, device=device)\n",
+    "    write_json(complete, outcome)\n",
+    "    gc.collect()\n",
+    "    if torch.cuda.is_available():\n",
+    "        torch.cuda.empty_cache()\n",
+    "    return outcome\n",
+    "\n",
+    "\n",
+    "def score_predictions(target, prediction, task):\n",
+    "    target, prediction = np.asarray(target), np.asarray(prediction)\n",
+    "    assert target.shape == prediction.shape and np.isfinite(prediction).all()\n",
+    "    if task == \"regression\":\n",
+    "        return dict(mae=float(mean_absolute_error(target, prediction)), rmse=float(np.sqrt(mean_squared_error(target, prediction))),\n",
+    "                    r2=float(r2_score(target, prediction)))\n",
+    "    assert set(target) == {0, 1} and ((prediction >= 0) & (prediction <= 1)).all()\n",
+    "    return dict(auroc=float(roc_auc_score(target, prediction)), auprc=float(average_precision_score(target, prediction)))\n",
+    "\n",
+    "\n",
+    "def paired_table(results):\n",
+    "    pairs = []\n",
+    "    for (dataset, seed), group in results.groupby([\"dataset\", \"seed\"]):\n",
+    "        assert set(group.model) == {\"GRASP\", \"A1\"} and len(group) == 2\n",
+    "        group = group.set_index(\"model\")\n",
+    "        metric = \"mae\" if group.loc[\"A1\", \"task\"] == \"regression\" else \"auroc\"\n",
+    "        baseline, a1 = float(group.loc[\"GRASP\", metric]), float(group.loc[\"A1\", metric])\n",
+    "        delta = a1 - baseline\n",
+    "        pairs.append(dict(dataset=dataset, seed=int(seed), metric=metric, GRASP=baseline, A1=a1,\n",
+    "                          delta_A1_minus_GRASP=delta, direction=\"lower is better\" if metric == \"mae\" else \"higher is better\",\n",
+    "                          A1_better=bool(delta < 0 if metric == \"mae\" else delta > 0)))\n",
+    "    return pd.DataFrame(pairs)\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-08"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import importlib.metadata, platform\n",
+    "from rdkit import RDLogger\n",
+    "from grasp.encoder import read_vocab\n",
+    "from IPython.display import display\n",
+    "RDLogger.DisableLog(\"rdApp.warning\")\n",
+    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "torch.set_num_threads(2)\n",
+    "enable_matched_initialization()\n",
+    "MODEL_DIRS, registry = {}, {}\n",
+    "for name, model in MODELS.items():\n",
+    "    folder = Path(snapshot_download(model[\"repo\"], revision=model[\"revision\"],\n",
+    "                 allow_patterns=[\"model.safetensors\", \"config.json\", \"vocab.json\"]))\n",
+    "    MODEL_DIRS[name] = folder\n",
+    "    config = json.loads((folder / \"config.json\").read_text())\n",
+    "    assert config[\"architecture\"][\"position_biased_input\"] == (name == \"GRASP\")\n",
+    "    registry[name] = {**model, \"files_sha256\": {file: sha256(folder / file)\n",
+    "                     for file in (\"model.safetensors\", \"config.json\", \"vocab.json\")}}\n",
+    "vocabulary = read_vocab(MODEL_DIRS[\"GRASP\"] / \"vocab.json\")\n",
+    "assert vocabulary == read_vocab(MODEL_DIRS[\"A1\"] / \"vocab.json\")\n",
+    "runtime = {\"python\": platform.python_version(), \"device\": DEVICE,\n",
+    "           \"gpu\": torch.cuda.get_device_name(0) if DEVICE == \"cuda\" else None,\n",
+    "           \"bf16\": DEVICE == \"cuda\" and torch.cuda.is_bf16_supported(),\n",
+    "           \"packages\": {p: importlib.metadata.version(p) for p in\n",
+    "               (\"torch\", \"rdkit\", \"numpy\", \"scipy\", \"scikit-learn\", \"grasp-molecular\")}}\n",
+    "print(runtime)\n",
+    "if DEVICE == \"cpu\":\n",
+    "    print(\"CPU selected: full fine-tuning can be slow; a Colab GPU is recommended.\")\n",
+    "write_json(RUN_ROOT / \"models.json\", registry)\n",
+    "write_json(RUN_ROOT / \"runtime.json\", runtime)\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-09"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Prepare data and create splits once\n",
+    "\n",
+    "Sources and checksums are fixed. Parsing, missing labels, unsupported size and featurizer failures are audited jointly **before** splitting; excluded rows are listed in `data_audit.csv`. No molecule is repaired or truncated. Duplicates are retained within one partition. Canonical isomeric SMILES are shared by both models; salts, charges and tautomers are not standardized away.\n",
+    "\n",
+    "The 80/10/10 **target** proportions use scaffold groups; all acyclic structures stay together. Actual counts can differ from these proportions. The notebook checks row, connectivity and scaffold disjointness. These are newly generated comparison splits, not an official test-set modification or a reproduction of DeepChem's default split for every dataset.\n"
+   ],
+   "id": "a1-comparison-10"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "assert set(TASKS) <= set(DATASETS) and len(TASKS) == len(set(TASKS))\n",
+    "split_manifests, overview = {}, []\n",
+    "for dataset in TASKS:\n",
+    "    directory = RUN_ROOT / \"splits\" / dataset\n",
+    "    frame = prepare_task(dataset, RUN_ROOT / \"raw\", directory, vocabulary)\n",
+    "    split_manifests[dataset] = save_splits(frame, DATASETS[dataset][\"task\"], directory, SPLIT_SEED)\n",
+    "    audit = json.loads((directory / \"data_manifest.json\").read_text())\n",
+    "    overview.append({\"dataset\": dataset, \"task\": DATASETS[dataset][\"task\"],\n",
+    "                     \"excluded\": audit[\"excluded_rows\"], **split_manifests[dataset][\"counts\"]})\n",
+    "display(pd.DataFrame(overview))\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-11"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "head_checks = {}\n",
+    "for seed in SEEDS:\n",
+    "    checks = {name: head_fingerprint(folder, seed, METHOD) for name, folder in MODEL_DIRS.items()}\n",
+    "    assert checks[\"GRASP\"] == checks[\"A1\"], \"Initial head/RNG mismatch: comparison blocked.\"\n",
+    "    head_checks[str(seed)] = checks\n",
+    "write_json(RUN_ROOT / \"head_initialization_checks.json\", head_checks)\n",
+    "experiment = dict(protocol=PROTOCOL_VERSION, code_revision=CODE_REVISION, models=registry,\n",
+    "                  tasks=TASKS, method=METHOD, seeds=SEEDS, split_seed=SPLIT_SEED,\n",
+    "                  epochs=EPOCHS, patience=PATIENCE, batch_size=BATCH_SIZE, learning_rate=LEARNING_RATE,\n",
+    "                  splits=split_manifests, runtime=runtime)\n",
+    "experiment_file = RUN_ROOT / \"experiment.json\"\n",
+    "if experiment_file.is_file():\n",
+    "    assert json.loads(experiment_file.read_text()) == experiment, \"Experiment changed: use a new RUN_ROOT.\"\n",
+    "else:\n",
+    "    write_json(experiment_file, experiment)\n",
+    "print(\"Same splits, checkpoint files, initial heads and post-initialization RNG verified.\")\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-12"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Train both models on the same files\n",
+    "\n",
+    "The learned readout mixes layers 4/8/10/12 and uses atom-attention pooling. Target scaling comes only from training labels. Validation controls early stopping separately for each model under the same rule; **no test predictions are used here**. The fitting runner saves a predictor per endpoint, model and seed.\n"
+   ],
+   "id": "a1-comparison-13"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "training_rows = []\n",
+    "for dataset in TASKS:\n",
+    "    split_dir = RUN_ROOT / \"splits\" / dataset\n",
+    "    for seed in SEEDS:\n",
+    "        for model_name in (\"GRASP\", \"A1\"):\n",
+    "            output = RUN_ROOT / \"runs\" / dataset / f\"seed_{seed}\" / model_name\n",
+    "            print(f\"TRAIN {dataset} seed={seed} {model_name}\", flush=True)\n",
+    "            outcome = train_one(MODEL_DIRS[model_name], split_dir, output,\n",
+    "                task=DATASETS[dataset][\"task\"], method=METHOD, seed=seed,\n",
+    "                epochs=EPOCHS, patience=PATIENCE, batch_size=BATCH_SIZE,\n",
+    "                learning_rate=LEARNING_RATE, device=DEVICE,\n",
+    "                identity={\"checkpoint\": registry[model_name], \"runtime\": runtime,\n",
+    "                          \"initialization\": head_checks[str(seed)][model_name]})\n",
+    "            training_rows.append({\"dataset\": dataset, \"seed\": seed, \"model\": model_name,\n",
+    "                                  \"best_epoch\": outcome[\"best_epoch\"], \"validation_score\": outcome[\"validation_score\"]})\n",
+    "pd.DataFrame(training_rows).to_csv(RUN_ROOT / \"training.csv\", index=False)\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-14"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Evaluate and compare\n",
+    "\n",
+    "MAE/RMSE are in the target's original units; lower is better. BBBP AUROC/AUPRC use predicted class-1 probabilities; higher is better. Each paired delta is **A1 minus GRASP**. Do not average incompatible metrics across tasks. With one seed, sample standard deviation is undefined. Repeated training seeds use the same held-out structures and do not quantify split or pretraining uncertainty.\n"
+   ],
+   "id": "a1-comparison-15"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from grasp import GRASPPredictor\n",
+    "results, predictions = [], []\n",
+    "for dataset in TASKS:\n",
+    "    split_dir = RUN_ROOT / \"splits\" / dataset\n",
+    "    test = pd.read_csv(split_dir / \"test.csv\")\n",
+    "    assert sha256(split_dir / \"test.csv\") == split_manifests[dataset][\"sha256\"][\"test\"]\n",
+    "    for seed in SEEDS:\n",
+    "        for model_name in (\"GRASP\", \"A1\"):\n",
+    "            predictor = GRASPPredictor.from_pretrained(RUN_ROOT / \"runs\" / dataset / f\"seed_{seed}\" / model_name, device=DEVICE)\n",
+    "            pred = predictor.predict(test.smiles.tolist(), batch_size=BATCH_SIZE)\n",
+    "            metrics = score_predictions(test.target.to_numpy(), pred, DATASETS[dataset][\"task\"])\n",
+    "            results.append({\"dataset\": dataset, \"seed\": seed, \"model\": model_name,\n",
+    "                            \"task\": DATASETS[dataset][\"task\"], \"test_rows\": len(test), **metrics})\n",
+    "            block = test[[\"source_row\", \"smiles\", \"target\"]].copy()\n",
+    "            block[\"prediction\"], block[\"dataset\"], block[\"seed\"], block[\"model\"] = pred, dataset, seed, model_name\n",
+    "            predictions.append(block)\n",
+    "            del predictor\n",
+    "            gc.collect()\n",
+    "            if DEVICE == \"cuda\": torch.cuda.empty_cache()\n",
+    "results = pd.DataFrame(results)\n",
+    "paired = paired_table(results)\n",
+    "summary = paired.groupby([\"dataset\", \"metric\", \"direction\"])[[\"GRASP\", \"A1\", \"delta_A1_minus_GRASP\"]].agg([\"mean\", \"std\", \"count\"])\n",
+    "results.to_csv(RUN_ROOT / \"results.csv\", index=False)\n",
+    "paired.to_csv(RUN_ROOT / \"paired_comparison.csv\", index=False)\n",
+    "summary.to_csv(RUN_ROOT / \"summary.csv\")\n",
+    "pd.concat(predictions, ignore_index=True).to_csv(RUN_ROOT / \"test_predictions.csv\", index=False)\n",
+    "display(paired)\n",
+    "display(summary)\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-16"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Download the evidence\n",
+    "\n",
+    "The compact ZIP contains scores, predictions, shared splits, exclusions, runtime settings, checkpoint hashes and head checks. It excludes raw data and model/predictor weights. Keep the `runs/` directory separately if you need the trained predictors after Colab disconnects. Test data must remain untouched for an honest comparison; new experiments selected after viewing these scores need a new held-out evaluation.\n"
+   ],
+   "id": "a1-comparison-17"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import zipfile\n",
+    "archive = RUN_ROOT / \"grasp_a1_comparison_results.zip\"\n",
+    "with zipfile.ZipFile(archive, \"w\", compression=zipfile.ZIP_DEFLATED) as bundle:\n",
+    "    for path in sorted(RUN_ROOT.rglob(\"*\")):\n",
+    "        if path.is_file() and path.suffix in (\".csv\", \".json\") and \"raw\" not in path.relative_to(RUN_ROOT).parts:\n",
+    "            bundle.write(path, path.relative_to(RUN_ROOT).as_posix())\n",
+    "print(f\"Saved {archive}\")\n",
+    "try:\n",
+    "    from google.colab import files\n",
+    "except ImportError:\n",
+    "    pass\n",
+    "else:\n",
+    "    files.download(str(archive))\n"
+   ],
+   "execution_count": null,
+   "outputs": [],
+   "id": "a1-comparison-18"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sources and scope\n",
+    "\n",
+    "Dataset URLs/targets follow the primary DeepChem loaders: [ESOL](https://github.com/deepchem/deepchem/blob/master/deepchem/molnet/load_function/delaney_datasets.py), [FreeSolv](https://github.com/deepchem/deepchem/blob/master/deepchem/molnet/load_function/freesolv_dataset.py), [BBBP](https://github.com/deepchem/deepchem/blob/master/deepchem/molnet/load_function/bbbp_datasets.py). This notebook deliberately applies one disclosed scaffold-split recipe to all three, rather than claiming to reproduce each loader's defaults. Data use remains subject to the source datasets' terms.\n",
+    "\n",
+    "Model weights: [GRASP](https://huggingface.co/caithmac/GRASP), [A1](https://huggingface.co/caithmac/GRASP-A1). Code and model terms are CC BY-NC 4.0 with the DeBERTa implementation's separate MIT attribution. See the model cards for training data and limitations. This notebook has no embedded tokens and is saved without outputs.\n"
+   ],
+   "id": "a1-comparison-19"
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "provenance": [],
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "protocol_helper_sha256": "5bff4850d4b7da22944c50ac1a741c6b1ca8738347ec5a55ec3bc9140d3aaaae"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/comparison_helpers.py

@@ -0,0 +1,205 @@
+"""Shared, auditable protocol embedded in the self-contained Colab notebook."""
+import gc
+import hashlib
+import io
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import requests
+import torch
+from rdkit import Chem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+from sklearn.metrics import average_precision_score, mean_absolute_error, mean_squared_error, r2_score, roc_auc_score
+
+DATASETS = {
+    "esol": dict(filename="delaney-processed.csv", target="measured log solubility in mols per litre", task="regression",
+                 sha256="8c06a76f0c6487d29ab0f903e6a7a7139f189ab3c1178f159c8be8964602f189"),
+    "freesolv": dict(filename="freesolv.csv.gz", target="y", task="regression",
+                     sha256="008e9a8a9e018c6b77eb00df7153d0bf2e35a1d663f18f301de522fe3f8711d7"),
+    "bbbp": dict(filename="BBBP.csv", target="p_np", task="binary",
+                 sha256="d07a38487aeac5cee5508413e468043ef3097451d2a112701c2d60be9ec6b662"),
+}
+PROTOCOL_VERSION = "grasp-a1-paired-colab-20261015-v1"
+
+
+def sha256(path):
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as stream:
+        for block in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def write_json(path, value):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    Path(path).write_text(json.dumps(value, indent=2, allow_nan=False) + "\n")
+
+
+def prepare_task(name, cache_dir, output_dir, vocabulary):
+    from mole.training.data.datasets import MolDataset
+    spec = DATASETS[name]
+    cache_dir, output_dir = Path(cache_dir), Path(output_dir)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    raw = cache_dir / spec["filename"]
+    url = "https://deepchemdata.s3.us-west-1.amazonaws.com/datasets/" + spec["filename"]
+    if not raw.is_file():
+        response = requests.get(url, timeout=90)
+        response.raise_for_status()
+        raw.write_bytes(response.content)
+    assert sha256(raw) == spec["sha256"], f"Source changed for {name}; review it before updating its checksum."
+    frame = pd.read_csv(io.BytesIO(raw.read_bytes()), compression="gzip" if raw.suffix == ".gz" else None)
+    accepted, audit = [], []
+    for source_row, row in frame.iterrows():
+        reason = ""
+        try:
+            target = float(row[spec["target"]])
+            if not np.isfinite(target):
+                raise ValueError("nonfinite_target")
+            if spec["task"] == "binary" and target not in (0, 1):
+                raise ValueError("nonbinary_target")
+            mol = Chem.MolFromSmiles(str(row["smiles"]))
+            if mol is None or not mol.GetNumAtoms():
+                raise ValueError("unparsable_or_empty")
+            if mol.GetNumAtoms() > 511:
+                raise ValueError("above_loader_atom_limit")
+            smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
+            # Check the exact deployed featurizer once, before either model receives a split.
+            MolDataset(pd.Series([smiles]), vocabulary, cls_token=True)[0]
+            connectivity = Chem.MolToSmiles(mol, isomericSmiles=False)
+            scaffold = Chem.MolToSmiles(MurckoScaffold.GetScaffoldForMol(mol), isomericSmiles=False) or "<acyclic>"
+            accepted.append(dict(source_row=int(source_row), smiles=smiles, target=target,
+                                 connectivity=connectivity, scaffold=scaffold))
+        except Exception as exc:
+            reason = f"{type(exc).__name__}: {exc}"
+        audit.append(dict(source_row=int(source_row), original_smiles=str(row["smiles"]), excluded=bool(reason), reason=reason))
+    pd.DataFrame(audit).to_csv(output_dir / "data_audit.csv", index=False)
+    result = pd.DataFrame(accepted)
+    assert len(result), f"No usable rows for {name}"
+    result.to_csv(output_dir / "eligible_rows.csv", index=False)
+    write_json(output_dir / "data_manifest.json", dict(dataset=name, url=url, source_sha256=spec["sha256"],
+               source_rows=len(frame), eligible_rows=len(result), excluded_rows=len(frame) - len(result),
+               policy="Canonical isomeric SMILES; no charge/salt/tautomer repair; retain duplicates within one scaffold partition."))
+    return result
+
+
+def scaffold_split(frame, seed):
+    groups = [group.index.to_list() for _, group in frame.groupby("scaffold", sort=True)]
+    np.random.default_rng(seed).shuffle(groups)
+    groups.sort(key=len, reverse=True)
+    partitions = {"train": [], "valid": [], "test": []}
+    for group in groups:
+        if len(partitions["train"]) + len(group) <= 0.8 * len(frame):
+            partitions["train"].extend(group)
+        elif len(partitions["valid"]) + len(group) <= 0.1 * len(frame):
+            partitions["valid"].extend(group)
+        else:
+            partitions["test"].extend(group)
+    result = {key: frame.loc[sorted(rows)].reset_index(drop=True) for key, rows in partitions.items()}
+    assert sum(map(len, result.values())) == len(frame)
+    assert all(len(part) for part in result.values()), "Too few scaffold groups to create three nonempty splits."
+    for column in ("source_row", "connectivity", "scaffold"):
+        sets = [set(part[column]) for part in result.values()]
+        assert not (sets[0] & sets[1] or sets[0] & sets[2] or sets[1] & sets[2]), f"Split leakage: {column}"
+    return result
+
+
+def save_splits(frame, task, directory, seed):
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    parts = scaffold_split(frame, seed)
+    if task == "binary":
+        assert all(set(part.target) == {0, 1} for part in parts.values()), "A binary partition lacks one class; do not silently resplit."
+    csv_bytes = {split: part.to_csv(index=False).encode("utf-8") for split, part in parts.items()}
+    for split, data in csv_bytes.items():
+        path = directory / f"{split}.csv"
+        assert not path.is_file() or path.read_bytes() == data, "Existing splits differ: choose a new RUN_ROOT."
+    for split, data in csv_bytes.items():
+        (directory / f"{split}.csv").write_bytes(data)
+    manifest = dict(split_seed=seed, policy="Shared scaffold-grouped 80/10/10 target proportions; all acyclic molecules form one group.",
+                    counts={key: len(part) for key, part in parts.items()},
+                    sha256={key: sha256(directory / f"{key}.csv") for key in parts},
+                    disjoint_source_rows=True, disjoint_connectivity=True, disjoint_scaffolds=True)
+    write_json(directory / "split_manifest.json", manifest)
+    return manifest
+
+
+def enable_matched_initialization():
+    import grasp.encoder as encoder_module
+    import grasp.predictor as predictor_module
+    from mole.training.models.mole import AtomEnvEmbeddings
+    class MatchedEmbeddings(AtomEnvEmbeddings):
+        def __init__(self, config):
+            # Loading a different encoder architecture must not shift task-head or loader RNG.
+            with torch.random.fork_rng(devices=[]):
+                super().__init__(config)
+    encoder_module.AtomEnvEmbeddings = predictor_module.AtomEnvEmbeddings = MatchedEmbeddings
+
+
+def head_fingerprint(model_dir, seed, method):
+    from grasp import GRASPEncoder, GRASPPredictor
+    torch.manual_seed(seed)
+    encoder = GRASPEncoder.from_pretrained(model_dir, device="cpu")
+    predictor = GRASPPredictor.from_encoder(encoder, task="regression", method=method)
+    digest = hashlib.sha256()
+    for key, value in sorted(predictor.model.state_dict().items()):
+        if not key.startswith("encoder."):
+            digest.update(key.encode())
+            digest.update(value.detach().cpu().contiguous().numpy().tobytes())
+    state = hashlib.sha256(torch.get_rng_state().numpy().tobytes()).hexdigest()
+    del predictor, encoder
+    gc.collect()
+    return dict(readout_sha256=digest.hexdigest(), post_initialization_rng_sha256=state)
+
+
+def train_one(model_dir, split_dir, output_dir, *, task, method, seed, epochs, patience, batch_size, learning_rate, device, identity):
+    from grasp.train import fit
+    split_dir, output_dir = Path(split_dir), Path(output_dir)
+    split_hashes = {key: sha256(split_dir / f"{key}.csv") for key in ("train", "valid", "test")}
+    assert split_hashes == json.loads((split_dir / "split_manifest.json").read_text())["sha256"], "Shared split files changed."
+    plan = dict(protocol=PROTOCOL_VERSION, identity=identity, split_sha256=split_hashes, task=task, method=method,
+                seed=seed, epochs=epochs, patience=patience, batch_size=batch_size, learning_rate=learning_rate, device=device)
+    plan_file = output_dir / "run_plan.json"
+    if plan_file.is_file():
+        assert json.loads(plan_file.read_text()) == plan, "Settings/splits changed: choose a new RUN_ROOT."
+    else:
+        write_json(plan_file, plan)
+    complete = output_dir / "TRAINING_COMPLETE.json"
+    if complete.is_file():
+        assert (output_dir / "predictor.safetensors").is_file()
+        return json.loads(complete.read_text())
+    # Test labels are never supplied to fitting or early stopping.
+    outcome = fit(model_source=model_dir, train_csv=split_dir / "train.csv", valid_csv=split_dir / "valid.csv",
+                  output_dir=output_dir, task=task, method=method, seed=seed, epochs=epochs, patience=patience,
+                  batch_size=batch_size, learning_rate=learning_rate, device=device)
+    write_json(complete, outcome)
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return outcome
+
+
+def score_predictions(target, prediction, task):
+    target, prediction = np.asarray(target), np.asarray(prediction)
+    assert target.shape == prediction.shape and np.isfinite(prediction).all()
+    if task == "regression":
+        return dict(mae=float(mean_absolute_error(target, prediction)), rmse=float(np.sqrt(mean_squared_error(target, prediction))),
+                    r2=float(r2_score(target, prediction)))
+    assert set(target) == {0, 1} and ((prediction >= 0) & (prediction <= 1)).all()
+    return dict(auroc=float(roc_auc_score(target, prediction)), auprc=float(average_precision_score(target, prediction)))
+
+
+def paired_table(results):
+    pairs = []
+    for (dataset, seed), group in results.groupby(["dataset", "seed"]):
+        assert set(group.model) == {"GRASP", "A1"} and len(group) == 2
+        group = group.set_index("model")
+        metric = "mae" if group.loc["A1", "task"] == "regression" else "auroc"
+        baseline, a1 = float(group.loc["GRASP", metric]), float(group.loc["A1", metric])
+        delta = a1 - baseline
+        pairs.append(dict(dataset=dataset, seed=int(seed), metric=metric, GRASP=baseline, A1=a1,
+                          delta_A1_minus_GRASP=delta, direction="lower is better" if metric == "mae" else "higher is better",
+                          A1_better=bool(delta < 0 if metric == "mae" else delta > 0)))
+    return pd.DataFrame(pairs)

notebooks/verification.json

@@ -0,0 +1,46 @@
+{
+  "datasets": {
+    "esol": {
+      "eligible_rows": 1128,
+      "exclusions": 0,
+      "split_counts": {
+        "train": 902,
+        "valid": 112,
+        "test": 114
+      },
+      "deterministic": true,
+      "disjoint_scaffolds_and_connectivity": true
+    },
+    "freesolv": {
+      "eligible_rows": 642,
+      "exclusions": 0,
+      "split_counts": {
+        "train": 513,
+        "valid": 64,
+        "test": 65
+      },
+      "deterministic": true,
+      "disjoint_scaffolds_and_connectivity": true
+    },
+    "bbbp": {
+      "eligible_rows": 2039,
+      "exclusions": 11,
+      "split_counts": {
+        "train": 1631,
+        "valid": 203,
+        "test": 205
+      },
+      "deterministic": true,
+      "disjoint_scaffolds_and_connectivity": true
+    }
+  },
+  "status": "passed",
+  "notebook_cells_compiled": true,
+  "paired_full_adaptation_smoke_fits": 4,
+  "initial_heads_and_rng_identical": true,
+  "resume_skip_and_mismatch_rejection": true,
+  "compact_archive": true,
+  "scope": "Real full-dataset split audits plus tiny CPU notebook fits; no Colab GPU or full benchmark run",
+  "auth_public_private_and_network_error_checks": true,
+  "notebook_sha256": "cc71af7c0c1572a0718cbd5276e2a76cc769ca72fb46f02447cfd383d28fd6dd"
+}