predict.py

"""
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}")