| """
|
| predict.py — Amphora NeuroText inference
|
|
|
| Models (honest cross-subject holdout val_R — no data leakage)
|
| ------
|
| text2roi_whisper_v4.pt Audio → 56 ROIs (Whisper-large-v3, 1280d, val R=0.217, holdout R=0.257)
|
| text2roi_combined_v4.pt Text+Audio → 56 ROIs (whisper|qwen3, 3840d, val R=0.192)
|
| text2roi_qwen3_v8.pt Text → 56 ROIs (Qwen3-Embedding-4B, 2560d, val R=0.115)
|
|
|
| All v4 models use per-subject z-scoring and a per-subject train/val split.
|
| Previous models (v2/v3) had inflated val_R from within-subject splits — do not compare directly.
|
|
|
| Whisper v4 beats TRIBE v2 (Meta AI, Algonauts 2025 winner) by +4.2% on a 23-subject holdout.
|
| Single shared model — no per-subject fine-tuning required.
|
|
|
| Quick start
|
| -----------
|
| from predict import predict_text, predict_audio, top_rois
|
|
|
| # Audio → brain regions (recommended — strongest model)
|
| roi_map = predict_audio("clip.wav")
|
| print(top_rois(roi_map, n=5))
|
| # → [('ACC', 0.44), ('STG', 0.42), ('Thalamus', 0.41), ...]
|
|
|
| # Text → brain regions
|
| roi_map = predict_text("I am terrified of the dark")
|
| print(top_rois(roi_map, n=5))
|
| # → [('Amygdala_L', 0.xx), ('AI', 0.xx), ('dACC', 0.xx), ...]
|
|
|
| # Text + Audio → brain regions
|
| roi_map = predict_combined("narration text", "clip.wav")
|
| print(top_rois(roi_map, n=5))
|
| """
|
| 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
|
|
|
|
|
| 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",
|
| ]
|
|
|
| DEFAULT_CHECKPOINT = "text2roi_whisper_v4.pt"
|
|
|
|
|
| class Text2ROI(nn.Module):
|
| def __init__(self, in_dim: int = 1280, 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)
|
|
|
|
|
|
|
| _model_cache: Dict[str, tuple] = {}
|
|
|
|
|
| def _load_model(checkpoint: str) -> tuple:
|
| if checkpoint in _model_cache:
|
| return _model_cache[checkpoint]
|
|
|
|
|
| ckpt_path = Path(checkpoint)
|
|
|
| if not ckpt_path.exists():
|
| ckpt_path = Path(__file__).parent / checkpoint
|
|
|
| if not ckpt_path.exists():
|
| try:
|
| from huggingface_hub import hf_hub_download
|
| ckpt_path = Path(hf_hub_download("ffh92r32rm0/Amphora_NeuroText", checkpoint))
|
| except Exception as e:
|
| raise FileNotFoundError(
|
| f"Checkpoint '{checkpoint}' not found locally or on HuggingFace.\n"
|
| f"Make sure the .pt file is in the same folder as predict.py.\n"
|
| f"Original error: {e}"
|
| )
|
|
|
| state = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
|
| in_dim = state.get("in_dim", 1280)
|
| n_roi = state.get("n_roi", 56)
|
| model = Text2ROI(in_dim=in_dim, out_dim=n_roi)
|
| model.load_state_dict(state["state_dict"])
|
| model.eval()
|
| _model_cache[checkpoint] = (model, in_dim)
|
| return model, in_dim
|
|
|
|
|
|
|
| _qwen3_model = None
|
| _qwen3_tokenizer = None
|
|
|
|
|
| def _embed_text_qwen3(text: str) -> np.ndarray:
|
| global _qwen3_model, _qwen3_tokenizer
|
| if _qwen3_model is None:
|
| from transformers import AutoTokenizer, AutoModel
|
| _qwen3_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-4B")
|
| _qwen3_model = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-4B")
|
| _qwen3_model.eval()
|
| inputs = _qwen3_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
|
| with torch.no_grad():
|
| out = _qwen3_model(**inputs)
|
| emb = out.last_hidden_state[:, 0, :].squeeze(0).cpu().numpy()
|
| return emb.astype(np.float32)
|
|
|
|
|
|
|
| _whisper_model = None
|
| _whisper_processor = None
|
|
|
|
|
| def _embed_audio_whisper(audio_path: str) -> np.ndarray:
|
| global _whisper_model, _whisper_processor
|
| if _whisper_model is None:
|
| from transformers import WhisperProcessor, WhisperModel
|
| _whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
|
| _whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3")
|
| _whisper_model.eval()
|
| import librosa
|
| audio, sr = librosa.load(audio_path, sr=16000, mono=True)
|
| inputs = _whisper_processor(audio, sampling_rate=16000, return_tensors="pt")
|
| with torch.no_grad():
|
| enc = _whisper_model.encoder(inputs.input_features)
|
| emb = enc.last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy()
|
| return emb.astype(np.float32)
|
|
|
|
|
|
|
|
|
| def predict_audio(
|
| audio_path: str,
|
| checkpoint: str = "text2roi_whisper_v4.pt",
|
| ) -> Dict[str, float]:
|
| """Predict 56 brain ROI activations from an audio file.
|
|
|
| Args:
|
| audio_path: Path to audio file (wav, mp3, flac, …).
|
| checkpoint: Model file. Default: text2roi_whisper_v4.pt (val R=0.217, holdout R=0.257).
|
|
|
| Returns:
|
| dict mapping ROI name → predicted activation (z-scored units).
|
| """
|
| model, in_dim = _load_model(checkpoint)
|
| emb = _embed_audio_whisper(audio_path)
|
| if emb.shape[0] != in_dim:
|
| raise ValueError(f"Audio embedding dim {emb.shape[0]} != model in_dim {in_dim}")
|
| feat = torch.from_numpy(emb).unsqueeze(0)
|
| with torch.no_grad():
|
| pred = model(feat).squeeze(0).numpy()
|
| return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}
|
|
|
|
|
| def predict_text(
|
| text: str,
|
| checkpoint: str = "text2roi_combined_v4.pt",
|
| ) -> Dict[str, float]:
|
| """Predict 56 brain ROI activations from a text string.
|
|
|
| Uses text2roi_combined_v4.pt by default: qwen3 embedding zero-padded into
|
| the whisper slot (positions 0:1280 = zeros, 1280:3840 = qwen3 2560d).
|
| The combined model was trained with modality dropout so text-only works.
|
|
|
| Args:
|
| text: Input text string.
|
| checkpoint: Model file. Default: text2roi_combined_v4.pt (val R=0.192).
|
|
|
| Returns:
|
| dict mapping ROI name → predicted activation (z-scored units).
|
| """
|
| model, in_dim = _load_model(checkpoint)
|
| emb = _embed_text_qwen3(text)
|
|
|
| if in_dim == 3840:
|
|
|
| feat_np = np.concatenate([np.zeros(1280, dtype=np.float32), emb])
|
| elif in_dim == 2560:
|
| feat_np = emb
|
| else:
|
| raise ValueError(f"Unexpected model in_dim {in_dim} for text inference")
|
|
|
| feat = torch.from_numpy(feat_np).unsqueeze(0)
|
| with torch.no_grad():
|
| pred = model(feat).squeeze(0).numpy()
|
| return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}
|
|
|
|
|
| def predict_combined(
|
| text: str,
|
| audio_path: str,
|
| checkpoint: str = "text2roi_combined_v4.pt",
|
| ) -> Dict[str, float]:
|
| """Predict 56 brain ROI activations from both text and audio.
|
|
|
| Args:
|
| text: Text string.
|
| audio_path: Path to audio file.
|
| checkpoint: Must be the combined model (in_dim=3840).
|
|
|
| Returns:
|
| dict mapping ROI name → predicted activation (z-scored units).
|
| """
|
| model, in_dim = _load_model(checkpoint)
|
| if in_dim != 3840:
|
| raise ValueError("predict_combined requires the combined model (in_dim=3840).")
|
| w_emb = _embed_audio_whisper(audio_path)
|
| q_emb = _embed_text_qwen3(text)
|
| feat_np = np.concatenate([w_emb, q_emb]).astype(np.float32)
|
| feat = torch.from_numpy(feat_np).unsqueeze(0)
|
| with torch.no_grad():
|
| pred = model(feat).squeeze(0).numpy()
|
| return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}
|
|
|
|
|
|
|
|
|
| def top_rois(roi_map: Dict[str, float], n: int = 10) -> List[tuple]:
|
| """Return top-n ROIs sorted by predicted activation."""
|
| return sorted(roi_map.items(), key=lambda x: -x[1])[:n]
|
|
|
|
|
| def network_summary(roi_map: Dict[str, float]) -> Dict[str, float]:
|
| """Return mean activation per brain network."""
|
| NETWORKS = {
|
| "Visual": ["V1","V2","V3","V4","V3A","V3B","LO1","LO2","MT","MST","V7","IPS1","FFA-1","FFA-2","PPA","RSC","OFA","EBA"],
|
| "Parietal": ["IPS2","IPS3","IPS4","IPS5","SPL1","hIP1","hIP2","hIP3"],
|
| "Frontal": ["dlPFC","vlPFC","OFC","ACC","mPFC","FP1","FP2"],
|
| "Language": ["IFG","IFGorb","STG","STS","MTG","AG"],
|
| "DefaultMode": ["PCC","mPFC_dmn","LP_L","LP_R","HPC_L","HPC_R"],
|
| "Salience": ["AI","dACC","sgACC","vmPFC","Amygdala_L","Amygdala_R"],
|
| "Subcortical": ["Caudate_L","Caudate_R","Putamen_L","Putamen_R","Thalamus"],
|
| }
|
| return {
|
| net: float(np.mean([roi_map[r] for r in rois if r in roi_map]))
|
| for net, rois in NETWORKS.items()
|
| }
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| import argparse, json
|
|
|
| parser = argparse.ArgumentParser(description="Amphora NeuroText — brain ROI prediction")
|
| sub = parser.add_subparsers(dest="cmd")
|
|
|
| p_audio = sub.add_parser("audio", help="Audio file → brain ROIs")
|
| p_audio.add_argument("audio", help="Path to audio file")
|
| p_audio.add_argument("--model", default="text2roi_whisper_v4.pt")
|
| p_audio.add_argument("--top", type=int, default=10)
|
|
|
| p_text = sub.add_parser("text", help="Text → brain ROIs")
|
| p_text.add_argument("text", help="Input text string")
|
| p_text.add_argument("--model", default="text2roi_combined_v4.pt")
|
| p_text.add_argument("--top", type=int, default=10)
|
|
|
| p_combo = sub.add_parser("combined", help="Text + Audio → brain ROIs")
|
| p_combo.add_argument("text")
|
| p_combo.add_argument("audio")
|
| p_combo.add_argument("--model", default="text2roi_combined_v4.pt")
|
| p_combo.add_argument("--top", type=int, default=10)
|
|
|
| args = parser.parse_args()
|
|
|
| if args.cmd == "audio":
|
| result = predict_audio(args.audio, args.model)
|
| elif args.cmd == "text":
|
| result = predict_text(args.text, args.model)
|
| elif args.cmd == "combined":
|
| result = predict_combined(args.text, args.audio, args.model)
|
| else:
|
| parser.print_help()
|
| sys.exit(0)
|
|
|
| tops = top_rois(result, args.top)
|
| nets = network_summary(result)
|
| print(f"\nTop {args.top} ROIs:")
|
| for roi, val in tops:
|
| print(f" {roi:<16} {val:+.4f}")
|
| print("\nNetwork summary:")
|
| for net, val in sorted(nets.items(), key=lambda x: -x[1]):
|
| print(f" {net:<16} {val:+.4f}")
|
|
|
