Repository navigation
Add BrainTokenizer, the BrainOmni EEG/MEG VQ-VAE tokenizer - #1230
Merged
bruAristimunha merged 6 commits intoOct 6, 2026
Merged
Conversation
BrainOmni is ported from OpenTSLab/BrainOmni (MIT); its SEANet codec from Meta's EnCodec (MIT) and its residual VQ from lucidrains/vector-quantize-pytorch (MIT). Per the model-port convention a port links to its licence rather than bundling the text, so add only the file lines for brainomni.py and the shared modules/quantization.py under the existing MIT section (the MIT link is already present) and leave pyproject.toml/MANIFEST.in untouched.
Port of the OpenTSLab BrainOmni tokenizer (MIT) with strict loading of the authors' raw BrainTokenizer.pt beside the architecture (BrainTokenizer.from_opentslab_config + load_state_dict key mapping). Sensor geometry, SEANet codec, residual VQ (EMA codebooks, K-means init, rotation trick) and the cross-attention bridges stay private in braindecode/models/brainomni.py; no public braindecode.modules or models.util API is added. Tests live in test/unit_tests/models/test_brainomni.py (Hub checks are marked network). Registration: models export, summary.csv, util registry lists, docs/api.rst, whats_new, NOTICE, TorchScript skip.
Resolved NOTICE.txt by keeping both MIT entries.
Resolved docs/whats_new.rst by keeping both entries.
bruAristimunha
added
needs-replication
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Dead-code replacements are discarded, and multi-window reconstruction is corrupted for non-divisible window lengths.
Review effort: Balanced
Findings: 1
Open (2)
What changed in this PR
Adds BrainTokenizer, BrainOmni’s geometry-aware EEG/MEG VQ-VAE tokenizer, with checkpoint compatibility and shared EMA quantization components.
Changes:
- Implements SEANet encoding/decoding, sensor geometry, tokenization, and checkpoint remapping.
- Adds reusable EMA and residual vector quantizers.
- Registers, documents, licenses, and comprehensively tests the model.
| File | Description |
|---|---|
braindecode/models/brainomni.py |
Implements BrainTokenizer and its codec architecture. |
braindecode/modules/quantization.py |
Adds shared EMA/RVQ components. |
braindecode/models/__init__.py |
Exports BrainTokenizer. |
braindecode/modules/__init__.py |
Exports quantization modules. |
braindecode/models/util.py |
Registers model integration metadata. |
braindecode/models/summary.csv |
Adds model summary metadata. |
test/unit_tests/models/test_brainomni.py |
Tests geometry, quantization, checkpoints, and inference. |
test/unit_tests/models/test_integration.py |
Records TorchScript incompatibility. |
docs/api.rst |
Adds API documentation entry. |
docs/whats_new.rst |
Adds release-note entry. |
NOTICE.txt |
Records MIT-derived components. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #1230 +/- ##
==========================================
+ Coverage 87.99% 88.17% +0.18%
==========================================
Files 151 153 +2
Lines 17628 18265 +637
==========================================
+ Hits 15511 16106 +595
- Misses 2117 2159 +42 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Fresh compiled codebooks and dead-code replacement currently produce incorrect quantization behavior.
Review effort: Balanced
Findings: 2
Open (5)
bruAristimunha
added a commit
to bruAristimunha/braindecode
that referenced
this pull request
Oct 6, 2026
…via w41/brainomni-2-tokenizer)
5 tasks done
bruAristimunha
added a commit
to bruAristimunha/braindecode
that referenced
this pull request
Oct 6, 2026
whats_new: keep both; refs for NeuroRVQ (braindecode#1218), BrainTokenizer (braindecode#1230), BrainOmni (braindecode#1231) point to the merged PRs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Split 2/3 of #1043, stacked on #1229 (its diff shrinks to the tokenizer once #1229 merges). Adds
BrainTokenizer, the EEG/MEG VQ-VAE tokenizer of BrainOmni (NeurIPS 2025): SEANet (EnCodec-style) codec + residual vector quantization, strictly loading the authors' rawBrainTokenizer.pt.The EMA vector-quantization layers are extracted into a shared
braindecode/modules/quantization.py(EMACodebook,VectorQuantizer,ResidualVectorQuantizer) so TFMTokenizer (#1202) and NeuroRVQ (#1223) can drop their private copies; the extraction is behaviour-preserving (seeded random-init forward: identical state-dict keys, max-abs diff 0.0). Reusesfunctional.rotate_pairs,modules.FeedForwardBlockand nativenn.RMSNorm.license="mit".Licence: MIT.
Replication: BrainOmni-tiny on PhysioNet-MI, test balanced accuracy 0.5796 ± 0.0164 vs the paper's 0.580 ± 0.019 (released protocol, 30/30 cells).