scripts/foundation_models.py

"""Official REVE / CSBrain adapters; real electrode subsets, no virtual channels."""
import os,sys,types,importlib.util,ast,json
from pathlib import Path
os.environ.setdefault('HF_HOME',str(Path('.cache/huggingface').resolve()))
import torch
from torch import nn


def csbrain_class():
 # Namespaced import preserves upstream source and avoids CBraMod's `models` collision.
 root=Path('external/CSBrain/models').resolve();package='_setu_csbrain'
 if package+'.CSBrain' in sys.modules:return sys.modules[package+'.CSBrain'].CSBrain
 pkg=types.ModuleType(package);pkg.__path__=[str(root)];sys.modules[package]=pkg
 for name in ['CSBrain_transformerlayer','CSBrain_transformer','CSBrain']:
  p=root/(name+'.py');code=p.read_text().replace('from models.CSBrain_',f'from {package}.CSBrain_')
  mod=types.ModuleType(package+'.'+name);mod.__file__=str(p);mod.__package__=package;sys.modules[mod.__name__]=mod;exec(compile(code,str(p),'exec'),mod.__dict__)
 return sys.modules[package+'.CSBrain'].CSBrain


def cs_topology(names):
 tree=ast.parse(Path('external/CSBrain/models/model_for_speech.py').read_text());values={}
 for node in ast.walk(tree):
  if isinstance(node,ast.Assign) and len(node.targets)==1 and isinstance(node.targets[0],ast.Name) and node.targets[0].id in ['brain_regions','electrode_labels','topology']:
   values[node.targets[0].id]=ast.literal_eval(node.value)
 mapping=dict(zip(values['electrode_labels'],values['brain_regions']));regions=[mapping[n] for n in names]
 order=sorted(range(len(names)),key=lambda i:(regions[i],values['topology'][regions[i]].index(names[i])))
 return regions,order


class Foundation(nn.Module):
 def __init__(self,kind,names,pretrained=True,n_times=400):
  super().__init__();self.kind=kind;self.names=list(names);self.n_times=n_times
  if kind=='reve':
   from transformers import AutoModel,AutoConfig
   directory='models/reve-base'
   if pretrained:self.backbone,info=AutoModel.from_pretrained(directory,trust_remote_code=True,local_files_only=True,output_loading_info=True)
   else:self.backbone=AutoModel.from_config(AutoConfig.from_pretrained(directory,trust_remote_code=True,local_files_only=True),trust_remote_code=True);info={}
   assert not info.get('missing_keys') and not info.get('unexpected_keys') and not info.get('mismatched_keys'),str(info)
   bank=AutoModel.from_pretrained('models/reve-positions',trust_remote_code=True,local_files_only=True)
   positions=bank(names).detach().float();assert positions.shape==(len(names),3) and torch.isfinite(positions).all()
   self.register_buffer('positions',positions);self.dimension=len(names)*(1+(n_times-200)//180)*512
  elif kind=='csbrain':
   regions,order=cs_topology(names);self.backbone=csbrain_class()(brain_regions=regions,sorted_indices=order)
   if pretrained:
    state=torch.load('models/csbrain_pretrained.pth',map_location='cpu',weights_only=True)
    state={k.removeprefix('module.').removeprefix('backbone.'):v for k,v in state.items()}
    self.backbone.load_state_dict(state,strict=True)
   self.backbone.proj_out=nn.Identity();self.dimension=len(names)*(n_times//200)*200
  elif kind=='cbramod':
   from cbramod_probe import backbone
   self.backbone=backbone();self.dimension=len(names)*(n_times//200)*200
  else:raise ValueError(kind)
 def forward(self,x):
  if self.kind=='reve':return self.backbone(x,self.positions.unsqueeze(0).expand(len(x),-1,-1)).flatten(1)
  # FFT in official CSBrain requires float32; incoming samples remain float32 under autocast.
  return self.backbone(x.reshape(len(x),len(self.names),self.n_times//200,200)).flatten(1)

class FoundationClassifier(nn.Module):
 def __init__(self,kind,names,mode='partial',pretrained=True):
  super().__init__();self.encoder=Foundation(kind,names,pretrained);self.head=nn.Sequential(nn.LayerNorm(self.encoder.dimension),nn.Dropout(.5),nn.Linear(self.encoder.dimension,5))
  if mode!='full':
   for p in self.encoder.parameters():p.requires_grad=False
   if mode=='partial':
    layers=self.encoder.backbone.transformer.layers if kind=='reve' else self.encoder.backbone.encoder.layers
    for p in layers[-2:].parameters():p.requires_grad=True
 def forward(self,x):return self.head(self.encoder(x))