Upload text2roi_colab_v4.ipynb
Browse files- text2roi_colab_v4.ipynb +14 -130
text2roi_colab_v4.ipynb
CHANGED
|
@@ -45,14 +45,17 @@
|
|
| 45 |
"execution_count": null,
|
| 46 |
"metadata": {},
|
| 47 |
"outputs": [],
|
| 48 |
-
"source":
|
| 49 |
-
"# ── Install dependencies ──────────────────────────────────────────────────────\n",
|
| 50 |
-
"!pip install -q torch transformers huggingface_hub numpy matplotlib librosa\n",
|
| 51 |
-
"!pip install -q nilearn # optional: brain surface visualization\n",
|
| 52 |
-
"print('Done.')"
|
| 53 |
-
],
|
| 54 |
"id": "install"
|
| 55 |
},
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
{
|
| 57 |
"cell_type": "code",
|
| 58 |
"execution_count": null,
|
|
@@ -77,49 +80,7 @@
|
|
| 77 |
"execution_count": null,
|
| 78 |
"metadata": {},
|
| 79 |
"outputs": [],
|
| 80 |
-
"source":
|
| 81 |
-
"# ── Locate model files (local first, HuggingFace fallback) ───────────────────\n",
|
| 82 |
-
"#\n",
|
| 83 |
-
"# Priority order:\n",
|
| 84 |
-
"# 1. Files uploaded to this Colab session (e.g. from the self-contained zip)\n",
|
| 85 |
-
"# 2. HuggingFace Hub download (requires internet; ~7-17 MB per model)\n",
|
| 86 |
-
"#\n",
|
| 87 |
-
"# To use the bundled zip offline:\n",
|
| 88 |
-
"# - Extract NeuroText_v4_Demo.zip\n",
|
| 89 |
-
"# - Upload the .pt files and predict.py to Colab via the Files panel\n",
|
| 90 |
-
"# - Re-run this cell — it will find them locally and skip the HF download\n",
|
| 91 |
-
"\n",
|
| 92 |
-
"import os, shutil\n",
|
| 93 |
-
"from pathlib import Path\n",
|
| 94 |
-
"\n",
|
| 95 |
-
"REPO = 'ffh92r32rm0/Amphora_NeuroText'\n",
|
| 96 |
-
"\n",
|
| 97 |
-
"def _resolve(filename, repo=REPO):\n",
|
| 98 |
-
" \"\"\"Return local path to filename, downloading from HF if not found locally.\"\"\"\n",
|
| 99 |
-
" local = Path(filename)\n",
|
| 100 |
-
" if local.exists():\n",
|
| 101 |
-
" print(f' {filename}: found locally')\n",
|
| 102 |
-
" return str(local)\n",
|
| 103 |
-
" print(f' {filename}: not found locally, downloading from HuggingFace...')\n",
|
| 104 |
-
" from huggingface_hub import hf_hub_download\n",
|
| 105 |
-
" path = hf_hub_download(repo, filename)\n",
|
| 106 |
-
" print(f' {filename}: downloaded to {path}')\n",
|
| 107 |
-
" return path\n",
|
| 108 |
-
"\n",
|
| 109 |
-
"print('Resolving files...')\n",
|
| 110 |
-
"ckpt_path = _resolve(MODEL_FILE)\n",
|
| 111 |
-
"predict_path = _resolve('predict.py')\n",
|
| 112 |
-
"examples_path = _resolve('examples_cache.npz') if Path('examples_cache.npz').exists() else None\n",
|
| 113 |
-
"\n",
|
| 114 |
-
"# make predict.py importable from CWD\n",
|
| 115 |
-
"if predict_path != 'predict.py':\n",
|
| 116 |
-
" shutil.copy(predict_path, 'predict.py')\n",
|
| 117 |
-
"\n",
|
| 118 |
-
"print(f'\\nCheckpoint : {ckpt_path}')\n",
|
| 119 |
-
"print(f'predict.py : ready')\n",
|
| 120 |
-
"print(f'examples : {\"found\" if examples_path else \"not available (cached examples will be skipped)\"}')\n",
|
| 121 |
-
"print('Ready.')"
|
| 122 |
-
],
|
| 123 |
"id": "resolve-files"
|
| 124 |
},
|
| 125 |
{
|
|
@@ -166,33 +127,7 @@
|
|
| 166 |
"execution_count": null,
|
| 167 |
"metadata": {},
|
| 168 |
"outputs": [],
|
| 169 |
-
"source": [
|
| 170 |
-
"# ── Helper: run inference + plot ──────────────────────────────────────────────\n",
|
| 171 |
-
"import matplotlib.pyplot as plt\n",
|
| 172 |
-
"\n",
|
| 173 |
-
"def predict_from_features(feat_1d, model, in_dim):\n",
|
| 174 |
-
" \"\"\"Run model on a 1-D feature vector. For combined (3840d), zero-pads the whisper slot.\"\"\"\n",
|
| 175 |
-
" if in_dim == 3840 and feat_1d.shape[0] == 2560:\n",
|
| 176 |
-
" feat_1d = np.concatenate([np.zeros(1280, dtype=np.float32), feat_1d])\n",
|
| 177 |
-
" with torch.no_grad():\n",
|
| 178 |
-
" pred = model(torch.from_numpy(feat_1d[:in_dim]).unsqueeze(0)).squeeze(0).numpy()\n",
|
| 179 |
-
" return dict(zip(ROI_NAMES, pred.tolist()))\n",
|
| 180 |
-
"\n",
|
| 181 |
-
"def plot_rois(roi_map, title='', top_n=20):\n",
|
| 182 |
-
" items = sorted(roi_map.items(), key=lambda x: -x[1])[:top_n]\n",
|
| 183 |
-
" names, vals = zip(*items)\n",
|
| 184 |
-
" vmin, vmax = min(vals), max(vals)\n",
|
| 185 |
-
" colors = plt.cm.RdYlGn([(v - vmin) / (vmax - vmin + 1e-9) for v in vals])\n",
|
| 186 |
-
" fig, ax = plt.subplots(figsize=(9, 5))\n",
|
| 187 |
-
" ax.barh(range(len(names)), vals, color=colors)\n",
|
| 188 |
-
" ax.set_yticks(range(len(names))); ax.set_yticklabels(names, fontsize=9)\n",
|
| 189 |
-
" ax.axvline(0, color='gray', linewidth=0.7)\n",
|
| 190 |
-
" ax.set_xlabel('Predicted activation (z-score)')\n",
|
| 191 |
-
" ax.set_title(title[:90], fontsize=10)\n",
|
| 192 |
-
" ax.invert_yaxis(); plt.tight_layout(); plt.show()\n",
|
| 193 |
-
"\n",
|
| 194 |
-
"print('Helpers ready.')"
|
| 195 |
-
],
|
| 196 |
"id": "helpers"
|
| 197 |
},
|
| 198 |
{
|
|
@@ -201,7 +136,7 @@
|
|
| 201 |
"id": "brain-plotter-lib",
|
| 202 |
"metadata": {},
|
| 203 |
"outputs": [],
|
| 204 |
-
"source": "# ── Brain map plotter (fsaverage5 / HCP-MMP1, same mesh as TRIBE v2) ─────────\n# Commercial-friendly: MNE + Nilearn only (BSD). Does NOT import tribev2.\n
|
| 205 |
},
|
| 206 |
{
|
| 207 |
"cell_type": "code",
|
|
@@ -216,30 +151,7 @@
|
|
| 216 |
"execution_count": null,
|
| 217 |
"metadata": {},
|
| 218 |
"outputs": [],
|
| 219 |
-
"source": [
|
| 220 |
-
"# ── Option A: pre-cached text examples (CPU, instant, no encoder needed) ─────\n",
|
| 221 |
-
"# Uses the examples_cache.npz bundled in the zip, or downloaded from HF.\n",
|
| 222 |
-
"# Skipped automatically if the file is unavailable.\n",
|
| 223 |
-
"\n",
|
| 224 |
-
"if examples_path:\n",
|
| 225 |
-
" cache = np.load(examples_path, allow_pickle=True)\n",
|
| 226 |
-
" sentences = [s.decode() if isinstance(s, bytes) else str(s) for s in cache['sentences']]\n",
|
| 227 |
-
" embeddings = cache['embeddings'].astype(np.float32) # (N, 2560)\n",
|
| 228 |
-
" print(f'Loaded {len(sentences)} cached examples')\n",
|
| 229 |
-
"\n",
|
| 230 |
-
" # ── pick a sentence to visualise ──\n",
|
| 231 |
-
" idx = 0 # change index or replace sentence below\n",
|
| 232 |
-
" sentence = sentences[idx]\n",
|
| 233 |
-
" roi_map = predict_from_features(embeddings[idx], model, IN_DIM)\n",
|
| 234 |
-
"\n",
|
| 235 |
-
" print(f'\\nInput: \"{sentence}\"')\n",
|
| 236 |
-
" print('Top 10 ROIs:')\n",
|
| 237 |
-
" for roi, val in sorted(roi_map.items(), key=lambda x: -x[1])[:10]:\n",
|
| 238 |
-
" print(f' {roi:<16} {val:+.4f}')\n",
|
| 239 |
-
" plot_rois(roi_map, title=sentence)\n",
|
| 240 |
-
"else:\n",
|
| 241 |
-
" print('examples_cache.npz not found — skip to Option B (audio) or Option C (live text).')"
|
| 242 |
-
],
|
| 243 |
"id": "cached-examples"
|
| 244 |
},
|
| 245 |
{
|
|
@@ -288,35 +200,7 @@
|
|
| 288 |
"execution_count": null,
|
| 289 |
"metadata": {},
|
| 290 |
"outputs": [],
|
| 291 |
-
"source": [
|
| 292 |
-
"# ── Option C: Custom text inference with Qwen3 (GPU recommended, ~8 GB) ──────\n",
|
| 293 |
-
"import torch.nn.functional as F\n",
|
| 294 |
-
"from transformers import AutoModel, AutoTokenizer\n",
|
| 295 |
-
"\n",
|
| 296 |
-
"YOUR_TEXT = 'listening to a symphony building to its climax'\n",
|
| 297 |
-
"\n",
|
| 298 |
-
"print('Loading Qwen3-Embedding-4B (downloads ~8 GB on first run)...')\n",
|
| 299 |
-
"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
|
| 300 |
-
"tok = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-4B', padding_side='left')\n",
|
| 301 |
-
"qwen = AutoModel.from_pretrained(\n",
|
| 302 |
-
" 'Qwen/Qwen3-Embedding-4B',\n",
|
| 303 |
-
" torch_dtype=torch.bfloat16 if device != 'cpu' else torch.float32\n",
|
| 304 |
-
").to(device).eval()\n",
|
| 305 |
-
"\n",
|
| 306 |
-
"with torch.no_grad():\n",
|
| 307 |
-
" enc = tok([YOUR_TEXT], return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)\n",
|
| 308 |
-
" h = qwen(**enc).last_hidden_state[:, -1].float()\n",
|
| 309 |
-
" emb = F.normalize(h, p=2, dim=1).cpu().numpy()[0].astype(np.float32) # (2560,)\n",
|
| 310 |
-
"\n",
|
| 311 |
-
"del qwen; torch.cuda.empty_cache() if device == 'cuda' else None\n",
|
| 312 |
-
"\n",
|
| 313 |
-
"roi_map_text = predict_from_features(emb, model, IN_DIM)\n",
|
| 314 |
-
"print(f'\\nInput: \"{YOUR_TEXT}\"')\n",
|
| 315 |
-
"print('Top 10 ROIs:')\n",
|
| 316 |
-
"for roi, val in sorted(roi_map_text.items(), key=lambda x: -x[1])[:10]:\n",
|
| 317 |
-
" print(f' {roi:<16} {val:+.4f}')\n",
|
| 318 |
-
"plot_rois(roi_map_text, title=YOUR_TEXT)"
|
| 319 |
-
],
|
| 320 |
"id": "custom-text"
|
| 321 |
},
|
| 322 |
{
|
|
|
|
| 45 |
"execution_count": null,
|
| 46 |
"metadata": {},
|
| 47 |
"outputs": [],
|
| 48 |
+
"source": "# ── Install dependencies ──────────────────────────────────────────────────────\n!pip install -q torch transformers huggingface_hub numpy matplotlib librosa\n!pip install -q nilearn mne # brain surface visualization\nprint('Done.')\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 49 |
"id": "install"
|
| 50 |
},
|
| 51 |
+
{
|
| 52 |
+
"cell_type": "code",
|
| 53 |
+
"execution_count": null,
|
| 54 |
+
"id": "hf-auth",
|
| 55 |
+
"metadata": {},
|
| 56 |
+
"outputs": [],
|
| 57 |
+
"source": "# ── HuggingFace authentication (required — repo is private) ──────────────────\n# Option 1 (recommended): add your HF token to Colab Secrets as \"HF_TOKEN\"\n# Runtime → Secrets → Add new secret → Name: HF_TOKEN, Value: hf_...\n# Option 2: run notebook_login() interactively (prompts for token below)\nimport os\nfrom huggingface_hub import login\n\ntry:\n from google.colab import userdata\n _token = userdata.get('HF_TOKEN')\n if _token:\n login(token=_token, add_to_git_credential=False)\n print('Logged in via Colab Secrets (HF_TOKEN).')\n else:\n raise KeyError('HF_TOKEN not set in Colab Secrets')\nexcept Exception:\n print('HF_TOKEN not found in Colab Secrets — opening interactive login...')\n from huggingface_hub import notebook_login\n notebook_login()\n"
|
| 58 |
+
},
|
| 59 |
{
|
| 60 |
"cell_type": "code",
|
| 61 |
"execution_count": null,
|
|
|
|
| 80 |
"execution_count": null,
|
| 81 |
"metadata": {},
|
| 82 |
"outputs": [],
|
| 83 |
+
"source": "# ── Locate model files (local first, HuggingFace fallback) ───────────────────\n#\n# Priority order:\n# 1. Files uploaded to this Colab session (e.g. from the self-contained zip)\n# 2. HuggingFace Hub download (requires internet; ~7-17 MB per model)\n#\n# To use the bundled zip offline:\n# - Extract NeuroText_v4_Demo.zip\n# - Upload the .pt files and predict.py to Colab via the Files panel\n# - Re-run this cell — it will find them locally and skip the HF download\n\nimport os, shutil\nfrom pathlib import Path\n\nREPO = 'ffh92r32rm0/Amphora_NeuroText'\n\ndef _resolve(filename, repo=REPO):\n \"\"\"Return local path to filename, downloading from HF if not found locally.\"\"\"\n local = Path(filename)\n if local.exists():\n print(f' {filename}: found locally')\n return str(local)\n print(f' {filename}: not found locally, downloading from HuggingFace...')\n from huggingface_hub import hf_hub_download\n path = hf_hub_download(repo, filename)\n print(f' {filename}: downloaded to {path}')\n return path\n\ndef _resolve_optional(filename, repo=REPO):\n \"\"\"Like _resolve but returns None on failure instead of raising.\"\"\"\n try:\n return _resolve(filename, repo)\n except Exception as e:\n print(f' {filename}: not available ({type(e).__name__}) — will be skipped')\n return None\n\nprint('Resolving files...')\nckpt_path = _resolve(MODEL_FILE)\npredict_path = _resolve('predict.py')\nexamples_path = _resolve_optional('examples_cache.npz')\n\n# make predict.py importable from CWD\nif predict_path != 'predict.py':\n shutil.copy(predict_path, 'predict.py')\n\nprint(f'\\nCheckpoint : {ckpt_path}')\nprint(f'predict.py : ready')\nprint(f'examples : {\"found\" if examples_path else \"not available (cached examples will be skipped)\"}')\nprint('Ready.')\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 84 |
"id": "resolve-files"
|
| 85 |
},
|
| 86 |
{
|
|
|
|
| 127 |
"execution_count": null,
|
| 128 |
"metadata": {},
|
| 129 |
"outputs": [],
|
| 130 |
+
"source": "# ── Helper: run inference + plot ──────────────────────────────────────────────\nimport matplotlib.pyplot as plt\n\ndef predict_from_features(feat_1d, model, in_dim):\n \"\"\"Run model on a 1-D feature vector.\n\n For the combined model (in_dim=3840) with a 2560-d Qwen3 embedding,\n zero-pads the whisper slot [0:1280] automatically.\n Raises ValueError if the embedding cannot be adapted to in_dim.\n \"\"\"\n if in_dim == 3840 and feat_1d.shape[0] == 2560:\n feat_1d = np.concatenate([np.zeros(1280, dtype=np.float32), feat_1d])\n if feat_1d.shape[0] < in_dim:\n raise ValueError(\n f\"Embedding dim {feat_1d.shape[0]} < model in_dim {in_dim}. \"\n f\"Make sure you are using the correct model for this embedding type. \"\n f\"For text (Qwen3 2560-d) use the combined model (in_dim=3840); \"\n f\"for audio (Whisper 1280-d) use the audio model (in_dim=1280).\"\n )\n with torch.no_grad():\n pred = model(torch.from_numpy(feat_1d[:in_dim]).unsqueeze(0)).squeeze(0).numpy()\n return dict(zip(ROI_NAMES, pred.tolist()))\n\ndef plot_rois(roi_map, title='', top_n=20):\n items = sorted(roi_map.items(), key=lambda x: -x[1])[:top_n]\n names, vals = zip(*items)\n vmin, vmax = min(vals), max(vals)\n colors = plt.cm.RdYlGn([(v - vmin) / (vmax - vmin + 1e-9) for v in vals])\n fig, ax = plt.subplots(figsize=(9, 5))\n ax.barh(range(len(names)), vals, color=colors)\n ax.set_yticks(range(len(names))); ax.set_yticklabels(names, fontsize=9)\n ax.axvline(0, color='gray', linewidth=0.7)\n ax.set_xlabel('Predicted activation (z-score)')\n ax.set_title(title[:90], fontsize=10)\n ax.invert_yaxis(); plt.tight_layout(); plt.show()\n\nprint('Helpers ready.')\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 131 |
"id": "helpers"
|
| 132 |
},
|
| 133 |
{
|
|
|
|
| 136 |
"id": "brain-plotter-lib",
|
| 137 |
"metadata": {},
|
| 138 |
"outputs": [],
|
| 139 |
+
"source": "# ── Brain map plotter (fsaverage5 / HCP-MMP1, same mesh as TRIBE v2) ─────────\n# Commercial-friendly: MNE + Nilearn only (BSD). Does NOT import tribev2.\n\nfrom __future__ import annotations\n\n!pip install -q mne nilearn Pillow\n\nfrom functools import lru_cache\nfrom pathlib import Path\nfrom typing import Dict, List, Mapping, Sequence, Tuple\nimport numpy as np\n\nFSAVERAGE5_VERTS_PER_HEMI = 10242\nFSAVERAGE5_VERTICES = FSAVERAGE5_VERTS_PER_HEMI * 2\n\n# ROI → HCP-MMP1 parcel mapping.\n# PPA/OFA entries are anatomical approximations (see comments).\nROI56_TO_HCP: Dict[str, List[Tuple[str, str]]] = {\n \"V1\": [(\"V1\",\"both\")], \"V2\": [(\"V2\",\"both\")], \"V3\": [(\"V3\",\"both\")], \"V4\": [(\"V4\",\"both\")],\n \"V3A\": [(\"V3A\",\"both\")], \"V3B\": [(\"V3B\",\"both\")], \"LO1\": [(\"LO1\",\"both\")], \"LO2\": [(\"LO2\",\"both\")],\n \"MT\": [(\"MT\",\"both\")], \"MST\": [(\"MST\",\"both\")], \"V7\": [(\"V7\",\"both\")], \"IPS1\": [(\"IPS1\",\"both\")],\n \"FFA-1\": [(\"FFC\",\"both\")], \"FFA-2\": [(\"FFC\",\"both\")],\n \"PPA\": [(\"PHA1\",\"both\"),(\"PHA2\",\"both\"),(\"PHA3\",\"both\")], # approx; PIT is posterior IT\n \"RSC\": [(\"POS1\",\"both\"),(\"7m\",\"both\")],\n \"OFA\": [(\"FFC\",\"both\"),(\"VVC\",\"both\")], # approx ventral temporal\n \"EBA\": [(\"FST\",\"both\")],\n \"IPS2\": [(\"IPS2\",\"both\")], \"IPS3\": [(\"IPS3\",\"both\")], \"IPS4\": [(\"IPS4\",\"both\")], \"IPS5\": [(\"IPS5\",\"both\")],\n \"SPL1\": [(\"SPL1\",\"both\")], \"hIP1\": [(\"7AL\",\"both\")], \"hIP2\": [(\"7PC\",\"both\")], \"hIP3\": [(\"7Am\",\"both\")],\n \"dlPFC\": [(\"9-46d\",\"both\"),(\"46\",\"both\")], \"vlPFC\": [(\"47l\",\"both\")], \"OFC\": [(\"11l\",\"both\")],\n \"ACC\": [(\"a24\",\"both\"),(\"p24\",\"both\")], \"mPFC\": [(\"9m\",\"both\")], \"FP1\": [(\"10d\",\"both\")], \"FP2\": [(\"10d\",\"both\")],\n \"IFG\": [(\"44\",\"both\"),(\"45a\",\"both\")], \"IFGorb\": [(\"47l\",\"both\")],\n \"STG\": [(\"STGa\",\"both\"),(\"STGr\",\"both\")], \"STS\": [(\"STSda\",\"both\")], \"MTG\": [(\"TE1a\",\"both\")],\n \"AG\": [(\"PFm\",\"both\"),(\"PGs\",\"both\")], \"PCC\": [(\"PCC\",\"both\")],\n \"mPFC_dmn\": [(\"9m\",\"both\"),(\"10r\",\"both\")],\n \"LP_L\": [(\"PFm\",\"left\"),(\"PGs\",\"left\")], \"LP_R\": [(\"PFm\",\"right\"),(\"PGs\",\"right\")],\n \"HPC_L\": [(\"Entorhinal\",\"left\")], \"HPC_R\": [(\"Entorhinal\",\"right\")],\n \"AI\": [(\"Ig\",\"both\"),(\"FOP4\",\"both\")], \"dACC\": [(\"a24\",\"both\")],\n \"sgACC\": [(\"s32\",\"both\")], \"vmPFC\": [(\"25\",\"both\")],\n \"Amygdala_L\": [], \"Amygdala_R\": [], \"Caudate_L\": [], \"Caudate_R\": [],\n \"Putamen_L\": [], \"Putamen_R\": [], \"Thalamus\": [],\n}\n\n@lru_cache(maxsize=1)\ndef _hcp_label_vertices(mesh: str = \"fsaverage5\") -> Dict[str, np.ndarray]:\n import mne, os\n if mesh != \"fsaverage5\":\n raise ValueError(\"Only fsaverage5 is supported\")\n max_v = FSAVERAGE5_VERTS_PER_HEMI\n\n # Fix 4/10: use fetch_fsaverage (~50 MB) instead of sample dataset (~1.7 GB)\n subjects_dir = Path(mne.datasets.fetch_fsaverage(verbose=False)).parent\n mne.datasets.fetch_hcp_mmp_parcellation(\n subjects_dir=subjects_dir, accept=True, verbose=False\n )\n\n # Fix 5: auto-detect annotation name (differs across MNE versions)\n annot_dir = subjects_dir / \"fsaverage\" / \"label\"\n annot_files = list(annot_dir.glob(\"lh.*.annot\"))\n if any(\"HCP-MMP1\" in f.name for f in annot_files):\n parc = \"HCP-MMP1\"\n else:\n parc = \"HCPMMP1\"\n\n out: Dict[str, List[np.ndarray]] = {}\n for hemi_code, offset in ((\"lh\", 0), (\"rh\", max_v)):\n labels = mne.read_labels_from_annot(\n \"fsaverage\", parc, hemi=hemi_code, subjects_dir=subjects_dir\n )\n for lab in labels:\n name = (lab.name\n .replace(\"_ROI-lh\",\"\").replace(\"_ROI-rh\",\"\")\n .replace(\"_ROI\",\"\").replace(\"-lh\",\"\").replace(\"-rh\",\"\"))\n verts = np.asarray(lab.vertices, dtype=np.int64)\n verts = verts[verts < max_v] + offset\n if verts.size:\n out.setdefault(name, []).append(verts)\n return {k: np.concatenate(v) for k, v in out.items()}\n\n\ndef get_hcp_roi_indices(rois, *, hemi=\"both\", mesh=\"fsaverage5\") -> np.ndarray:\n labels = _hcp_label_vertices(mesh)\n names = [rois] if isinstance(rois, str) else list(rois)\n selected = []\n for roi in names:\n if roi.endswith(\"*\"): selected.extend(k for k in labels if k.startswith(roi[:-1]))\n elif roi.startswith(\"*\"): selected.extend(k for k in labels if k.endswith(roi[1:]))\n elif roi in labels: selected.append(roi)\n else: raise ValueError(f\"ROI {roi!r} not found in HCP-MMP labels\")\n idx_parts = []\n for name in selected:\n verts = labels[name]\n if hemi == \"left\": idx_parts.append(verts[verts < FSAVERAGE5_VERTS_PER_HEMI])\n elif hemi == \"right\": idx_parts.append(verts[verts >= FSAVERAGE5_VERTS_PER_HEMI])\n else: idx_parts.append(verts)\n if not idx_parts:\n return np.array([], dtype=np.int64)\n # Fix 9: deduplicate vertices (wildcard matches can overlap)\n return np.unique(np.concatenate(idx_parts))\n\n\ndef roi_dict_to_fsaverage5(roi_values: Mapping[str, float]) -> Tuple[np.ndarray, List[str]]:\n acc = np.zeros(FSAVERAGE5_VERTICES, dtype=np.float64)\n counts = np.zeros(FSAVERAGE5_VERTICES, dtype=np.float64)\n skipped = []\n for roi_name, value in roi_values.items():\n targets = ROI56_TO_HCP.get(roi_name, [(roi_name, \"both\")])\n if not targets:\n skipped.append(roi_name); continue\n painted = False\n for hcp_name, hemi in targets:\n try:\n idx = get_hcp_roi_indices(hcp_name, hemi=hemi)\n except ValueError:\n continue\n if idx.size == 0:\n continue\n acc[idx] += float(value); counts[idx] += 1.0; painted = True\n if not painted:\n skipped.append(roi_name)\n mask = counts > 0\n acc[mask] /= counts[mask]\n return acc.astype(np.float32), skipped\n\n\ndef plot_brain(\n roi_map: Mapping[str, float],\n *,\n views: Sequence[str] = (\"left\", \"right\", \"dorsal\"),\n cmap: str = \"RdBu_r\",\n vmax: float | None = None,\n threshold: float = 0.02,\n title: str | None = None,\n):\n \"\"\"Paint roi_map on fsaverage5 and plot with Nilearn.\n\n views: any subset of (\"left\", \"right\", \"dorsal\", \"ventral\").\n MNE HCP-MMP1 parcellation downloads on first use (~50 MB).\n \"\"\"\n import matplotlib.pyplot as plt\n from nilearn.datasets import fetch_surf_fsaverage\n from nilearn.plotting import plot_surf_stat_map\n\n vertex_map, skipped = roi_dict_to_fsaverage5(roi_map)\n if skipped:\n print(f\" [brain plotter] skipped (no HCP mapping): {skipped}\")\n\n v = np.asarray(vertex_map, dtype=np.float32)\n # Fix 7: nanpercentile is safe when some vertices are zero/NaN\n if vmax is None:\n vmax = float(np.nanpercentile(np.abs(v), 99)) or 1.0\n\n fsa = fetch_surf_fsaverage(mesh=\"fsaverage5\")\n left = v[:FSAVERAGE5_VERTS_PER_HEMI]\n right = v[FSAVERAGE5_VERTS_PER_HEMI:]\n\n VIEW_SPEC = {\n \"left\": (\"left\", \"lateral\", \"infl_left\", \"sulc_left\"),\n \"right\": (\"right\", \"lateral\", \"infl_right\", \"sulc_right\"),\n \"dorsal\": (\"left\", \"dorsal\", \"infl_left\", \"sulc_left\"),\n \"ventral\": (\"left\", \"ventral\", \"infl_left\", \"sulc_left\"),\n }\n n = len(views)\n fig, axes = plt.subplots(1, n, figsize=(4.2*n, 3.8), subplot_kw={\"projection\": \"3d\"})\n if n == 1:\n axes = [axes]\n for ax, view in zip(axes, views):\n hemi, nv, infl_key, sulc_key = VIEW_SPEC.get(view, VIEW_SPEC[\"left\"])\n stat = left if hemi == \"left\" else right\n # Fix 3: surf_mesh first (nilearn keyword argument order)\n plot_surf_stat_map(\n surf_mesh=fsa[infl_key], stat_map=stat, bg_map=fsa[sulc_key],\n view=nv, axes=ax, cmap=cmap, vmin=-vmax, vmax=vmax,\n threshold=threshold, colorbar=False, symmetric_cbar=True,\n )\n ax.set_title(view, fontsize=9)\n sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=-vmax, vmax=vmax))\n sm.set_array([])\n # Fix 6: ax=list(axes) avoids matplotlib deprecation warning\n fig.colorbar(sm, ax=list(axes), shrink=0.55, label=\"Predicted activation (a.u.)\")\n if title:\n fig.suptitle(title, fontsize=11)\n plt.tight_layout()\n plt.show()\n return fig\n\nprint(\"Brain plotter ready.\")\nprint(\" Usage: plot_brain(roi_map, title='My stimulus')\")\nprint(\" HCP-MMP1 parcellation downloads on first call (~50 MB via fetch_fsaverage).\")\n"
|
| 140 |
},
|
| 141 |
{
|
| 142 |
"cell_type": "code",
|
|
|
|
| 151 |
"execution_count": null,
|
| 152 |
"metadata": {},
|
| 153 |
"outputs": [],
|
| 154 |
+
"source": "# ── Option A: pre-cached text examples (CPU, instant, no encoder needed) ─────\n# Uses examples_cache.npz bundled in the zip or downloaded from HuggingFace.\n# Skipped automatically if the file is unavailable.\n#\n# NOTE: The cache holds 2560-d Qwen3 embeddings → this cell always uses\n# text2roi_combined_v4.pt (in_dim=3840), regardless of the MODALITY setting above.\n\nif examples_path is None:\n print('examples_cache.npz not available — skip to Option B (audio) or Option C (text).')\nelse:\n # Load combined model (in_dim=3840) — required for Qwen3 2560-d embeddings\n _A_MODEL_FILE = 'text2roi_combined_v4.pt'\n _A_IN_DIM = 3840\n _a_ckpt_path = _resolve(_A_MODEL_FILE)\n _a_ckpt = torch.load(_a_ckpt_path, map_location='cpu', weights_only=False)\n _a_model = Text2ROI(in_dim=_a_ckpt['in_dim'], out_dim=_a_ckpt.get('n_roi', 56))\n _a_model.load_state_dict(_a_ckpt['state_dict'])\n _a_model.eval()\n print(f'Option A model: in_dim={_a_ckpt[\"in_dim\"]} val_R={_a_ckpt.get(\"best_val_r\",\"N/A\")}')\n\n cache = np.load(examples_path, allow_pickle=True)\n sentences = [s.decode() if isinstance(s, bytes) else str(s) for s in cache['sentences']]\n embeddings = cache['embeddings'].astype(np.float32) # (N, 2560)\n print(f'Loaded {len(sentences)} cached examples embedding_dim={embeddings.shape[1]}')\n\n # ── pick a sentence to visualise ──\n idx = 0 # change index or replace sentence below\n sentence = sentences[idx]\n roi_map = predict_from_features(embeddings[idx], _a_model, _A_IN_DIM)\n\n print(f'\\nInput: \"{sentence}\"')\n print('Top 10 ROIs:')\n for roi, val in sorted(roi_map.items(), key=lambda x: -x[1])[:10]:\n print(f' {roi:<16} {val:+.4f}')\n plot_rois(roi_map, title=sentence)\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 155 |
"id": "cached-examples"
|
| 156 |
},
|
| 157 |
{
|
|
|
|
| 200 |
"execution_count": null,
|
| 201 |
"metadata": {},
|
| 202 |
"outputs": [],
|
| 203 |
+
"source": "# ── Option C: Custom text inference with Qwen3 (GPU recommended, ~8 GB) ──────\n# NOTE: this cell always uses text2roi_combined_v4.pt (in_dim=3840), regardless\n# of the MODALITY setting above. The combined model was trained on both audio and\n# text; for text-only input the whisper slot is zero-padded automatically.\n# For best results on audio content, prefer Option B (actual audio file → Whisper).\nimport torch.nn.functional as F\nfrom transformers import AutoModel, AutoTokenizer\n\nYOUR_TEXT = 'listening to a symphony building to its climax'\n\n# ── Load combined text model (always in_dim=3840) ─────────────────────────────\n_TEXT_MODEL_FILE = 'text2roi_combined_v4.pt'\n_TEXT_IN_DIM = 3840\n\n_text_ckpt_path = _resolve(_TEXT_MODEL_FILE)\n_text_ckpt = torch.load(_text_ckpt_path, map_location='cpu', weights_only=False)\n_text_model = Text2ROI(in_dim=_text_ckpt['in_dim'], out_dim=_text_ckpt.get('n_roi', 56))\n_text_model.load_state_dict(_text_ckpt['state_dict'])\n_text_model.eval()\nprint(f'Text model loaded in_dim={_text_ckpt[\"in_dim\"]} val_R={_text_ckpt.get(\"best_val_r\",\"N/A\")}')\n\n# ── Qwen3 embedding (last-token + L2 norm, matches training extraction) ───────\nprint('Loading Qwen3-Embedding-4B (downloads ~8 GB on first run)...')\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\ntok = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-4B', padding_side='left')\nqwen = AutoModel.from_pretrained(\n 'Qwen/Qwen3-Embedding-4B',\n torch_dtype=torch.bfloat16 if device != 'cpu' else torch.float32\n).to(device).eval()\n\nwith torch.no_grad():\n enc = tok([YOUR_TEXT], return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)\n h = qwen(**enc).last_hidden_state[:, -1].float() # last token (decoder-style)\n emb = F.normalize(h, p=2, dim=1).cpu().numpy()[0].astype(np.float32) # (2560,)\n\ndel qwen; torch.cuda.empty_cache() if device == 'cuda' else None\nprint(f'Text embedding: {emb.shape} norm={float(np.linalg.norm(emb)):.4f}')\n\n# ── Predict using the combined model ─────────────────────────────────────────\nroi_map_text = predict_from_features(emb, _text_model, _TEXT_IN_DIM)\nprint(f'\\nInput: \"{YOUR_TEXT}\"')\nprint('Top 10 ROIs:')\nfor roi, val in sorted(roi_map_text.items(), key=lambda x: -x[1])[:10]:\n print(f' {roi:<16} {val:+.4f}')\nplot_rois(roi_map_text, title=YOUR_TEXT)\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 204 |
"id": "custom-text"
|
| 205 |
},
|
| 206 |
{
|