{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" }, "colab": { "provenance": [], "gpuType": "T4" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Amphora NeuroText v4 — Audio & Text → Brain ROI Activation\n", "\n", "Predict which brain regions activate in response to any audio or text stimulus.\n", "\n", "Trained on **real naturalistic fMRI data** from 1,600+ subjects across 4,480 sessions. Zero-shot: no brain scan needed at inference.\n", "\n", "**Audio model (Whisper v4) beats TRIBE v2** (Meta AI, Algonauts 2025 winner) **by +4.2%** \n", "Holdout R = 0.257 vs TRIBE 0.215 · 23 held-out subjects · single shared model (no per-subject fine-tuning)\n", "\n", "**All v4 models use honest evaluation:** per-subject z-scoring + per-subject train/val split (no data leakage).\n", "\n", "**Runtime:** GPU recommended (T4 or better). The MLP projector runs on CPU instantly — only the encoder models (Whisper/Qwen3) need GPU.\n", "\n", "---\n", "Model repo: [ffh92r32rm0/Amphora_NeuroText](https://huggingface.co/ffh92r32rm0/Amphora_NeuroText)" ], "id": "markdown-intro" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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", "id": "install" }, { "cell_type": "code", "execution_count": null, "id": "hf-auth", "metadata": {}, "outputs": [], "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" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Choose modality ───────────────────────────────────────────────────────────\n", "# Set MODALITY to 'audio' or 'text'\n", "MODALITY = 'audio' # 'audio' = Whisper v4 (holdout R=0.257, beats TRIBE v2 by +4.2%)\n", " # 'text' = Combined v4 (val R=0.192, zero-padded whisper slot)\n", "\n", "MODEL_FILE = {'audio': 'text2roi_whisper_v4.pt', 'text': 'text2roi_combined_v4.pt'}[MODALITY]\n", "IN_DIM = {'audio': 1280, 'text': 3840 }[MODALITY]\n", "\n", "print(f'Modality : {MODALITY}')\n", "print(f'Model : {MODEL_FILE} (in_dim={IN_DIM})')" ], "id": "choose-model" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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", "id": "resolve-files" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Load model (tiny MLP, runs instantly on CPU) ──────────────────────────────\n", "import torch, torch.nn as nn, numpy as np\n", "\n", "class Text2ROI(nn.Module):\n", " def __init__(self, in_dim=1280, hidden=1024, out_dim=56, dropout=0.1):\n", " super().__init__()\n", " self.net = nn.Sequential(\n", " nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.LayerNorm(hidden),\n", " nn.Linear(hidden, hidden//2), nn.GELU(), nn.Dropout(dropout),\n", " nn.Linear(hidden//2, out_dim),\n", " )\n", " def forward(self, x): return self.net(x)\n", "\n", "ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)\n", "model = Text2ROI(in_dim=ckpt['in_dim'], out_dim=ckpt.get('n_roi', 56))\n", "model.load_state_dict(ckpt['state_dict'])\n", "model.eval()\n", "\n", "ROI_NAMES = [\n", " 'V1','V2','V3','V4','V3A','V3B','LO1','LO2','MT','MST','V7','IPS1',\n", " 'FFA-1','FFA-2','PPA','RSC','OFA','EBA','IPS2','IPS3','IPS4','IPS5','SPL1',\n", " 'hIP1','hIP2','hIP3','dlPFC','vlPFC','OFC','ACC','mPFC','FP1','FP2',\n", " 'IFG','IFGorb','STG','STS','MTG','AG','PCC','mPFC_dmn','LP_L','LP_R',\n", " 'HPC_L','HPC_R','AI','dACC','sgACC','vmPFC',\n", " 'Amygdala_L','Amygdala_R','Caudate_L','Caudate_R','Putamen_L','Putamen_R','Thalamus',\n", "]\n", "\n", "val_r = ckpt.get('best_val_r', 'N/A')\n", "print(f'Model loaded in_dim={ckpt[\"in_dim\"]} val_R={val_r}')\n", "print(f'Parameters: {sum(p.numel() for p in model.parameters()):,}')" ], "id": "load-model" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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", "id": "helpers" }, { "cell_type": "code", "execution_count": null, "id": "brain-plotter-lib", "metadata": {}, "outputs": [], "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" }, { "cell_type": "code", "execution_count": null, "id": "brain-plotter-usage", "metadata": {}, "outputs": [], "source": "# ── Plot predictions on the brain ─────────────────────────────────────────────\n# Run any inference cell above first (Option A, B, or C), then run this cell.\n\n_map = (\n roi_map if 'roi_map' in dir() else\n roi_map_audio if 'roi_map_audio' in dir() else\n roi_map_text if 'roi_map_text' in dir() else\n None\n)\n\nif _map is None:\n print(\"Run an inference cell first (Option A, B, or C) to produce a roi_map.\")\nelse:\n # 4-view: left/right lateral + dorsal + ventral\n plot_brain(\n _map,\n views=(\"left\", \"right\", \"dorsal\", \"ventral\"),\n cmap=\"RdBu_r\",\n title=\"Predicted brain activation (fsaverage5, HCP-MMP1)\",\n )\n\n# ---- Compare two stimuli side-by-side ----------------------------------------\n# fig1 = plot_brain(roi_map_audio, title=\"Audio stimulus\")\n# fig2 = plot_brain(roi_map_text, title=\"Text stimulus\")\n\n# ---- Custom views (e.g. medial-only) -----------------------------------------\n# plot_brain(_map, views=(\"dorsal\", \"ventral\"), cmap=\"hot\", vmax=0.5)\n" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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", "id": "cached-examples" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Option B: Audio inference with Whisper v4 (GPU recommended) ──────────────\n", "# Upload a .wav file to Colab, or set AUDIO_PATH to a URL.\n", "# Whisper-large-v3 downloads automatically (~3 GB, cached after first run).\n", "\n", "AUDIO_PATH = 'your_audio.wav' # <-- replace with your file\n", "\n", "# Uncomment to download a sample clip:\n", "# import urllib.request\n", "# urllib.request.urlretrieve('https://upload.wikimedia.org/wikipedia/commons/2/22/Beethoven_Moonlight_Sonata_First_Movement.ogg', 'sample.ogg')\n", "# AUDIO_PATH = 'sample.ogg'\n", "\n", "import librosa\n", "from transformers import WhisperProcessor, WhisperModel\n", "\n", "print('Loading Whisper-large-v3 (downloads ~3 GB on first run)...')\n", "processor = WhisperProcessor.from_pretrained('openai/whisper-large-v3')\n", "whisper_enc = WhisperModel.from_pretrained('openai/whisper-large-v3')\n", "whisper_enc.eval()\n", "print('Loaded.')\n", "\n", "audio, sr = librosa.load(AUDIO_PATH, sr=16000, mono=True)\n", "inputs = processor(audio, sampling_rate=16000, return_tensors='pt')\n", "with torch.no_grad():\n", " enc_out = whisper_enc.encoder(inputs.input_features)\n", "audio_emb = enc_out.last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy() # (1280,)\n", "print(f'Audio embedding: {audio_emb.shape}')\n", "\n", "roi_map_audio = predict_from_features(audio_emb, model, IN_DIM)\n", "print('\\nTop 10 ROIs (audio):')\n", "for roi, val in sorted(roi_map_audio.items(), key=lambda x: -x[1])[:10]:\n", " print(f' {roi:<16} {val:+.4f}')\n", "plot_rois(roi_map_audio, title=f'Audio: {AUDIO_PATH}')" ], "id": "audio-inference" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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", "id": "custom-text" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Network-level summary ─────────────────────────────────────────────────────\n", "NETWORKS = {\n", " 'Visual': ['V1','V2','V3','V4','V3A','V3B','LO1','LO2','MT','MST','V7','IPS1','FFA-1','FFA-2','PPA','RSC','OFA','EBA'],\n", " 'Parietal': ['IPS2','IPS3','IPS4','IPS5','SPL1','hIP1','hIP2','hIP3'],\n", " 'Frontal': ['dlPFC','vlPFC','OFC','ACC','mPFC','FP1','FP2'],\n", " 'Language': ['IFG','IFGorb','STG','STS','MTG','AG'],\n", " 'DefaultMode': ['PCC','mPFC_dmn','LP_L','LP_R','HPC_L','HPC_R'],\n", " 'Salience': ['AI','dACC','sgACC','vmPFC','Amygdala_L','Amygdala_R'],\n", " 'Subcortical': ['Caudate_L','Caudate_R','Putamen_L','Putamen_R','Thalamus'],\n", "}\n", "\n", "# swap roi_map for whichever you ran above\n", "active_map = roi_map if 'roi_map' in dir() else (roi_map_audio if 'roi_map_audio' in dir() else roi_map_text)\n", "\n", "nets = {net: float(np.mean([active_map[r] for r in rois if r in active_map]))\n", " for net, rois in NETWORKS.items()}\n", "\n", "print('Network activations:')\n", "for net, val in sorted(nets.items(), key=lambda x: -x[1]):\n", " bar = '\\u2588' * max(0, int((val + 0.3) * 18))\n", " print(f' {net:<14} {val:+.4f} {bar}')" ], "id": "network-summary" } ] }