File size: 4,187 Bytes
4fcc185 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """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))
|