{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2f61e03",
   "metadata": {},
   "source": [
    "# GRASP-A3 Colab quickstart\n",
    "\n",
    "Public model: [caithmac/GRASP-A3](https://huggingface.co/caithmac/GRASP-A3).\n",
    "Use a fresh Python 3.10–3.13 runtime. A GPU is optional for this small example.\n",
    "Run the three code cells in order. No Hub token is required.\n",
    "\n",
    "The loader and model hashes are verified against the audited release, and one Hub\n",
    "revision is used throughout. Installing an updated loader in a runtime that already\n",
    "imported an older `grasp` requires restarting that runtime before loading A3.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b593ed9",
   "metadata": {},
   "source": [
    "## 1. Install the A3-compatible loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "062b6490",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fresh Colab runtime recommended; restart if an older grasp was imported.\n",
    "import hashlib, json, shutil, subprocess, sys, urllib.request\n",
    "from pathlib import Path\n",
    "\n",
    "assert (3, 10) <= sys.version_info[:2] < (3, 14), \"Use Python 3.10–3.13.\"\n",
    "MODEL_ID = \"caithmac/GRASP-A3\"\n",
    "with urllib.request.urlopen(\"https://huggingface.co/api/models/\" + MODEL_ID, timeout=60) as response:\n",
    "    model_info = json.load(response)\n",
    "assert model_info[\"id\"] == MODEL_ID and not model_info[\"private\"]\n",
    "MODEL_REVISION = model_info[\"sha\"]\n",
    "assert len(MODEL_REVISION) == 40 and all(c in \"0123456789abcdef\" for c in MODEL_REVISION)\n",
    "WHEEL_FILENAME = \"grasp_molecular-0.1.2-py3-none-any.whl\"\n",
    "WHEEL_SHA256 = \"914eaf0b182ad3d70ab57f0b38d6cd7d8e772f11dd0d95bd2e40ee5c1476b6d6\"\n",
    "WEIGHTS_SHA256 = \"b3a16758d213b1fd792eca74649d814cb36e1f588a931ba50252911ed38fd3ac\"\n",
    "wheel_path = Path.cwd() / WHEEL_FILENAME\n",
    "wheel_url = \"https://huggingface.co/\" + MODEL_ID + \"/resolve/\" + MODEL_REVISION + \"/\" + WHEEL_FILENAME\n",
    "with urllib.request.urlopen(wheel_url, timeout=60) as response, wheel_path.open(\"wb\") as target:\n",
    "    shutil.copyfileobj(response, target)\n",
    "assert hashlib.sha256(wheel_path.read_bytes()).hexdigest() == WHEEL_SHA256, \"Loader hash mismatch.\"\n",
    "subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"--quiet\", str(wheel_path)])\n",
    "print(\"Hub revision:\", MODEL_REVISION)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9b1d01a",
   "metadata": {},
   "source": [
    "## 2. Verify and load the public encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ea6513",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from huggingface_hub import hf_hub_download\n",
    "from grasp import GRASPEncoder\n",
    "\n",
    "# Verify the downloaded public weights before loading the same pinned revision.\n",
    "weight_path = Path(hf_hub_download(MODEL_ID, \"model.safetensors\", revision=MODEL_REVISION))\n",
    "digest = hashlib.sha256()\n",
    "with weight_path.open(\"rb\") as stream:\n",
    "    for block in iter(lambda: stream.read(8 * 1024**2), b\"\"):\n",
    "        digest.update(block)\n",
    "assert digest.hexdigest() == WEIGHTS_SHA256, \"Weight hash mismatch.\"\n",
    "encoder = GRASPEncoder.from_pretrained(MODEL_ID, revision=MODEL_REVISION)\n",
    "assert encoder.config[\"architecture\"][\"factorized_atom_features\"]\n",
    "assert sum(p.numel() for p in encoder.model.parameters()) == 93_089_496\n",
    "\n",
    "smiles = [\"CCO\", \"c1ccccc1\", \"N[C@@H](C)C(=O)O\", \"F[C@H](Cl)Br\", \"F[C@@H](Cl)Br\"]\n",
    "embeddings = encoder.encode(smiles, batch_size=4)\n",
    "assert embeddings.shape == (5, 768) and np.isfinite(embeddings).all()\n",
    "print(\"Device:\", encoder.device, \"| Embedding shape:\", embeddings.shape)\n",
    "display(pd.DataFrame(embeddings[:, :5], index=smiles, columns=[f\"dim_{i}\" for i in range(5)]))\n",
    "np.save(\"a3_example_embeddings.npy\", embeddings)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a0b894c",
   "metadata": {},
   "source": [
    "## 3. Encode your own molecules and save the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eec7a254",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace these with your own valid SMILES. Outputs are embeddings, not potency predictions.\n",
    "my_smiles = [\"CC(=O)O\", \"CCN\", \"c1ccncc1\"]\n",
    "my_embeddings = encoder.encode(my_smiles, batch_size=4)\n",
    "assert my_embeddings.shape == (len(my_smiles), 768) and np.isfinite(my_embeddings).all()\n",
    "np.save(\"a3_my_embeddings.npy\", my_embeddings)\n",
    "pd.DataFrame({\"smiles\": my_smiles}).to_csv(\"a3_my_smiles.csv\", index=False)\n",
    "print(\"Saved a3_my_embeddings.npy and a3_my_smiles.csv\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d930d899",
   "metadata": {},
   "source": [
    "## Optional: train a task head or fine-tune on your CSVs\n",
    "\n",
    "Upload separate `train.csv` and `valid.csv` files with `smiles,target` columns.\n",
    "Copy this template into a new code cell only after those files are ready:\n",
    "\n",
    "```python\n",
    "from grasp.train import fit\n",
    "from grasp import GRASPPredictor\n",
    "\n",
    "fit(model_source=str(weight_path.parent), train_csv=\"train.csv\", valid_csv=\"valid.csv\",\n",
    "    output_dir=\"my_a3_predictor\", task=\"regression\", method=\"frozen\",\n",
    "    epochs=20, patience=5, batch_size=16)\n",
    "predictor = GRASPPredictor.from_pretrained(\"my_a3_predictor\")\n",
    "print(predictor.predict([\"CCO\", \"CCN\"]))\n",
    "```\n",
    "\n",
    "`task` is `regression` or `binary` (binary targets must be 0/1).\n",
    "`method` is `frozen`, `full` or `lora`. The template is illustrative; its settings\n",
    "are not a benchmark or a universal recommended tuning recipe. The predictor learns\n",
    "layer mixing, atom pooling and a nonlinear task head. The encoder package does not\n",
    "include endpoint-specific prediction heads. Use shared fixed splits for comparisons.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b2bb03b",
   "metadata": {},
   "source": [
    "## Interpretation and scope\n",
    "\n",
    "`encode` returns final-layer CLS vectors. Reported OpenADMET results use full\n",
    "fine-tuning with a mixed-layer predictor, not a frozen CLS probe. A3 changes both\n",
    "the atom representation and generator/corruption task. It is not a representation-only\n",
    "causal experiment. Source overlap, censored labels and assay comparability require\n",
    "task-specific handling; the checkpoint's Step 2 supervision is binary.\n",
    "\n",
    "See the model card for license (CC BY-NC 4.0), provenance, limits and validation.\n",
    "The publication smoke test checks hosted files/inference on local CPU; this notebook\n",
    "has not been run on a Google-hosted Colab GPU during publication.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "GRASP_A3_Colab_Quickstart.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}