Add MEG-GPT Cam-CAN weights converted for braindecode.models.MEGGPT (OHBA-analysis/MEG-GPT@78ff9439 models/, braindecode PR #1555 @ 2f253bc1)
5c49e27 verified | """Convert the released MEG-GPT + EphysTokenizer weights into one braindecode MEGGPT state dict. | |
| Inputs (OHBA-analysis/MEG-GPT @78ff9439, models/): | |
| meg-gpt/model_state.pt MEGGPT (meg_gpt/models/meg_gpt.py) state dict | |
| tokenizer/model_state.pt EphysTokenizer state dict | |
| tokenizer/vocab.pkl vocabulary ({"label_map": codebook entry -> token id, ...}) | |
| Dropped on purpose (documented in the PR): the Cam-CAN session embedding | |
| (input_embedding_layer.extra_embeds.0.*), the next-token prediction head, the | |
| attention-mask buffers (rebuilt from the config, checked equal in parity.py), | |
| the tokenizer decoder (token_basis_layer, only used to reconstruct signals) and | |
| the annealing temperature buffer. | |
| usage: python convert.py <models dir> <out dir> | |
| <models dir> is ``models/`` of the original repository; <out dir> receives the | |
| braindecode export (``MEGGPT.save_pretrained``: config.json, model.safetensors, | |
| pytorch_model.bin) that ``MEGGPT.from_pretrained(<out dir or Hub repo>)`` loads. | |
| The classification head (n_outputs=2) is not part of the release: it is | |
| randomly initialised (seed 0); pass ``n_outputs`` to ``from_pretrained`` to | |
| replace it. | |
| """ | |
| import pickle | |
| import re | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| _GPT_RULES = [ | |
| (r"input_embedding_layer\.token_embed\.base_module\.weight", "token_embedding.weight"), | |
| (r"input_embedding_layer\.pos_embed\.base_module\.position_embeddings", "time_embedding"), | |
| (r"input_embedding_layer\.channel_embed\.base_module\.position_embeddings", "channel_embedding"), | |
| (r"transformer_decoder\.layers\.(\d+)\.norm_time1\.norm_layer\.", r"layers.\1.norm1."), | |
| (r"transformer_decoder\.layers\.(\d+)\.norm_time2\.norm_layer\.", r"layers.\1.norm2."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_k\.", r"layers.\1.attn.patch_k."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_v\.", r"layers.\1.attn.patch_v."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_k\.", r"layers.\1.attn.k_proj."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_v\.", r"layers.\1.attn.v_proj."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.query_projection\.", r"layers.\1.attn.q_proj."), | |
| (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.output_projection\.", r"layers.\1.attn.out_proj."), | |
| (r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.0\.", r"layers.\1.mlp.0."), | |
| (r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.2\.", r"layers.\1.mlp.3."), | |
| ] | |
| _GPT_DROP = re.compile( | |
| r"input_embedding_layer\.extra_embeds\.|prediction_head\.|.*\.gasp_layer\.attention_mask$" | |
| ) | |
| _TOK_RULES = [ | |
| (r"encoder_layer\.rnn_layers\.0\.", "tokenizer.rnn."), | |
| (r"token_weights_layer\.dense_layer\.", "tokenizer.proj."), | |
| (r"token_weights_layer\.norm_layer\.", "tokenizer.norm."), | |
| ] | |
| _TOK_DROP = re.compile(r"decoder_layer\.token_basis_layer\.|token_weights_layer\._temperature$") | |
| def _rename(state, rules, drop): | |
| out, dropped = {}, [] | |
| for key, value in state.items(): | |
| if drop.match(key): | |
| dropped.append(key) | |
| continue | |
| for pattern, repl in rules: | |
| new, n = re.subn("^" + pattern, repl, key) | |
| if n: | |
| out[new] = value | |
| break | |
| else: | |
| raise KeyError(f"no conversion rule for {key}") | |
| return out, dropped | |
| def convert(gpt_state, tok_state, vocab): | |
| """Return (braindecode state dict, list of dropped original keys).""" | |
| gpt, dropped_gpt = _rename(gpt_state, _GPT_RULES, _GPT_DROP) | |
| tok, dropped_tok = _rename(tok_state, _TOK_RULES, _TOK_DROP) | |
| state = {**tok, **gpt} | |
| state["tokenizer.vocab"] = torch.as_tensor(vocab["label_map"], dtype=torch.long) | |
| return state, dropped_gpt + dropped_tok | |
| def load_released(models_dir): | |
| models_dir = Path(models_dir) | |
| gpt = torch.load(models_dir / "meg-gpt/model_state.pt", map_location="cpu", weights_only=True) | |
| tok = torch.load(models_dir / "tokenizer/model_state.pt", map_location="cpu", weights_only=True) | |
| with open(models_dir / "tokenizer/vocab.pkl", "rb") as f: | |
| vocab = pickle.load(f) | |
| return gpt, tok, vocab | |
| def export(models_dir, out_dir): | |
| from braindecode.models import MEGGPT | |
| state, dropped = convert(*load_released(models_dir)) | |
| torch.manual_seed(0) | |
| model = MEGGPT(n_chans=52, n_outputs=2, sfreq=250).eval() | |
| res = model.load_state_dict(state, strict=False) | |
| assert sorted(res.missing_keys) == ["final_layer.1.bias", "final_layer.1.weight"], res | |
| assert not res.unexpected_keys, res | |
| model.save_pretrained(out_dir) | |
| return state, dropped | |
| if __name__ == "__main__": | |
| state, dropped = export(sys.argv[1], sys.argv[2]) | |
| print(f"exported {len(state)} converted tensors to {sys.argv[2]}; dropped {len(dropped)}: {dropped}") | |