Upload NativeSpecZ-296M.ipynb

NativeSpecZ-296M.ipynb

@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# NativeSpecZ-296M — Demo Notebook (scale-up ablation)\n",
+    "\n",
+    "A 296M-parameter unimodal foundation model for DESI spectra, trained from scratch (no AION pretrained weights) with wavelength-grid jitter augmentation. Hits the spec's ~300M target and is competitive with AION-base on DESI in-distribution, but does NOT beat AION on out-of-distribution data — the smaller NativeSpecZ-FM-76M is the recommended headline.\n",
+    "\n",
+    "**Approach A**: z head trained jointly with the encoder.  \n",
+    "**Approach B**: `[Z_MASK]` token always masked.\n",
+    "\n",
+    "NOTE: load with `strict=False` — the current `hybrid_redshift.py` has `z_rerank_head` / `z_calib_head` modules added after this checkpoint was trained; they are unused in the `bin_residual` prediction path."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys, json\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "sys.path.append(\"code\")\n",
+    "from hybrid_redshift import HybridSpecZ, RawCollatorConfig, RawSpectraCollator, move_to_device\n",
+    "from data import SpectraListDataset, parse_mmu_example\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Device: {device}\")\n",
+    "torch.manual_seed(42); np.random.seed(42)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the 296M model (strict=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ckpt = torch.load(\"weights/best.pt\", map_location=device, weights_only=False)\n",
+    "a = ckpt[\"args\"]\n",
+    "model = HybridSpecZ(\n",
+    "    d_model=a[\"d_model\"], conv_width=a[\"conv_width\"], layers=a[\"layers\"],\n",
+    "    heads=a[\"heads\"], dropout=a[\"dropout\"], z_bins=a[\"z_bins\"],\n",
+    "    stem_stride=a[\"stem_stride\"], rec_hidden_mult=a[\"rec_hidden_mult\"],\n",
+    "    rec_refine_width=a[\"rec_refine_width\"], rec_refine_kernel=a[\"rec_refine_kernel\"],\n",
+    "    layerscale_init=a[\"layerscale_init\"], prediction_mode=a[\"prediction_mode\"],\n",
+    "    bin_temperature=a[\"bin_temperature\"], residual_scale=a[\"residual_scale\"],\n",
+    "    candidate_topk=a[\"candidate_topk\"],\n",
+    ").to(device)\n",
+    "missing, unexpected = model.load_state_dict(ckpt[\"model\"], strict=False)\n",
+    "model.eval()\n",
+    "print(f\"Params: {sum(p.numel() for p in model.parameters()):,}\")\n",
+    "print(f\"Missing keys (unused rerank/calib heads): {len(missing)}; unexpected: {len(unexpected)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load DESI test data + run inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "from tqdm.auto import tqdm\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "ds = load_dataset(\"MultimodalUniverse/desi\", split=\"train\", streaming=True).shuffle(seed=4242, buffer_size=5000)\n",
+    "samples = []\n",
+    "for ex in ds:\n",
+    "    p = parse_mmu_example(ex)\n",
+    "    if p is not None: samples.append(p)\n",
+    "    if len(samples) >= 500: break\n",
+    "print(f\"Loaded {len(samples)} DESI spectra\")\n",
+    "\n",
+    "cfg = RawCollatorConfig(target_length=a[\"target_length\"], eval_mask_ratio=0.30, mask_mode=\"mixed_span\", mask_span_min=a[\"mask_span_min\"], mask_span_max=a[\"mask_span_max\"])\n",
+    "loader = DataLoader(SpectraListDataset(samples, np.arange(len(samples))), batch_size=8, shuffle=False, num_workers=2, collate_fn=RawSpectraCollator(cfg, train=False, seed=31415))\n",
+    "z_true_all, y_pred_all = [], []\n",
+    "with torch.no_grad():\n",
+    "    for batch in tqdm(loader):\n",
+    "        batch = move_to_device(batch, device)\n",
+    "        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type==\"cuda\"):\n",
+    "            out = model(batch[\"x\"], batch[\"valid\"], batch[\"loglam\"])\n",
+    "        z_true_all.append(batch[\"z\"].cpu().numpy())\n",
+    "        y_pred_all.append(out.get(\"y_pred\", out[\"y_mu\"]).float().cpu().numpy())\n",
+    "z_true = np.concatenate(z_true_all)\n",
+    "z_pred = np.clip(np.expm1(np.concatenate(y_pred_all)), 0, 6)\n",
+    "mae = float(np.mean(np.abs(z_pred - z_true)))\n",
+    "r = float(np.corrcoef(z_true, z_pred)[0,1])\n",
+    "print(f\"DESI MAE={mae:.4f}  Pearson r={r:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Held-out results (precomputed, n=2500/2000/2000) + comparison\n",
+    "\n",
+    "| Dataset | NativeSpecZ-296M | NativeSpecZ-76M (headline) | AION-base |\n",
+    "|---|---:|---:|---:|\n",
+    "| **DESI** | **0.067** | 0.069 | 0.074 |\n",
+    "| SDSS (real non-DESI) | 0.314 | 0.382 | **0.127** |\n",
+    "| VIPERS (real non-DESI) | 0.316 | **0.172** | 0.274 |\n",
+    "\n",
+    "The 296M ties the 76M on DESI (both beat AION) and improves SDSS over the 76M, but loses the 76M's strong VIPERS result. It beats AION-base on no OOD dataset, so the 76M remains the recommended headline; the 296M is the scale-up ablation that hits the ~300M target."
+   ]
+  }
+ ],
+ "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.11"}},
+ "nbformat": 4,
+ "nbformat_minor": 4
+}