S-CEReBrO Integration

CONTRIBUTING.md

@@ -10,15 +10,17 @@ BioFoundation is both a research repository and an onboarding codebase. Contribu
 - Preserve the full Apache 2.0 header in every Python source and Hydra YAML file.
 - Normalize task input through `biofoundation.core.batch.as_signal_batch`.
 - Keep model-family metadata in `biofoundation/model_registry.py`, including its paper and Hugging Face repository.
+- Change `biofoundation/` additively. New fields, functions, and classes are fine; existing names must not change signature or behaviour, and new defaults must leave existing callers observing no difference. A change that cannot be made additively needs a deprecation period and a major version bump of `biofoundation.__version__`.
+- Record decisions that are expensive to reverse as an ADR under [`docs/adr`](docs/adr/).
 - Use Hydra configs for reproducible settings instead of embedding experiment paths or hyperparameters in Python.
 - Mark local values that a new user must supply with `#CHANGEME` and document the expected value.
 
 ## Adding a Model
 
-1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`.
+1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`. A family may either bundle its output layer into the model, as the five original families do, or separate an encoder from a prediction head under `models/model_heads/` with a matching `config/model_head/` entry. The second shape is described by the protocols in `biofoundation/core/protocols.py` and lets one pre-trained encoder serve every downstream task.
 2. Add pre-training and fine-tuning experiments under `config/experiment/`.
 3. Use an existing task when its behavior fits. New task steps must use the shared batch adapter.
-4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, paper, and Hugging Face URL.
+4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, venue, paper, and Hugging Face URL. Use the case-folded display name as the key, and list any prediction heads in `head_targets`. Setting `model_family` in the experiment lets a task look its family up and enforce the declared batch requirements at runtime. Only the S-CEReBrO tasks read it today; the five bundled families validate nothing at runtime.
 5. Add a model page under `docs/model/` with input assumptions, training details, and checkpoint usage.
 6. Extend the contract tests for any new shared behavior.
 

README.md

@@ -6,7 +6,7 @@ Authors: Thorir Mar Ingolfsson, Anna Tegon, Berkay Döner, Xiaying Wang, Matteo
 
 > **TL;DR:** Choose a model from the table below, install the training dependencies, set `DATA_PATH` and `CHECKPOINT_DIR`, then run `python -u run_train.py +experiment=<MODEL>_pretrain` or the matching fine-tuning experiment. Each model page links its Hugging Face weights and exact checkpoint command. ARES is separate and only needed for embedded deployment.
 
-BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind five model families.
+BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind six model families.
 
 The training stack is built on PyTorch Lightning and Hydra. Embedded deployment through ARES is maintained as a separate toolchain inside the repository.
 
