""" predict.py — Amphora NeuroText inference for all three models. Models ------ text2roi_projector.pt Text → 56 ROIs (Qwen3-Embedding-4B, val R=0.212) text2roi_whisper.pt Audio → 56 ROIs (Whisper-large-v3, val R=0.413) text2roi_dual_v2.pt Text+Audio → 56 ROIs (dual-tower, val R=0.170) Quick start ----------- from predict import predict_text, predict_audio, predict_combined, ROI_NAMES_56 roi_map = predict_text("I am terrified of the dark", "text2roi_projector.pt") print(sorted(roi_map.items(), key=lambda x: -x[1])[:5]) roi_map = predict_audio("clip.wav", "text2roi_whisper.pt") roi_map = predict_combined("narration text", "clip.wav", "text2roi_dual_v2.pt") """ from __future__ import annotations import sys from pathlib import Path from typing import Dict, List, Optional import numpy as np import torch import torch.nn as nn # ── Canonical 56-ROI schema ─────────────────────────────────────────────────── ROI_NAMES_56: List[str] = [ "V1", "V2", "V3", "V4", "V3A", "V3B", "LO1", "LO2", "MT", "MST", "V7", "IPS1", "FFA-1", "FFA-2", "PPA", "RSC", "OFA", "EBA", "IPS2", "IPS3", "IPS4", "IPS5", "SPL1", "hIP1", "hIP2", "hIP3", "dlPFC", "vlPFC", "OFC", "ACC", "mPFC", "FP1", "FP2", "IFG", "IFGorb", "STG", "STS", "MTG", "AG", "PCC", "mPFC_dmn", "LP_L", "LP_R", "HPC_L", "HPC_R", "AI", "dACC", "sgACC", "vmPFC", "Amygdala_L", "Amygdala_R", "Caudate_L", "Caudate_R", "Putamen_L", "Putamen_R", "Thalamus", ] # ── Shared MLP architecture (qwen3 + whisper projectors) ───────────────────── class Text2ROI(nn.Module): def __init__(self, in_dim: int = 2560, hidden: int = 1024, out_dim: int = 56, dropout: float = 0.1): super().__init__() self.net = nn.Sequential( nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.LayerNorm(hidden), nn.Linear(hidden, hidden // 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden // 2, out_dim), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) # ── Dual-tower architecture ─────────────────────────────────────────────────── _WHISPER_ROI_IDX = [33, 34, 35, 36, 45, 46] _QWEN3_ROI_IDX = [14, 15, 30, 37, 38, 39, 40, 41, 42, 43, 44, 47, 48] _SHARED_ROI_IDX = [i for i in range(56) if i not in _WHISPER_ROI_IDX and i not in _QWEN3_ROI_IDX] class SpecializedDualTower(nn.Module): """ Dual-tower model with modality-specialized heads. whisper_tower : 1280 → tower_dim → 6 auditory/salience ROIs qwen3_tower : 2560 → tower_dim → 13 semantic/DMN ROIs fusion_head : tower_dim×2 → hidden → 37 shared ROIs """ def __init__(self, n_roi: int = 56, tower_dim: int = 512, hidden: int = 512, dropout: float = 0.15): super().__init__() self.whisper_tower = nn.Sequential( nn.Linear(1280, tower_dim), nn.GELU(), nn.LayerNorm(tower_dim)) self.qwen3_tower = nn.Sequential( nn.Linear(2560, tower_dim), nn.GELU(), nn.LayerNorm(tower_dim)) self.whisper_head = nn.Linear(tower_dim, len(_WHISPER_ROI_IDX)) self.qwen3_head = nn.Linear(tower_dim, len(_QWEN3_ROI_IDX)) self.fusion_head = nn.Sequential( nn.Linear(tower_dim * 2, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, len(_SHARED_ROI_IDX))) self.n_roi = n_roi self.register_buffer('w_idx', torch.tensor(_WHISPER_ROI_IDX)) self.register_buffer('q_idx', torch.tensor(_QWEN3_ROI_IDX)) self.register_buffer('s_idx', torch.tensor(_SHARED_ROI_IDX)) def forward(self, whisper_emb: torch.Tensor, qwen3_emb: torch.Tensor) -> torch.Tensor: wt = self.whisper_tower(whisper_emb) qt = self.qwen3_tower(qwen3_emb) out = torch.zeros(whisper_emb.shape[0], self.n_roi, device=whisper_emb.device) out[:, self.w_idx] = self.whisper_head(wt) out[:, self.q_idx] = self.qwen3_head(qt) out[:, self.s_idx] = self.fusion_head(torch.cat([wt, qt], dim=-1)) return out # ── Checkpoint loaders ──────────────────────────────────────────────────────── def load_projector(checkpoint_path: str, device: torch.device) -> tuple[Text2ROI, List[str], dict]: ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) roi_names = [s.decode() if isinstance(s, bytes) else str(s) for s in ckpt.get("roi_names", ROI_NAMES_56)] in_dim = ckpt.get("in_dim", 2560) hidden = ckpt.get("hidden", 1024) n_roi = ckpt.get("n_roi", len(roi_names)) dropout = ckpt.get("args", {}).get("dropout", 0.1) model = Text2ROI(in_dim=in_dim, hidden=hidden, out_dim=n_roi, dropout=dropout).to(device) model.load_state_dict(ckpt["state_dict"]) model.eval() return model, roi_names, ckpt def load_dual_tower(checkpoint_path: str, device: torch.device) -> tuple[SpecializedDualTower, List[str], dict]: ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) roi_names = [s.decode() if isinstance(s, bytes) else str(s) for s in ckpt.get("roi_names", ROI_NAMES_56)] args = ckpt.get("args", {}) model = SpecializedDualTower( n_roi=len(roi_names), tower_dim=args.get("tower_dim", 512), hidden=args.get("hidden", 512), dropout=args.get("dropout", 0.15), ).to(device) model.load_state_dict(ckpt["state_dict"]) model.eval() return model, roi_names, ckpt # ── Embedding helpers ───────────────────────────────────────────────────────── def _embed_qwen3(texts: List[str], device: str = "cuda") -> np.ndarray: from transformers import AutoTokenizer, AutoModel import torch.nn.functional as F tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-4B", padding_side="left") dtype = torch.bfloat16 if device != "cpu" else torch.float32 model = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-4B", torch_dtype=dtype).to(device).eval() results = [] with torch.no_grad(): for i in range(0, len(texts), 8): batch = texts[i:i+8] enc = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device) h = model(**enc).last_hidden_state[:, -1].float() results.append(F.normalize(h, p=2, dim=1).cpu().numpy()) return np.concatenate(results).astype(np.float32) def _embed_whisper(audio_paths: List[str], device: str = "cuda") -> np.ndarray: import librosa from transformers import WhisperProcessor, WhisperModel processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3") dtype = torch.bfloat16 if device != "cpu" else torch.float32 model = WhisperModel.from_pretrained("openai/whisper-large-v3", torch_dtype=dtype).to(device).eval() results = [] with torch.no_grad(): for path in audio_paths: wav, sr = librosa.load(path, sr=16000, mono=True) inp = processor(wav, sampling_rate=16000, return_tensors="pt").input_features.to(device) enc = model.encoder(inp.to(dtype)).last_hidden_state emb = enc.mean(dim=1).float().cpu().numpy() results.append(emb) return np.concatenate(results).astype(np.float32) # ── High-level inference API ────────────────────────────────────────────────── def predict_text(text: str, checkpoint_path: str, device: Optional[str] = None) -> Dict[str, float]: """Text → 56 ROI activation scores using the qwen3 projector.""" if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" dev = torch.device(device) model, roi_names, _ = load_projector(checkpoint_path, dev) emb = _embed_qwen3([text], device) with torch.no_grad(): pred = model(torch.from_numpy(emb).to(dev)).cpu().numpy()[0] return dict(zip(roi_names, pred.tolist())) def predict_audio(audio_path: str, checkpoint_path: str, device: Optional[str] = None) -> Dict[str, float]: """Audio file → 56 ROI activation scores using the whisper projector.""" if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" dev = torch.device(device) model, roi_names, _ = load_projector(checkpoint_path, dev) emb = _embed_whisper([audio_path], device) with torch.no_grad(): pred = model(torch.from_numpy(emb).to(dev)).cpu().numpy()[0] return dict(zip(roi_names, pred.tolist())) def predict_combined(text: str, audio_path: str, checkpoint_path: str, device: Optional[str] = None) -> Dict[str, float]: """Text + audio → 56 ROI activation scores using the dual-tower model.""" if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" dev = torch.device(device) model, roi_names, _ = load_dual_tower(checkpoint_path, dev) q_emb = torch.from_numpy(_embed_qwen3([text], device)).to(dev) w_emb = torch.from_numpy(_embed_whisper([audio_path], device)).to(dev) with torch.no_grad(): pred = model(w_emb, q_emb).cpu().numpy()[0] return dict(zip(roi_names, pred.tolist())) def top_rois(roi_map: Dict[str, float], n: int = 10) -> List[tuple]: return sorted(roi_map.items(), key=lambda x: -x[1])[:n] # ── CLI ─────────────────────────────────────────────────────────────────────── if __name__ == "__main__": import argparse ap = argparse.ArgumentParser(description="Amphora NeuroText — predict ROI activations") ap.add_argument("text", help="Input text stimulus") ap.add_argument("--checkpoint", default="text2roi_projector.pt") ap.add_argument("--audio", default=None, help="Audio file (for whisper/dual models)") ap.add_argument("--top", type=int, default=10) ap.add_argument("--device", default="auto") args = ap.parse_args() device = ("cuda" if torch.cuda.is_available() else "cpu") if args.device == "auto" else args.device if args.audio and "dual" in args.checkpoint: roi_map = predict_combined(args.text, args.audio, args.checkpoint, device) mode = "dual-tower (text+audio)" elif args.audio: roi_map = predict_audio(args.audio, args.checkpoint, device) mode = "whisper (audio)" else: roi_map = predict_text(args.text, args.checkpoint, device) mode = "qwen3 (text)" print(f"\nModel: {mode}") print(f"Input: {args.text[:80]!r}") print(f"\nTop {args.top} ROIs:") for roi, score in top_rois(roi_map, args.top): bar = "█" * max(0, int((score + 1) * 15)) print(f" {roi:<16} {score:+.3f} {bar}")