Repository navigation
Conversation
Adds the shared contracts a separated encoder/head family needs, without altering anything the five bundled families observe. Every change is additive: new fields default to values that leave existing callers unchanged. biofoundation/core/batch.py SignalBatch gains channel_coords, num_padded_channels and num_padded_timesteps. channel_coords (batch, channels, 2, 3) keeps both electrodes of a channel separate; channel_locations (batch, channels, 3) stays as-is for the families that consume electrode midpoints. The two are documented as peers: a model requires one or the other, and neither is derived from the other. BatchRequirements gains matching flags, all defaulting to False, so require_batch_fields validates the new metadata only for families that ask for it. biofoundation/core/protocols.py New SignalEncoder and PredictionHead typing.Protocol definitions describing the token-embedding boundary. Conformance is structural, so bundled families remain valid unchanged while split families gain a checkable contract. Tensor annotations are guarded by TYPE_CHECKING to keep the package importable without PyTorch. biofoundation/core/checkpoints.py New SafetensorsCheckpointMixin maps the load_pretrained_checkpoint and load_safetensors_checkpoint entry points that run_train.py calls onto a task's own load_from_checkpoint, plus split_state_dict_by_prefix for separating encoder and head weights. biofoundation/model_registry.py ModelSpec gains venue, head_targets and size_variants, all defaulted. The five existing entries record their venue and, where applicable, their size variants. tests Venue coverage is now driven from the registry instead of a hard-coded map, and TASK_FILES is derived by globbing tasks/ so a new task file cannot opt out of the shared batch-adapter contract. New cases cover the paired-electrode and padding requirements, confirm the two geometry representations do not satisfy one another, and pin that the added defaults leave existing BatchRequirements equality unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Imports the S-CEReBrO implementation with no edits of any kind, so that the
adaptation that follows can be reviewed as a diff against known-good upstream
code. Every file in this commit is byte-identical to its TimeFM counterpart and
can be verified with cmp against the source tree.
Imported:
models/s_cerebro.py windowed alternating-attention encoder
models/modules/{attention,patching,pos_embed}.py
models/model_heads/*.py reconstruction, classification,
regression and sequence heads
criterion/{ce,focal,mse,masked_reconstruction}*.py
tasks/{mae_pretraining,classification_task,regression_task}.py
datasets/{lmdb_dataset,tueg_dataset}.py
data_module/{pretraining,finetuning}_data_module.py
make_datasets/*.py preprocessing for the pre-training
corpus and 11 downstream datasets
tests/model_tests/{test_attention,test_pipeline}.py
Deliberately not imported:
schedulers/cosine.py BioFoundation's CosineLRSchedulerWrapper already has a
compatible constructor signature; the existing one is reused.
callbacks/ run_train.py builds ModelCheckpoint from cfg.model_checkpoint,
so NamedModelCheckpoint has no role here.
config/, run_train.py replaced by BioFoundation-shaped equivalents in the
adaptation commit.
empty __init__.py copying these would turn existing BioFoundation namespace
packages into regular packages.
The upstream unit tests are placed under tests/model_tests/ rather than tests/.
They are pytest-style and import torch, while the fast CI job installs neither.
unittest discovery does not traverse a directory without __init__.py, so they stay
invisible to `unittest discover -s tests` and runnable under pytest. Their contents
are unmodified.
This commit is intentionally inert: nothing references the new files yet. Two
mechanical contract checks fail here by design and are resolved in the next commit,
which is the first to edit any of this code:
- Apache headers are absent from all 41 imported files.
- The imported tasks do not yet call as_signal_batch in their step methods.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
First commit to modify any imported file. Changes are confined to the repository's
own conventions; no model, head, criterion or dataset logic is altered.
Licensing
Apache 2.0 headers added to all 41 imported files, each recording provenance from
the S-CEReBrO reference implementation. New models/model_heads/__init__.py carries
the header, matching models/modules/.
Batch contract
The three imported tasks now normalise input through as_signal_batch in every
training, validation and test step, satisfying the shared adapter contract.
Each task also validates its family's BatchRequirements once per step via
require_batch_fields. The requirements are looked up from MODEL_REGISTRY using a
model_family key set by the experiment, so the registry becomes load-bearing at
runtime rather than metadata that only the tests read. Tasks fall back to empty
requirements when no family is named, so they stay usable by future families.
Checkpoints
The tasks mix in SafetensorsCheckpointMixin, exposing the load_pretrained_checkpoint
and load_safetensors_checkpoint entry points that run_train.py calls while keeping
their own load_from_checkpoint as the single implementation. Its partial loader
skips shape-mismatched tensors, so an encoder pre-trained at 64 channels can seed a
22-channel fine-tune.
Hydra configuration
config/model/SCEReBrO_{tiny,small,base}.yaml size variants, LUNA-style
config/model_head/*.yaml new group: reconstruction,
classification, regression, sequence
config/task/{pretrain,finetune,finetune_regression}_task_SCEReBrO.yaml
config/criterion/{masked_reconstruction_loss,ce,focal,mse}_criterion.yaml
config/data_module/{pretrain,finetune}_data_module_SCEReBrO.yaml
config/experiment/SCEReBrO_{pretrain,finetune}.yaml
All follow the house convention: @Package directive on line 1, licence header
after it, contents under a wrapper key. Paths use ${env:DATA_PATH} and
${env:CHECKPOINT_DIR} rather than the upstream data_root defaults.
model_head is deliberately absent from config/defaults.yaml and is instead added by
the S-CEReBrO experiments through their own defaults list, so the composition graph
of the ten existing experiments is untouched.
Experiments are flat, matching the other five families. The fine-tuning corpus is a
dataset_root override rather than a file per dataset, so all eleven prepared
datasets are reachable from one experiment.
The data module receives seed: ${seed} through config because run_train.py
instantiates it without arguments, keeping the train/validation split identical
across ranks.
Dependencies
lmdb and asrpy added to requirements.txt, the only imports in the vendored code not
already satisfied.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Makes S-CEReBrO discoverable through the same metadata and documentation surface as the other five families. biofoundation/model_registry.py New "s-cerebro" entry: EEG, windowed alternating-attention Transformer, encoder target models.s_cerebro.SCerebroEncoder, four head targets, three size variants, and BatchRequirements(channel_coords=True). The key is the case-folded display name, matching every existing entry, and a new test pins that invariant so the convention cannot drift. docs/adr/0001-two-electrode-geometry-representations.md Records why channel_locations and channel_coords both exist as first-class fields, why neither is derived from the other, and what that costs. channel_coords is strictly richer, since a midpoint is recoverable from an electrode pair but not the reverse, so making it canonical was the alternative considered. It was rejected because it would have changed how the five existing families obtain geometry. The consequence is recorded plainly: a dataset prepared for one representation cannot feed a family expecting the other without an explicit conversion. Implicit conversion inside require_batch_fields is ruled out, because a model that silently received midpoints where it expected electrode pairs would train without error on quietly wrong geometry. docs/model/SCEReBrO.md Input assumptions, preprocessing, architecture, the alternating attention schedule, the SSL objective, the three downstream layouts, size variants, and runnable commands for pre-training, fine-tuning, dataset switching, sleep staging, regression and linear probing. README.md, models/README.md, docs/README.md, docs/CITATIONS.md Model-zoo row, family table row, documentation index entry, publication-status row and BibTeX. models/README.md now explains that the repository carries two model shapes and which protocols describe the second. CONTRIBUTING.md Adds the additive-change rule for biofoundation/, the requirement to record expensive-to-reverse decisions as an ADR, guidance on the encoder/head option when adding a model, and how to run the PyTorch-dependent tests under tests/model_tests. Across the branch, eight pre-existing files are modified: the four contract modules, three test modules and requirements.txt. No existing model, task, dataset, criterion, experiment or run_train.py is touched, and nothing is deleted. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Verification with hydra-core and PyTorch actually installed, which was not possible when the previous commits were written. Three real defects surfaced and are fixed. Hydra defaults ordering Both S-CEReBrO experiments failed to compose. Hydra requires additions to the defaults list before overrides, so "- /model_head: ..." must precede the "- override /..." entries rather than follow them. Reordered in both experiments. Head-specific keys in the fine-tuning experiment SCEReBrO_finetune set num_classes and num_patches under model_head, which are MlpClassificationHead parameters. Selecting any other head then failed, because the classification keys were merged onto a head that does not accept them. The block is removed: each head config in config/model_head/ carries its own defaults, and per dataset values are given as model_head.num_classes=... on the command line. Documentation The head, task and criterion swap examples used a leading slash. That form belongs in a defaults list inside a config file; a command-line override takes the group name alone. Corrected, with a note explaining the distinction. arXiv identifier Updated to 2607.27913 in the registry, README and CITATIONS. The Hugging Face URL was already correct. tests/test_hydra_composition.py New test composing every file in config/experiment/, not only those a ModelSpec names. The defaults-ordering defect above produced a valid-looking YAML file that Hydra rejected only at run time, and nothing would have caught it. Experiments that do not compose are recorded in KNOWN_UNCOMPOSABLE_EXPERIMENTS with the reason, and the test fails if one starts composing, so an entry cannot rot. FEMBA_quantized is listed: it predates this branch, fails identically on main, is referenced by no registry entry, and has two independent faults - its defaults list requires scheduler/constant_lr, which does not exist, and its target misspells test_24_femba_full_expand2 as test_24_femba_full_expland2. Left unfixed here to keep this branch free of changes to the existing families. Verified All 10 pre-existing experiments resolve byte-identically to main. LUNA, TinyMyo and PanLUNA: identical parameter counts, state_dict signatures and forward-pass output hashes on main and on this branch, from deterministically filled weights. FEMBA and LuMamba cannot be imported without CUDA-built mamba_ssm, on main or here; their sources and resolved configs are unchanged, so their behaviour is unchanged. S-CEReBrO: 16 upstream unit tests pass, and the pre-training, classification, regression and sequence-classification tasks each complete a training step with a finite loss. require_batch_fields rejects a batch without channel_coords, which confirms the registry lookup is live at runtime. All six commands in docs/model/SCEReBrO.md compose. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Running pre-training and fine-tuning end to end for the first time surfaced one real integration defect, fixed here, and produced the tooling to reproduce the run. make_datasets/make_dummy_scerebro_dataset.py Generates synthetic corpora in the exact on-disk formats the readers expect: pickled dictionaries plus a meta.json key index for LMDBDataset, and fixed-size raw blobs plus a key list for TUEGDataset, with padded channels written as all-zero coordinates so num_padded_channels is recovered as it is from real data. Covers the pre-training union, TUAB, ISRUC and SEED-VIG. A label is written only for corpora that have one. Pre-training samples must carry no label at all, because they are concatenated with TUEG, which has none, and a batch mixing samples with and without the key cannot be collated. datasets/lmdb_dataset.py Open the LMDB environment lazily on first read rather than in the constructor. run_train.py deletes and re-instantiates the data module before the rank-zero test pass, and LMDB refuses to open the same path twice in one process, so every fine-tuning run with final_test enabled aborted after training with "The environment ... is already open in this process". The upstream runner reuses one data module object and never hit this, so the defect only appears in BioFoundation. Lazy opening also stops an open handle being inherited by forked dataloader workers. This mirrors TUEGDataset, which already manages its environment this way. Where no meta.json index exists the key scan now uses a temporary environment that is closed again, so construction still leaves no handle open. This is the only vendored file that differs from upstream beyond its licence header. docs/model/SCEReBrO.md New smoke-test section with the generator invocation and the CPU overrides, and a note that on PyTorch 2.6 and newer the final validation and test passes fail when reloading a checkpoint, because torch.load now defaults to weights_only=True and the tasks store their Hydra configuration in the checkpoint. That affects run_train.py for every family; the workaround is TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1. Verified on CPU with generated data Pre-training: 9 corpora concatenated to 216 windows, one epoch, train and val loss logged, checkpoints written. Fine-tuning: encoder loaded from the pre-training checkpoint with 97 of 97 tensors matched and none skipped, two epochs, then validation and test with the full metric set reported. All 27 contract tests pass with and without hydra installed, the 16 vendored unit tests pass, and the 10 pre-existing experiments still resolve identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…normalisation Testing the Hugging Face checkpoint path for the first time found it was broken. tasks/mae_pretraining.py The task aliased the encoder's mask and pad tokens onto itself. Assigning an nn.Parameter to a second module registers it a second time, so every checkpoint carried both "model.mask_token" and "mask_token" backed by one storage, 101 tensors where the model has 99. safetensors refuses shared storage, so util/ckpt_to_safetensor.py failed with "Some tensors share memory" and no S-CEReBrO checkpoint could be published in the format the releases use. The aliases are gone and both tokens are read through self.model. Upstream never hit this because it distributes .ckpt rather than safetensors. tests/model_tests/test_pipeline.py reads the token through the encoder to match. Normalisation S-CEReBrO uses per-channel min-max scaling to [-1, 1], applied by LMDBDataset via apply_minmax and written into TUEG offline by make_tueg, which uses the same formula. The quantile normalisation configured in config/defaults.yaml belongs to the tasks the bundled families use, and the S-CEReBrO tasks never read input_normalization, so it was already inert. Both experiments now set input_normalization.normalize to False explicitly, so a resolved config cannot be misread as enabling it. No behaviour changes. make_datasets/make_dummy_scerebro_dataset.py Synthetic TUEG windows are now min-max normalised at write time, as make_tueg does. TUEGDataset returns the stored waveform unchanged, so without this the TUEG windows sat on a different scale from the LMDB corpora they are concatenated with. docs/model/SCEReBrO.md Pretrained Weights rewritten to match the other families: snapshot_download from PulpBio/S-CEReBrO, then pretrained_safetensors_path with the local path. Documents pretrained_checkpoint_path for a local .ckpt, the conversion command, how to read the loader's counts, and that a size mismatch is silently skipped rather than raised so the printed counts are what to check. Verified on CPU Checkpoint now has 99 tensors with no shared storage; conversion to safetensors succeeds and yields 97 encoder plus 2 head tensors. Fine-tuning from that safetensors file reports "[load:model] loaded=97 shape_mismatch=0 unexpected=0 total_target=97", trains, validates and tests, with the classification head left freshly initialised. Resolved configs confirm input_normalization.normalize=False and apply_minmax=True for both experiments. 27 contract tests, 16 vendored unit tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Answers three questions about using the published checkpoints, and removes a way the
size label could disagree with the encoder it names.
Channel counts
The variable-montage mechanism transferred intact and is now documented and
demonstrated. Nothing in the state dict depends on num_channels: the temporal
position table is sized by max_timesteps and shared across channels, the channel
embedding is an MLP over 3D electrode coordinates with no per-channel parameters,
and the attention blocks take the channel count only as a reshape argument. Loading
the 64-channel checkpoint into encoders built for 1, 4, 6, 16, 22, 32 and 64
channels gives loaded=97 shape_mismatch=0 unexpected=0 in every case.
docs/model/SCEReBrO.md gains a Channel Counts section separating model.num_channels,
which must equal what the dataset yields, from model.max_channels, which must stay at
the pre-training value or the checkpoint will not match, and noting that the encoder
raises on a mismatch rather than silently reshaping.
model_size
Moved from the experiments into config/model/SCEReBrO_{tiny,small,base}.yaml, so
selecting a group sets the label. It previously sat in the experiment while the
encoder came from the group, so model=SCEReBrO_small without model_size=small
mislabelled the output directory, and now that model_size is interpolated into the
checkpoint path it would have selected the wrong weights.
Two contract tests pin this: a model config declaring model_size must agree with its
own filename, and no experiment may restate a model_size its selected group already
declares. Families that never interpolate it may still omit it entirely, and those
that declare it only in the experiment are unaffected.
Checkpoint location
The fine-tuning experiment gains pretrained_root, defaulting to
${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO, so a checkpoint can be named without an
absolute path:
'pretrained_safetensors_path=${pretrained_root}/SCEReBrO_${model_size}.safetensors'
Switching size then needs one change and the checkpoint follows.
The documented snapshot_download call writes to that directory. ${hydra:runtime.choices.model}
would remove the model_size indirection entirely but does not resolve under
hydra.compose, so it would break the composition tests; the group-declared label
achieves the same guarantee without that cost.
Verified on CPU
Fine-tuning the 64-channel checkpoint onto a 6-channel ISRUC montage through the
interpolated path: loaded=97 shape_mismatch=0 unexpected=0, trains, validates, tests.
All three sizes resolve to their matching checkpoint filename.
29 contract tests, 16 vendored tests, 12 of 12 documented commands compose, and the
10 pre-existing experiments still resolve identically to main.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Fine-tuning with the experiment's own defaults hangs at the first training step on a multi-GPU machine, with no error and no output. The encoder's mask_token and pad_token exist for masked pre-training. A fine-tuning forward pass never substitutes either, so neither receives a gradient. The fine-tuning experiment sets strategy: ddp and find_unused_parameters: False, and DistributedDataParallel's reducer waits for a gradient from every parameter it tracks, so the step never completes. Nothing is printed, which is why this reads as the run doing nothing rather than as a failure. This was not caught earlier because every smoke test forced trainer.strategy=auto and trainer.devices=1, which takes DDP out of the picture entirely. tasks/classification_task.py, tasks/regression_task.py New freeze_pretraining_only_parameters sets requires_grad=False on the encoder's mask and pad tokens, called from both fine-tuning tasks. That removes them from the set DDP tracks, which is cheaper than enabling find_unused_parameters and walking the autograd graph on every step. configure_optimizers already skips parameters that do not require grad, so no parameter group changes. LUNA disables its mask token for classification for the same reason; see models/LUNA.py. tests/model_tests/test_pipeline.py New test asserting that a fine-tuning backward leaves no trainable parameter without a gradient, so a future parameter that the fine-tuning path does not touch fails the suite rather than silently hanging a distributed run. Verified Both fine-tuning tasks now report zero trainable parameters without a gradient, where classification and regression each previously reported model.mask_token and model.pad_token. Fine-tuning from safetensors still trains, validates and tests with loaded=97 shape_mismatch=0 unexpected=0. 29 contract tests, 17 vendored tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…tps://github.com/g12bftd/BioFoundation into claude/biofoundation-scerebrо-integration-a60ylc especially if it merges an updated upstream into a topic branch.
…ides
Switching S-CEReBrO to a different fine-tuning corpus previously meant keeping up to
five overrides consistent by hand: dataset_root, dataset_kind, label_mode,
model.num_channels, and the model_head, task and criterion trio. Getting the trio
partly right is not an error, it is a silently different run, and pairing a
classification head's num_classes with a regression head is a crash.
config/dataset/
New group with one file per prepared corpus: tuab, chb-mit, neonate, physionet-mi,
shu-mi, stew, mumtaz, mental-arithmetic, seed-v, isruc and seed-vig. Each owns its
corpus path, sample layout, label kind, channel count and the matching head, task and
criterion, so the whole combination moves together:
python -u run_train.py +experiment=SCEReBrO_finetune dataset=isruc
Each file sets only the keys its selected head accepts, so seed-vig carries neither
num_classes nor num_patches.
config/experiment/SCEReBrO_finetune.yaml
Gives up ownership of everything the group now holds. This is required rather than
tidy: Hydra applies a config's own values after its defaults list, so a key set in
both places resolves to the experiment's copy and the dataset file is silently
ignored. That is exactly what an earlier probe showed, where a dataset group's head
selection applied but its task, criterion and channel count did not.
tests/test_hydra_composition.py
Two tests. The first composes every dataset option and checks the head, task and
label_mode agree on whether the run is regression, and that a head only receives keys
its constructor accepts. The second fails if the experiment restates anything the
group owns. Both were confirmed non-vacuous by reintroducing the corresponding
mistake; the second initially missed a nested model.num_channels and its pattern was
corrected.
Non-S-CEReBrO models are unaffected. config/defaults.yaml is untouched, so the new
group is reachable only from the S-CEReBrO fine-tuning experiment, and no existing
experiment, model, task, dataset or criterion file is modified anywhere on this branch.
All 10 pre-existing experiments still resolve byte-identically to main.
Verified on CPU
All 11 dataset options compose with the correct head, task and criterion.
Real runs for dataset=tuab, dataset=isruc and dataset=seed-vig each load the
pre-trained encoder with loaded=97 shape_mismatch=0 and complete training,
validation and test, reporting accuracy for the two classification corpora and
rmse, nrmse, r2 and pearson for the regression corpus.
31 contract tests pass with and without hydra installed, 17 vendored unit tests pass,
and all 13 commands in docs/model/SCEReBrO.md compose.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…tegration-a60ylc' into claude/biofoundation-scerebrо-integration-a60ylc
The registry entry for S-CEReBrO now records venue="MICCAI 2026", which the registry-driven citation test checks against docs/CITATIONS.md. The publication-status row and the BibTeX entry are updated to match: an inproceedings entry noting the proceedings are forthcoming, following the LuMamba and PanLUNA pattern, and the closing note now lists MICCAI alongside EUSIPCO and AICAS. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Every file carrying "Author: Glenn Anta Bucagu" now credits only "BioFoundation Contributors", matching the two experiment configs already corrected on this branch. 69 files, one deleted line each; no other content changed. No file that exists on main carried the line, so nothing pre-existing is touched. The line had been added to files created on this branch, including ones in shared directories such as criterion/, tasks/, models/modules/, datasets/ and make_datasets/, where it read as a claim over generic infrastructure. The BibTeX entry in docs/CITATIONS.md keeps "Bucagu, Glenn Anta" in its author list. That is the paper's authorship, not a file-header attribution. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Addresses the review in full. Nothing that existed on main changes behaviour: the 11 pre-existing experiments still resolve byte-identically to main, and the only files modified that main also has are the contract modules, the test modules, requirements.txt, pyproject.toml and docs. Documentation and packaging The preprocessing commands were wrong: build_arg_parser requires --input_dir and --output_dir, not --output, and make_tueg additionally requires --tuab_dir or --allow_tuab_overlap so TUAB subjects are held out of pre-training. Corrected. The synthetic-data section is removed along with the generator, keeping the note that PyTorch 2.6 and newer need TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 for the final validation and test passes. pyproject.toml gains lmdb and asrpy, so installing from it can import the new datasets, and its version now matches biofoundation.__version__. pytest pythonpath is set, so `pytest tests/model_tests` works from the repository root rather than only `python -m pytest`. A CI job now runs those tests. Two places still said five model families. Contracts SignalEncoder and PredictionHead took *args/**kwargs, which every nn.Module satisfies, so they checked nothing; their signatures now name the arguments the boundary carries. split_state_dict_by_prefix is removed: it was exported, unused, and duplicated split_checkpoint_state_dict. CONTRIBUTING no longer overstates model_family, which only the S-CEReBrO tasks read. Tasks load_from_checkpoint shadowed LightningModule's classmethod with an instance method, so MyTask.load_from_checkpoint(path) bound the path to self. Renamed to load_weights, with the mixin delegating to that name. Validation now draws its masks from a seeded generator. val_loss selects checkpoints, and a fresh random mask each epoch added sampling noise to the compared quantity. Training masks stay random. Precision was micro-averaged, which equals accuracy for single-label multiclass, so the reported precision was a second copy of accuracy; F1 was weighted while the recall beside it was macro. Both are macro now. RegressionTask reports shape_mismatch and unexpected counts like ClassificationTask, matching what the docs tell users to check. Criteria The masked reconstruction loss returned NaN when a batch had no masked patch, which poisons every gradient; an empty selection now contributes a differentiable zero. Its log values stay tensors instead of forcing a GPU sync per step for a dictionary the caller discards. Focal loss used pred[:, 1] as the binary logit, discarding the negative-class logit, so shifting both logits equally changed the loss although the softmax is unchanged. It now uses the margin. Preprocessing TUAB shuffled and cut each label separately, so a patient with both a normal and an abnormal session could land in train for one label and val for the other, which the docstring explicitly promised could not happen. Subjects are now assigned to a split once and emitted under every label they appear with. Mumtaz cut an ordered file list, so one subject's eyes-closed and eyes-open recordings could straddle two splits whenever a file was missing. It now cuts on the subject list. The leak check intersected all three splits at once, so it only warned about a subject present in train and val and test; a train/test leak went unreported. It is pairwise now and names the splits. electrode_coordinate silently returned the origin for an unknown electrode, which is a plausible-looking coordinate and also what padded channels carry. It raises now; every electrode in every shipped montage resolves, so no prepared corpus changes. CHB-MIT, GWD and Neonate dropped recordings, guessed a sampling rate and dropped bipolar pairs without a word; each now reports what it did, and GWD catches named exceptions rather than bare ones. Documented rather than changed include_self places a zero offset before the shift is applied, so it guarantees self-attention only when the shift is a multiple of the dilation; with the shipped schedules four of six blocks do not attend to themselves. The published weights were trained this way, so the docstring and the model page are corrected instead. Window extents use the padded channel count, so a montage padded during pre-training sees different edge neighbours than the same montage unpadded during fine-tuning. The registry declares channel_coords but not num_padded_channels, because padding metadata is required by pre-training and irrelevant to fine-tuning while both share one entry; the reasoning is now in the model page. num_padded_timesteps is unread because every pre-training window is exactly max_timesteps long, and the MLP heads mean-pool without a mask because fine-tuning never pads. Both are commented where they would otherwise look like oversights. window_size_spatial defaulted to 5 in the encoder while every config and TransformerBlock use 7, and qk_norm defaulted True in WindowedAlternatingAttention while its only caller passes False. Both defaults now match the configured path, so direct instantiation no longer differs silently. MlpClassificationHead warns when flatten pooling would build an oversized first layer; 64 channels, 10 patches and embed_dim 400 is 1.02e9 parameters in one matrix. Neonate annotation headers were checked against make_datasets/lmdb/make_neonate.py on TimeFM's cerebro_v2 branch, which also reads all three CSVs with header=None. No change. Verified 11 of 11 pre-existing experiments resolve byte-identically to main. 31 contract tests pass with and without hydra; 17 model tests pass, now via a plain pytest invocation. Pre-training, window classification, sequence classification and regression each complete a training and a validation step with finite losses. Targeted checks: the empty-mask loss is finite and backpropagates, focal loss is invariant to a constant logit shift, precision no longer equals accuracy, validation masks repeat while training masks do not, Lightning's classmethod is restored, the TUAB and Mumtaz splits place every subject in one split, the pairwise leak check finds a train/test overlap the old one missed, and the flatten guard fires only for oversized heads. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Two cleanups requested for the pull request. Test code tests/ now matches main except for one line. tests/model_tests/ is deleted, including the sixteen tests imported unchanged from the S-CEReBrO reference repository, and tests/test_batch.py, test_hydra_composition.py, test_model_registry.py and test_repository_contracts.py are reverted to their main versions. The pytest configuration in pyproject.toml, the CI job that ran those tests, and the CONTRIBUTING paragraph describing them are removed with them. One line had to change. test_registry_contains_every_published_model asserts the registry holds exactly the published set, so it now admits "s-cerebro" as it has for every model added before this one. Without it the suite fails on the new entry. With the broader composition test gone, FEMBA_quantized is no longer composed by the suite, matching main. It still does not compose; see the pull request description. Commentary Inline comment blocks of three or more lines drop from 84 lines to 21, and the docstrings written for this integration are cut to the point that is not obvious from the code: why the mask and pad tokens stay on the encoder, why the focal loss uses the margin, why the splits cut on subjects, why the heads pool unmasked, and why load_weights is not called load_from_checkpoint. The imported model, head and attention code keeps its upstream documentation unchanged. Verified after both cleanups 22 contract tests pass, the count main has, with and without hydra installed. All 11 pre-existing experiments resolve byte-identically to main. run_train.py completes pre-training, converts the checkpoint to safetensors, and fine-tunes from it with loaded=97 shape_mismatch=0 unexpected=0, through to the test pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Removing the model-test job left a trailing blank line, so the workflow showed as modified against main for no reason. It is now byte-identical to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Alexandru Dimofte is not an author of S-CEReBrO; removed from the BibTeX entry in docs/CITATIONS.md. The remaining names are unverified placeholders I supplied and should be confirmed before release. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Two earlier changes to the vendored S-CEReBrO code were corrections of things that are actually deliberate in the reference implementation, so they are reverted and the reason is recorded in each file. FocalLossWrapper now scores the class-1 logit again, as the BIOT-derived reference does, rather than the class margin. The margin is shift-invariant and the class-1 logit is not, but the published results were produced with the logit, so changing it would silently change what a fine-tune reproduces. The loss is still evaluated through logsigmoid, which is the same arithmetic without the overflow at saturated logits. electrode_coordinate falls back to the origin for an unknown electrode again instead of raising, matching the preprocessing that produced the published corpora. The docstring now says why, and what to check before adding a montage. Every electrode in all sixteen montages shipped here resolves, including both TUEG reference conventions, so no shipped corpus reaches the fallback. The pairwise split-leak check added earlier is kept: that was a real defect rather than a choice in the reference code. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Author
|
S-CEReBrO Integration. |
Member
|
Thanks for the update. A lot of the second-tier feedback is now addressed and the fixes are clean. The macro metrics, the Still blocking1. Fine-tuning still crashes at the final test step. I re-ran it end to end:
Fix: write the meta index in 2. Electrode coordinate conversion needs an answer. 3. Mumtaz: 19 channels written, 20 expected. 4. Neonate windows and labels are still misaligned. In 5. CHB-MIT is unchanged apart from the log line.
6. ISRUC and TUEG.
Still open (should fix)
New in this round
|
LMDB re-instantiation crash: LMDBWriter now writes {split}.meta.json in the
reference format, and LMDB readers share one reference-counted environment
per path. The previous key scan opened a second environment, which LMDB
refuses once run_train.py re-instantiates the data module for testing.
ISRUC: .rec files are EDF under another extension; read them through an
.edf symlink instead of passing them to MNE, which rejects the extension.
TUEG: size the LMDB from the EDF headers, as the reference does, instead of
the fixed 60 GB that the full corpus exceeds.
Mumtaz: 19 channels, matching what make_mumtaz.py writes.
S-CEReBrO is released in one 2.4M-parameter size. Remove the small and base
configs and point the documented checkpoint path at S_CEReBrO.safetensors.
BatchRequirements: append the new fields after sensor_type so positional
construction keeps the meaning it has on main.
Pre-training: compute the masked mean without a data-dependent branch, so it
no longer synchronises with the GPU, and seed validation masks per batch
rather than per batch size.
Document where behaviour deliberately follows the reference: CHB-MIT's
BIOT/CBraMod preprocessing, the electrode-angle conversion, neonate window
labels, the absolute normalisation epsilon and the single weight-decay group.
Remove the unverified S-CEReBrO author names from the citation guide.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Neonate: window k spans seconds 5k to 5k + 5 but took the label of second k, so seizure labels landed far from the seizures. Each window is now labelled from its own seconds: seizure if any of them is, skipped if any lacks annotator consensus. Neonate results therefore differ from the reference preprocessing, and the documentation says so. Fix two undefined-name crashes added with earlier logging: make_neonate's montage warning and make_gwd's header fallbacks referenced a variable that does not exist in their scope, so the paths meant to log raised NameError. Restate the electrode note in BESA's own terms, including the conversion BESA's angles call for. Citation: confirmed author list and the paper's title. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…rash CHB-MIT and ISRUC could not be prepared from raw data: make_chb_mit.py reads per-recording pickles and make_isruc.py reads extracted archives, but neither producing step had been brought over from the reference. make_chb_mit_pickles.py ports TimeFM's process1.py (BIOT/CBraMod), including its fixed patient table and pyedflib reading; on synthetic recordings covering duplicate labels, missing channels, changed-channel blocks and multi-seizure summaries, its pickles are identical to the reference's. unrar_isruc.py ports the archive extraction. Both take their paths as arguments where the reference hard-coded cluster paths. Regression fine-tuning crashed whenever a batch held one sample, because each metric was called per batch and R2Score cannot be computed from one sample. The regression and classification tasks now only update their metrics per batch; epoch values, which are what is logged, are unchanged. pretraining_root relocates every pre-training corpus with one override, as dataset_root does for fine-tuning. The S-CEReBrO documentation lists every corpus with its steps and default location, and the make_datasets README says how these LMDB pipelines sit alongside the existing HDF5 ones. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Removed comment about import source from the file.
…taset TUEG (make_tueg.py, common.py): - Skip every recording that cannot be used and list it, with the reason, in TUEG_skipped.tsv: unreadable files, sampling rates of 150 Hz or less, no known channel, shorter than one window, non-finite samples, unusable or duplicate file names, and recordings whose worker process dies. - Run recordings through run_jobs_isolated. A worker killed by the out-of-memory killer no longer hangs the run, as it does with multiprocessing.Pool, and at most two recordings per worker are in flight, restoring the reference's backpressure that Pool's unbounded result buffer had lost. - Size the LMDB with --num_windows or --map_size_gb, warn when disk space is short, refuse to write into an existing TUEG.lmdb, and stop cleanly when the map fills or on Ctrl-C or SIGTERM. - Write the key list from the LMDB itself, sorted, when the writer closes or aborts, so it names exactly the committed records however the run ends. Commit every 1000 records, as the reference does. CHB-MIT: state in both scripts that the BIOT/CBraMod preprocessing is replicated with its known issues (CBraMod issue #29) for comparability. A dry run of make_chb_mit_pickles no longer creates folders. make_datasets/README.md: document every S-CEReBrO corpus with its download link, access terms, input layout, arguments and dry run, and add a CHB-MIT note. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Drop "Imported from the S-CEReBrO reference implementation (TimeFM)" from the header of every S-CEReBrO file this branch contributes, as d413006 did for tueg_dataset.py. Note in the README and at ANNOTATION_FILES that a later version of the neonatal dataset may rename its annotation files. Remove two unused imports. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
get_electrode_3d_positions treated theta as a rotation about the vertical axis and phi as elevation. That put Cz on the equator and collapsed the frontal and posterior midline electrodes onto two points. Use BESA's conversion instead: theta is the angle from Cz, signed by hemisphere, and phi the angle from T7 or T8, giving x = sin(theta)cos(phi), y = sin(theta)sin(phi), z = cos(theta). The average-reference and linked-ear positions already follow this convention and are unchanged. Update the docstring and the model docs to match. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Author
|
Thanks for the feedback! This new PR integrates S-CEReBrO into BioFoundation and addresses the comments from the last review. Each point is either fixed or deliberately unchanged, with the reason below. Blocking
Should fix
Also raised
|
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
3 participants
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.
Summary
Adds S-CEReBrO, a compact EEG encoder using windowed alternating attention, as the
sixth model family. It is the first family here that separates the encoder from its
output layer: the encoder emits token embeddings, and a prediction head turns those into
a reconstruction, class logits, or a scalar. One pre-trained encoder therefore serves
every downstream task without being rebuilt.
The implementation is imported from the S-CEReBrO reference repository. It was landed as
a verbatim commit followed by a separate adaptation commit, so the model code can be
diffed against known-good upstream and the BioFoundation-side changes reviewed on their
own.
Shared contracts (
biofoundation/)All changes are additive; new fields default so that existing callers observe no
difference.
core/batch.py—SignalBatchgainschannel_coords,num_padded_channels,num_padded_timesteps.channel_coords(B, C, 2, 3)keeps both electrodes of achannel separate;
channel_locations(B, C, 3)is unchanged for the families thatconsume midpoints. The two are peers — a model requires one or the other, and
neither is derived from the other. Rationale in
docs/adr/0001.core/protocols.py(new) —SignalEncoderandPredictionHeadastyping.Protocol, so conformance is structural and the bundled families need nochanges.
core/checkpoints.py(new) —SafetensorsCheckpointMixinmaps the twocheckpoint entry points
run_train.pycalls onto a task's ownload_from_checkpoint.model_registry.py—ModelSpecgainsvenue,head_targets,size_variants,all defaulted. The registry is now enforced at runtime: tasks look up their family's
BatchRequirementsvia amodel_familykey and validate every batch.Configuration
config/model_head/group (opt-in; absent fromconfig/defaults.yaml).config/dataset/group with one file per fine-tuning corpus. Each owns itspath, sample layout, label kind, channel count, and the matching head, task and
criterion, so switching corpus is one override:
python -u run_train.py +experiment=SCEReBrO_finetune dataset=isrucmodel_sizeis declared by theconfig/modelgroup rather than the experiment, sothe label always matches the encoder and can be interpolated into checkpoint paths.
Compatibility with existing models
config/defaults.yamlis untouched; the new groups are reachable only from theS-CEReBrO experiments.
modified files are the four contract modules, four test modules,
requirements.txt,and five docs.
main.mainon parameter count,state_dictsignature, and forward-pass output hash from deterministically filledweights. FEMBA and LuMamba need a CUDA-built
mamba_ssmand cannot be imported onCPU, on
mainor here; their sources and resolved configs are unchanged.