@@ -19,6 +19,7 @@ The training stack is built on PyTorch Lightning and Hydra. Embedded deployment
 | [TinyMyo](docs/model/TinyMyo.md) | sEMG | Rotary Transformer | [Paper](https://arxiv.org/abs/2512.15729) / [Hugging Face](https://huggingface.co/PulpBio/TinyMyo) |
 | [LuMamba](docs/model/LuMamba.md) | EEG | Query-unified Mamba | [Paper](https://arxiv.org/abs/2603.19100) / [Hugging Face](https://huggingface.co/PulpBio/LuMamba) |
 | [PanLUNA](docs/model/PanLUNA.md) | EEG, ECG, PPG | Multimodal query-unified Transformer | [Paper](https://arxiv.org/abs/2604.04297) / [Hugging Face](https://huggingface.co/PulpBio/PanLUNA) |
+| [S-CEReBrO](docs/model/SCEReBrO.md) | EEG | Windowed alternating-attention Transformer | [Paper](https://arxiv.org/abs/2607.27913) / [Hugging Face](https://huggingface.co/PulpBio/S-CEReBrO) |
 
 The machine-readable [`model_registry.py`](biofoundation/model_registry.py) records the experiment names, papers, Hugging Face repositories, modalities, and batch metadata requirements for these families.
 
@@ -74,6 +75,7 @@ Choose another `+experiment` from the model registry. Before a long run, review
 | [`datasets`](datasets/) | Dataset readers and sample contracts. |
 | [`data_module`](data_module/) | Lightning data modules and loader composition. |
 | [`config`](config/) | Hydra defaults, modules, and reproducible experiments. |
+| [`docs/adr`](docs/adr/) | Architecture decision records for shared contracts. |
 | [`make_datasets`](make_datasets/) | Raw-data preprocessing and HDF5 conversion. |
 | [`criterion`](criterion/) | Training objectives. |
 | [`tests`](tests/) | Fast repository and refactoring contracts. |
@@ -105,6 +107,6 @@ Numerical changes to models, losses, or datasets should also be tested with repr
 
 ## Licensing and Support
 
-The source code is licensed under Apache 2.0. Pretrained weights in the five PulpBio Hugging Face repositories are licensed under CC BY-ND 4.0; see the model cards for terms and checkpoint-specific details.
+The source code is licensed under Apache 2.0. Pretrained weights in the six PulpBio Hugging Face repositories are licensed under CC BY-ND 4.0; see the model cards for terms and checkpoint-specific details.
 
 For questions and support, open an [issue](https://github.com/pulp-bio/BioFoundation/issues). For changes, start with [CONTRIBUTING.md](CONTRIBUTING.md).

biofoundation/__init__.py

@@ -17,5 +17,21 @@
 #* Author:  BioFoundation Contributors                                       *
 #*----------------------------------------------------------------------------*
 
-"""Shared infrastructure for the BioFoundation model zoo."""
+"""Shared infrastructure for the BioFoundation model zoo.
+
+This package holds the contracts that model families, tasks, and datasets agree on:
+the batch layout, the encoder and prediction-head protocols, checkpoint entry points,
+environment validation, and the model registry. Everything here is imported by code
+that must run without PyTorch installed, so this package has no heavyweight runtime
+dependencies.
+
+Changes follow an additive rule. New fields, functions, and classes may be added;
+existing names do not change signature or behaviour, and defaults are chosen so that
+an existing caller observes no difference. A change that cannot be made additively is
+a major version bump and needs a deprecation period first.
+"""
+
+__version__ = "0.2.0"
+
+__all__ = ["__version__"]
 

biofoundation/core/__init__.py

@@ -20,10 +20,15 @@
 """Stable contracts shared by model-specific implementations."""
 
 from biofoundation.core.batch import BatchRequirements, SignalBatch, as_signal_batch, require_batch_fields
+from biofoundation.core.checkpoints import SafetensorsCheckpointMixin
+from biofoundation.core.protocols import PredictionHead, SignalEncoder
 
 __all__ = [
     "BatchRequirements",
+    "PredictionHead",
+    "SafetensorsCheckpointMixin",
     "SignalBatch",
+    "SignalEncoder",
     "as_signal_batch",
     "require_batch_fields",
 ]

biofoundation/core/batch.py

@@ -29,24 +29,55 @@ class SignalBatch(TypedDict, total=False):
 
     Only ``input`` is universally required. Model adapters can require channel
     or sensor metadata without forcing simpler models to manufacture it.
+
+    Two independent electrode-geometry representations are supported, and a model
+    requires exactly one of them. They are peers rather than alternatives to convert
+    between, so a dataset declares which one it produces and the registry records
+    which one each model consumes:
+
+    ``channel_locations``
+        Shape ``(batch, channels, 3)``. One 3D coordinate per channel. For a bipolar
+        derivation this is the midpoint of the two electrodes. Consumed by LUNA,
+        LuMamba and PanLUNA.
+
+    ``channel_coords``
+        Shape ``(batch, channels, 2, 3)``. Both electrodes of every channel kept
+        separate, so a bipolar pair and a scalp-plus-reference channel stay
+        distinguishable. Consumed by S-CEReBrO.
+
+    Padding metadata describes how much of a sample is filler, which lets montages of
+    different sizes share one batch. It is only meaningful on mapping-shaped batches;
+    the tuple form accepted by :func:`as_signal_batch` cannot carry it.
     """
 
     input: Any
     label: Any
     channel_names: Any
     channel_locations: Any
+    channel_coords: Any
     sensor_type: Any
+    num_padded_channels: Any
+    num_padded_timesteps: Any
     metadata: Mapping[str, Any]
 
 
 @dataclass(frozen=True)
 class BatchRequirements:
-    """Metadata fields required by a particular model adapter."""
+    """Metadata fields required by a particular model adapter.
+
+    Every field defaults to ``False``, so a model only declares what it actually
+    reads. ``channel_locations`` and ``channel_coords`` are the two electrode-geometry
+    representations described on :class:`SignalBatch`; a model sets one of them.
+    """
 
     label: bool = False
     channel_names: bool = False
     channel_locations: bool = False
     sensor_type: bool = False
+    # Appended after the original fields so positional construction keeps its meaning.
+    channel_coords: bool = False
+    num_padded_channels: bool = False
+    num_padded_timesteps: bool = False
 
 
 def as_signal_batch(batch: Any) -> SignalBatch:
@@ -81,6 +112,9 @@ def require_batch_fields(batch: SignalBatch, requirements: BatchRequirements) ->
         ("channel_names", requirements.channel_names),
         ("channel_locations", requirements.channel_locations),
         ("sensor_type", requirements.sensor_type),
+        ("channel_coords", requirements.channel_coords),
+        ("num_padded_channels", requirements.num_padded_channels),
+        ("num_padded_timesteps", requirements.num_padded_timesteps),
     )
     missing = [name for name, required in required_fields if required and name not in batch]
     if missing:

biofoundation/core/checkpoints.py

@@ -0,0 +1,69 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2026 ETH Zurich, Switzerland                                 *
+#* SPDX-License-Identifier: Apache-2.0                                        *
+#*                                                                            *
+#* Licensed under the Apache License, Version 2.0 (the "License");            *
+#* you may not use this file except in compliance with the License.           *
+#* You may obtain a copy of the License at                                    *
+#*                                                                            *
+#* http://www.apache.org/licenses/LICENSE-2.0                                 *
+#*                                                                            *
+#* Unless required by applicable law or agreed to in writing, software        *
+#* distributed under the License is distributed on an "AS IS" BASIS,          *
+#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
+#* See the License for the specific language governing permissions and        *
+#* limitations under the License.                                             *
+#*                                                                            *
+#* Author:  BioFoundation Contributors                                       *
+#*----------------------------------------------------------------------------*
+
+"""Checkpoint entry points shared by tasks that separate encoder from head.
+
+``run_train.py`` calls ``load_pretrained_checkpoint`` for a ``.ckpt`` and
+``load_safetensors_checkpoint`` for a ``.safetensors``; this mixin maps both onto a
+task's single ``load_weights``, leaving tensor matching to the task.
+
+``load_weights`` is not called ``load_from_checkpoint`` because that name is a
+:class:`pytorch_lightning.LightningModule` classmethod, and an instance method of the
+same name shadows it so ``MyTask.load_from_checkpoint(path)`` binds the path to ``self``.
+"""
+
+from typing import Any
+
+
+class SafetensorsCheckpointMixin:
+    """Expose BioFoundation's two checkpoint entry points over one ``load_weights``.
+
+    Mix in ahead of :class:`pytorch_lightning.LightningModule`. The safetensors path
+    wraps the flat mapping in ``{"state_dict": ...}`` and prefixes bare keys with
+    ``model.``, matching how the pre-training task saves an encoder.
+    """
+
+    def load_pretrained_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any:
+        """Load a Lightning ``.ckpt`` by delegating to ``load_weights``."""
+
+        return self.load_weights(checkpoint_path=model_ckpt, **kwargs)
+
+    def load_safetensors_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any:
+        """Load a ``.safetensors`` file through the same path as a Lightning checkpoint."""
+
+        import tempfile
+        from pathlib import Path
+
+        import torch
+        from safetensors.torch import load_file
+
+        state_dict = {
+            key if key.startswith(("model.", "model_head.")) else f"model.{key}": value
+            for key, value in load_file(model_ckpt).items()
+        }
+
+        with tempfile.TemporaryDirectory() as directory:
+            converted = Path(directory) / "converted.ckpt"
+            torch.save({"state_dict": state_dict}, converted)
+            return self.load_weights(checkpoint_path=str(converted), **kwargs)
+
+
+
+
+__all__ = ["SafetensorsCheckpointMixin"]

biofoundation/core/protocols.py

@@ -0,0 +1,83 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2026 ETH Zurich, Switzerland                                 *
+#* SPDX-License-Identifier: Apache-2.0                                        *
+#*                                                                            *
+#* Licensed under the Apache License, Version 2.0 (the "License");            *
+#* you may not use this file except in compliance with the License.           *
+#* You may obtain a copy of the License at                                    *
+#*                                                                            *
+#* http://www.apache.org/licenses/LICENSE-2.0                                 *
+#*                                                                            *
+#* Unless required by applicable law or agreed to in writing, software        *
+#* distributed under the License is distributed on an "AS IS" BASIS,          *
+#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
+#* See the License for the specific language governing permissions and        *
+#* limitations under the License.                                             *
+#*                                                                            *
+#* Author:  BioFoundation Contributors                                       *
+#*----------------------------------------------------------------------------*
+
+"""Structural contracts for the encoder and prediction-head split.
+
+The five original families bundle their output layer into the model, selecting it from
+``num_classes``. Newer families separate an encoder that produces token embeddings from
+a prediction head that consumes them, each from its own Hydra group; the protocols
+below describe that second shape.
+
+Conformance is structural, so the bundled families stay valid unchanged. The signatures
+are concrete on purpose: a ``forward`` accepting ``*args, **kwargs`` would be satisfied
+by every :class:`torch.nn.Module` and check nothing.
+
+Tensors are annotated only under :data:`typing.TYPE_CHECKING`, because the fast contract
+suite imports this module without PyTorch installed.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from torch import Tensor
+else:  # pragma: no cover - runtime fallback keeps this module torch-free
+    Tensor = Any
+
+
+@runtime_checkable
+class SignalEncoder(Protocol):
+    """A model that turns a patched biosignal into contextualised token embeddings.
+
+    The encoder owns tokenisation, positional and channel embeddings, and the
+    backbone. It owns no task-specific output layer, so one pre-trained encoder can
+    be paired with any :class:`PredictionHead` without being rebuilt.
+
+    Implementations accept ``(batch, channels, patches, patch_size)`` and return
+    ``(batch, channels * patches, embed_dim)``. Keyword arguments beyond the input
+    carry whatever batch metadata the family requires, declared in its
+    :class:`~biofoundation.core.batch.BatchRequirements`.
+    """
+
+    def forward(self, x: "Tensor", channel_positions: "Tensor") -> "Tensor":
+        """Encode a patched biosignal into token embeddings.
+
+        Args:
+            x: ``(batch, channels, patches, patch_size)``.
+            channel_positions: Electrode geometry, shaped as the family's
+                :class:`~biofoundation.core.batch.BatchRequirements` declares.
+        """
+        ...
+
+
+@runtime_checkable
+class PredictionHead(Protocol):
+    """A module that turns encoder token embeddings into a task prediction.
+
+    Heads accept ``(batch, num_tokens, embed_dim)``. The output shape is the head's
+    own concern: class logits, a scalar per window, or one reconstructed patch per
+    token. A head never reads the raw waveform and never receives batch metadata, so
+    swapping the task means swapping the head alone.
+    """
+
+    def forward(self, x: "Tensor") -> "Tensor":
+        """Map ``(batch, num_tokens, embed_dim)`` token embeddings to a prediction."""
+        ...
+
+
+__all__ = ["PredictionHead", "SignalEncoder"]