diff --git a/.github/workflows/core-smoke.yml b/.github/workflows/core-smoke.yml
new file mode 100644
index 0000000..f067389
--- /dev/null
+++ b/.github/workflows/core-smoke.yml
@@ -0,0 +1,46 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+name: Core smoke tests
+
+on:
+ pull_request:
+ push:
+ branches: [main]
+
+permissions:
+ contents: read
+
+jobs:
+ contracts:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions/setup-python@v5
+ with:
+ python-version: "3.11"
+ - name: Install configuration test dependencies
+ run: python -m pip install hydra-core
+ - name: Run contract and configuration tests
+ run: python -m unittest discover -s tests -p "test_*.py" -v
+ - name: Compile Python sources
+ run: >-
+ python -m compileall -q
+ biofoundation run_train.py models tasks datasets data_module
+ criterion schedulers util make_datasets
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..acdaad4
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.coverage
+htmlcov/
+.DS_Store
+
+# Local environments and experiment output
+.venv/
+venv/
+outputs/
+checkpoints/
+wandb/
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 0000000..74592ef
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,48 @@
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE for details.
+
+# Contributing to BioFoundation
+
+BioFoundation is both a research repository and an onboarding codebase. Contributions should keep published experiments reproducible while making shared behavior easier to discover.
+
+## Repository Contracts
+
+- Keep changes scoped to one model, dataset, or shared contract where practical.
+- 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.
+- 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/`.
+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.
+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.
+
+## Adding Data or Training Components
+
+1. Put dataset implementations in `datasets/` and preprocessing scripts in `make_datasets/`.
+2. Compose datasets through a Lightning data module in `data_module/` and a matching `config/data_module/` file.
+3. Put reusable losses, tasks, and schedulers in their matching source and config directories.
+4. Document data layout, labels, channel or sensor metadata, and a reproducible command.
+
+## Pull Request Checklist
+
+- Explain the research or engineering reason for the change.
+- Include the exact configs, seeds, environment, and hardware needed to reproduce reported results.
+- Report the evaluation protocol, splits, task metrics, and relevant efficiency metrics.
+- Add focused tests and run the fast suite (`hydra-core` enables experiment composition checks):
+
+```bash
+python -m unittest discover -s tests -p 'test_*.py' -v
+python -m compileall -q biofoundation run_train.py models tasks datasets data_module
+```
+
+## Official Checkpoint Improvements
+
+The published weights are licensed under CC BY-ND 4.0. Modified weights, including adapters, deltas, pruned variants, and quantized variants, may not be redistributed.
+
+To propose an improvement as an official release, open a PR with the implementation, configs, seeds, logs, environment, hardware, evaluation protocol, and full metrics. Maintainers will review accepted changes, retrain and validate them, and publish the resulting checkpoint to the corresponding PulpBio Hugging Face repository.
diff --git a/README.md b/README.md
index 56b9245..e6e6de6 100644
--- a/README.md
+++ b/README.md
@@ -1,532 +1,110 @@
# BioFoundation
-
-
-Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE file for details.
-
-Authors: Thorir Mar Ingolfsson, Anna Tegon, Berkay Döner, Xiaying Wang, Matteo Fasulo, Danaé Broustail, Marija Zelic, Yawei Li & Luca Benini.
-
-## About
-
-**BioFoundation** is a flexible and extensible codebase for deep learning with biological signals. This repository is designed to support a variety of research projects, and currently hosts the work of multiple papers on EEG analysis.
-
-This repository is built on PyTorch Lightning and Hydra to enable reproducible and scalable research.
-
-## 🤗 Pretrained Weights on Hugging Face
-
-Looking for ready-to-use weights of models? We host them on Hugging Face:
-
-### Currently available
-
-- **FEMBA** ([paper](https://arxiv.org/abs/2502.06438)) [](https://huggingface.co/PulpBio/FEMBA)
-- **LUNA** ([paper](https://arxiv.org/abs/2510.22257)) [](https://huggingface.co/PulpBio/LUNA)
-- **TinyMyo** ([paper](https://arxiv.org/abs/2512.15729)) [](https://huggingface.co/PulpBio/TinyMyo)
-- **LuMamba** ([paper](https://arxiv.org/abs/2603.19100)) [](https://huggingface.co/PulpBio/LuMamba)
-- **PanLUNA** ([paper](https://arxiv.org/pdf/2604.04297)) [](https://huggingface.co/PulpBio/PanLUNA)
-
-#### Why FEMBA?
-
-- **Scales to long EEG** with linear-time Mamba (no quadratic attention).
-- **Strong results** on TUAB/TUAR/TUSL with ready task-specific checkpoints.
-- **Simple fine-tune path:** set `CHECKPOINT_DIR`, run `+experiment=FEMBA_finetune`.
-
-**➡️ Model hub:**
-**📄 Model card:** [FEMBA on Hugging Face](https://huggingface.co/PulpBio/FEMBA) — benchmarks, protocols, and efficiency notes.
-**📜 Weights license:** CC BY-ND 4.0 (use + redistribute **unmodified** weights with attribution; no redistribution of **modified** weights)
-**🧑🍳 PR-gated improvements:** If you fine-tune internally and want your variant to become an **official** FEMBA release, open a PR with configs, logs, and evals. We’ll review together; if it looks good, we’ll retrain/validate and publish an **official** FEMBA checkpoint.
-
-**What you’ll find on the hub**
-
-- `TUAB/` → abnormal EEG (base/large)
-- `TUAR/` → artifact detection (tiny/base/large)
-- `TUSL/` → slowing classification (variants as in the paper)
-
-Quick download with `huggingface_hub`:
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See [LICENSE](LICENSE) for details.
-```bash
-pip install huggingface_hub
-```
-
-```python
-from huggingface_hub import snapshot_download
-
-# downloads all task folders (TUAB/TUAR/TUSL) and safetensors into ./checkpoints/FEMBA
-snapshot_download(repo_id="PulpBio/FEMBA", repo_type="model", local_dir="checkpoints/FEMBA")
-```
-
-Use the paths directly in your runs, e.g.:
-
-```bash
-export DATA_PATH=/path/to/data
-export CHECKPOINT_DIR=checkpoints/FEMBA/TUAR/base.safetensors
-python -u run_train.py +experiment=FEMBA_finetune
-```
-
-#### Why LUNA?
-
-- **Topology-agnostic** EEG via **query-based channel unification** (consistent latent across arbitrary montages).
-- **Linear-in-channels** compute & memory (unifies channels **before** temporal modeling; no quadratic spatio-temporal attention).
-- **Pretrained on >21k hours** (TUEG + Siena) with masked-patch reconstruction; strong transfer across datasets/montages.
-- **Simple fine-tune path:** pick model size with `LUNA_{base,large,huge}.yaml`, set `pretrained_safetensors_path`, run `+experiment=LUNA_finetune`.
-
-**➡️ Model hub:**
-**📄 Model card:** [LUNA on Hugging Face](https://huggingface.co/PulpBio/LUNA) — variants, configs, and fine-tuning walkthrough.
-**📜 Weights license:** CC BY-ND 4.0 (use + redistribute **unmodified** weights with attribution; no redistribution of **modified** weights)
-**🧑🍳 PR-gated improvements:** If you fine-tune internally and want your variant to become an **official** LUNA release, open a PR with configs, logs, and evals. We’ll review; if it looks good, we’ll retrain/validate and publish an **official** LUNA checkpoint.
-
-**What you’ll find on the hub**
-
-- `Base/`, `Large/`, `Huge/` → LUNA size variants (matching `config/model/LUNA_{base,large,huge}.yaml`)
-- Task-specific heads/checkpoints for common TUH downstream tasks (TUAB / TUAR / TUSL)
-
-Quick download with `huggingface_hub`:
-
-```bash
-pip install huggingface_hub
-```
-
-```python
-from huggingface_hub import snapshot_download
-
-# downloads LUNA folders and .safetensors into ./checkpoints/LUNA
-snapshot_download(repo_id="PulpBio/LUNA", repo_type="model", local_dir="checkpoints/LUNA")
-```
-
-Use the paths directly in your runs like here below:
-
-```bash
-python -u run_train.py +experiment=LUNA_finetune /model=LUNA_base \
- pretrained_safetensors_path=/absolute/path/to/checkpoints/LUNA/Base/LUNA_base.safetensors
-
-python -u run_train.py +experiment=LUNA_finetune /model=LUNA_large \
- pretrained_safetensors_path=/absolute/path/to/checkpoints/LUNA/Large/LUNA_large.safetensors
-
-python -u run_train.py +experiment=LUNA_finetune /model=LUNA_huge \
- pretrained_safetensors_path=/absolute/path/to/checkpoints/LUNA/Huge/LUNA_huge.safetensors
-```
-
-*If your checkpoint path contains spaces, wrap it in quotes.*
-
-Tips:
-
-- TUH datasets (TUAB/TUAR/TUSL): keep `- override /data_module: finetune_data_module` and set `data_module.*.hdf5_file` to your `{train,val,test}.h5`.
-- Non-TUH (e.g., SEED-V): use `- override /data_module: subject_independent_data_module` and remove the TUH-specific `data_module` block.
-- Match task settings: `classification_type` (`bc`, `mc`, `mmc`, `mcc`) and `model.num_classes` (e.g., TUSL=4, TUAB=2).
-
-#### Why TinyMyo?
+Authors: Thorir Mar Ingolfsson, Anna Tegon, Berkay Döner, Xiaying Wang, Matteo Fasulo, Danaé Broustail, Marija Zelic, Yawei Li, and Luca Benini.
-- **Ultra-lightweight**: only 3.6M parameters, suitable for microcontroller deployment.
-- **Broad generalization**: pretrained on multiple large-scale EMG datasets for versatility across tasks and sensor configurations.
-- **Strong results** on surface EMG tasks with ready task-specific checkpoints.
+> **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=_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.
-**➡️ Model hub:**
-**📄 Model card:** [TinyMyo on Hugging Face](https://huggingface.co/PulpBio/TinyMyo) — benchmarks, protocols, and efficiency notes.
-**📜 Weights license:** CC BY-ND 4.0 (use + redistribute **unmodified** weights with attribution; no redistribution of **modified** weights)
-**🧑🍳 PR-gated improvements:** If you fine-tune internally and want your variant to become an **official** TinyMyo release, open a PR with configs, logs, and evals. We’ll review together; if it looks good, we’ll retrain/validate and publish an **official** TinyMyo checkpoint.
+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.
-**What you’ll find on the hub**
+The training stack is built on PyTorch Lightning and Hydra. Embedded deployment through ARES is maintained as a separate toolchain inside the repository.
-- `DB5/` → gesture classification
-- `UCI_EMG/` → gesture classification
-- `EPN612/` → gesture classification
+## Model Zoo
-> The scripts to download and preprocess the datasets are available at:
+| Model | Signals | Architecture | Resources |
+| --- | --- | --- | --- |
+| [FEMBA](docs/model/FEMBA.md) | EEG | Bidirectional Mamba | [Paper](https://arxiv.org/abs/2502.06438) / [Hugging Face](https://huggingface.co/PulpBio/FEMBA) |
+| [LUNA](docs/model/LUNA.md) | EEG | Query-unified Transformer | [Paper](https://arxiv.org/abs/2510.22257) / [Hugging Face](https://huggingface.co/PulpBio/LUNA) |
+| [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) |
-Quick download with `huggingface_hub`:
+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.
-```bash
-pip install huggingface_hub
-```
-
-```python
-from huggingface_hub import snapshot_download
-
-# downloads all task folders (DB5/UCI_EMG/EPN612) and safetensors into ./checkpoints/TinyMyo
-snapshot_download(repo_id="MatteoFasulo/TinyMyo", repo_type="model", local_dir="checkpoints/TinyMyo")
-```
-
-Use the paths directly in your runs, e.g.:
-
-```bash
-export DATA_PATH=/path/to/data
-export CHECKPOINT_DIR=checkpoints/TinyMyo/UCI_EMG/base.safetensors
-python -u run_train.py +experiment=TinyMyo_finetune \
- pretrained_safetensors_path=/path/to/model.safetensors
-```
-
-**What you won’t find on the hub**
-
-- **Silent Speech**
- - Codebase: [MatteoFasulo/silent_speech](https://github.com/MatteoFasulo/silent_speech)
-- **Generic Neuromotor Interface**
- - Codebase: [MatteoFasulo/generic-neuromotor-interface](https://github.com/MatteoFasulo/generic-neuromotor-interface)
-
-#### Why LuMamba?
-
-- **Merging topology-invariant** EEG representations with **Mamba-based** temporal modeling to jointly achieve **linear-time efficiency** and **channel invariance**.
-- **Pretrained on >21k hours** (TUEG) with **LeJEPA** (Balestriero and LeCun, 2025), a recent joint-embedding predictive approach, adapted to biosignals in this repository to enhance cross-montage robustness.
-
-**➡️ Model hub:**
-**📄 Model card:** [LuMamba on Hugging Face](https://huggingface.co/PulpBio/LuMamba) — variants, configs, and fine-tuning walkthrough.
-**📜 Weights license:** CC BY-ND 4.0 (use + redistribute **unmodified** weights with attribution; no redistribution of **modified** weights)
-**🧑🍳 PR-gated improvements:** If you fine-tune internally and want your variant to become an **official** LuMamba release, open a PR with configs, logs, and evals. We’ll review; if it looks good, we’ll retrain/validate and publish an **official** LuMamba checkpoint.
-
-**What you’ll find on the hub**
-
-- `LeJEPA-only`, `Reconstruction-only`, `Mixed LeJEPA-Reconstruction` pre-trained checkpoints → Pre-training design variants.
-- Instructions to get started on fine-tuning experiments.
+## Quick Start
-Quick download with `huggingface_hub`:
-
-```bash
-pip install huggingface_hub
-```
-
-```python
-from huggingface_hub import snapshot_download
-
-# downloads all pre-trained variants and safetensors into ./checkpoints/LuMamba
-snapshot_download(repo_id="PulpBio/LuMamba", repo_type="model", local_dir="checkpoints/LuMamba")
-```
-Include the safetensors checkpoint path as input and run fine-tuning in the commandline:
-```bash
-python -u run_train.py +experiment=LuMamba_finetune \
- pretrained_safetensors_path=/absolute/path/to/checkpoints/LuMamba/LuMamba.safetensors
-```
-
-Tips:
-- Data preprocessing scripts are provided in `/make_datasets` for various downstream datasets. See `make_datasets/README.md` for instructions on getting started.
-- Adapt configuration file `config/experiment/LuMamba_finetune.yaml` to your specific task with correct dataset paths, classification type (regression and multi-class classification `mcc`, binary `bc` and change `model.num_classes` accordingly), I/O settings, trainer parameters, etc.
- - Ensure `data_module:train/test/val` are initialized with the compatible dataset class.
- - Configuration file includes sufficient `#CHANGEME` tags and further instructions for a working example.
-
-#### Why PanLUNA?
-
-* Extending LUNA's channel-unification mechanism from topology invariance to **cross-modal fusion**, jointly processing EEG, ECG, and PPG within a single shared encoder via sensor-type embeddings - no modality-specific backbones, no paired multimodal data required during pretraining.
-* Pretrained on ~40,000 hours of heterogeneous biosignal data (TUEG, Siena, MIMIC-IV, CODE-15%, PulseDB) with a masked signal reconstruction objective.
-* Strong performance on unimodal and multimodal experiments.
-
-➡️ Model hub: __https://huggingface.co/PulpBio/PanLUNA__ 📄 Model card: __[PanLUNA on Hugging Face](https://huggingface.co/PulpBio/PanLUNA)__ — variants, configs, and fine-tuning walkthrough. 📜 Weights license: CC BY-ND 4.0 (use + redistribute unmodified weights with attribution; no redistribution of modified weights) 🧑🍳 PR-gated improvements: If you fine-tune internally and want your variant to become an official PanLUNA release, open a PR with configs, logs, and evals. We'll review; if it looks good, we'll retrain/validate and publish an official PanLUNA checkpoint.
-
-**What you'll find on the hub**
-
-* Pre-trained checkpoint.
-* Instructions to get started on fine-tuning experiments.
-
-Quick download with `huggingface_hub`:
-
-```
-pip install huggingface_hub
-```
-
-```python
-from huggingface_hub import snapshot_download
-
-# downloads all pre-trained variants and safetensors into ./checkpoints/PanLUNA
-snapshot_download(repo_id="PulpBio/PanLUNA", repo_type="model", local_dir="checkpoints/PanLUNA")
-```
-
-Include the safetensors checkpoint path as input and run fine-tuning in the commandline:
-
-```bash
-python -u run_train.py +experiment=PanLUNA_finetune \
- pretrained_safetensors_path=/absolute/path/to/checkpoints/PanLUNA/PanLUNA.safetensors
-```
-
-Tips:
-
-* Data preprocessing scripts are provided in `/make_datasets` for various downstream datasets. Download the corresponding dataset, locate its preprocessing script by name matching, and adjust key parameters.
-* Adapt configuration file `config/experiment/PanLUNA_finetune.yaml` to your specific task with correct data module (for unimodal experiments `config/data_module/finetune_data_module_unimodal_PanLUNA` or multimodal experiments `config/data_module/finetune_data_module_multimodal_PanLUNA`), classification type (binary `bc`, multi-class `mcc` and mulit-label `mlp`) and change `model.num_classes` accordingly. For different fine-tuning strategies adjust `finetuning.mode` parameter with `lora` for Low-Rank Adaptation, `freeze_encoder` for frozen backbone or select `full` for complete adaptation.
- * Ensure `data_module:train/test/val` are initialized with the compatible dataset class. Leverage `config/data_module/finetune_data_module_multimodal_PanLUNA:data_module.train.channel_groups` to specify channels available in the dataset and `channel_start`/`channel_end` to tweak the used subset.
- * Configuration file includes sufficient `#CHANGEME` tags and further instructions for a working example.
-* Results on the addtional cross-modal experiments provided in the `docs/model/PanLUNA.md`.
-
-## Features
-
-- **Modular Design**: The repository is organized into modules for data loading, models, training tasks, and more, making it easy to extend and adapt for new research projects.
-- **Flexible Configuration**: We use [Hydra](https://hydra.cc/docs/intro/) to manage experiment configurations, allowing for easy customization of models, data, and training parameters.
-- **Reproducibility**: Our use of `Hydra` and PyTorch Lightning helps ensure that our experiments are reproducible.
-- **Extensible**: The repository is designed to be easily extended with new datasets, models, and tasks.
-
-## Installation
-
-To use BioFoundation, clone the repository and install the required dependencies.
+BioFoundation requires Python 3.11 or newer. Create an isolated environment and install the training dependencies:
```bash
git clone https://github.com/pulp-bio/BioFoundation.git
-```
-
-We recommend using a virtual environment to manage dependencies. You can use `conda` or `virtualenv` for this purpose. We have provided a `requirements.txt` file that lists the necessary packages. You can install them using pip, and optionally, you can use `conda` to create a new environment.
-
-```bash
-conda create -n BioFoundation
-conda activate BioFoundation
+cd BioFoundation
+conda create -n biofoundation python=3.11
+conda activate biofoundation
pip install -r requirements.txt
```
-alternatively, you can create a venv using [`uv`](https://docs.astral.sh/uv/) and then install via `pyproject.toml` with:
+With `uv`, the equivalent setup is:
```bash
+uv venv --python 3.11
+source .venv/bin/activate
uv pip install -r pyproject.toml --torch-backend=auto
```
-where `--torch-backend=auto` will query for installed CUDA driver, AMD GPU versions, and Intel GPU presence and then use the most-compatible PyTorch index for all the torch-related packages. Have a look at [uv docs](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) for more details.
-
-### Path changes
-
-Throughout the repository, you may find paths that need to be adjusted based on your local setup. For example, the path to the datasets in the configuration files or the scripts that process the datasets. Make sure to update these paths accordingly. They have been named "#CHANGEME" to facilitate finding them.
-
-## Dataset Preparation
-
-The datasets used in this repository should be converted to HDF5 for efficient I/O. Other formats can work, but you’d need to adapt the dataloaders accordingly.
-To prepare the TUH EEG datasets (see the [official source](https://isip.piconepress.com/projects/nedc/html/tuh_eeg/index.shtml)), follow these steps:
-
-1. **Download raw data** from the official sources (e.g., TUH EEG corpus).
-2. **Preprocess to pickles** (windowing/labels):
-
- ```bash
- # examples (adjust paths)
- python make_datasets/process_raw_eeg.py tuab --root_dir /eeg_data/TUAB/edf --output_dir /processed_eeg
- python make_datasets/process_raw_eeg.py tusl --root_dir /eeg_data/TUSL/edf --output_dir /processed_eeg
- python make_datasets/process_raw_eeg.py tuar --root_dir /eeg_data/TUAR/edf --output_dir /processed_eeg
- ```
-
-3. **Bundle into HDF5:**: Use the provided script to process the raw data into HDF5 files.
-
- ```bash
- # all datasets found under /processed_eeg
- python make_datasets/make_hdf5.py --prepath /processed_eeg --dataset All --remove_pkl
-
- # or a single dataset
- python make_datasets/make_hdf5.py --prepath /processed_eeg --dataset TUSL --remove_pkl
- ```
-
- You may need to edit the `prepath` variable in the script to point to the directory where you have downloaded the raw data.
-4. **Update Configs**: so `data_module.*.hdf5_file` points to your `${DATA_PATH}/_data/{train,val,test}.h5`
-
-## How to Run
-
-### Pre-training
-
-To run a pre-training experiment, you can use the `run_train.py` script with the appropriate configuration file. For example in the case of pre-training FEMBA:
+Set the prepared data and experiment-output roots:
```bash
-python -u run_train.py +experiment=FEMBA_pretrain
-
+export DATA_PATH=/absolute/path/to/data
+export CHECKPOINT_DIR=/absolute/path/to/experiments
```
-### Fine-tuning
-
-To run a fine-tuning experiment, you can use the `run_train.py` script with the appropriate configuration file. For example in the case of fine-tuning FEMBA:
+Start a pre-training experiment:
```bash
-python -u run_train.py +experiment=FEMBA_finetune
-
-```
-
-> **Tip:** Pretrained FEMBA weights (TUAB/TUAR/TUSL folders) are available on 🤗 Hugging Face:
->
-> Set `CHECKPOINT_DIR` to the desired `.safetensors` (e.g., `.../TUAR/base.safetensors`) before launching.
-
-Note in both cases one needs to make sure that the dataset that specific experiment is using is downloaded and available in the correct path.
-
-## Repository Structure
-
-```
-BioFoundation/
-├── config # Hydra configuration files
-├── criterion # Loss functions
-├── data_module # PyTorch Lightning DataModules
-├── datasets # PyTorch Datasets
-├── docs # Detailed documentation
-├── models # Model implementations
-├── schedulers # Learning rate schedulers
-├── tasks # PyTorch Lightning tasks
-└── ...
+python -u run_train.py +experiment=FEMBA_pretrain
```
-## Contributing
-
-We welcome contributions to BioFoundation! If you have a new model, dataset, or task that you would like to add, please follow the guidelines below.
-
-### How to add a new dataset?
-
-1. Add the code of the dataset to [`datasets`](datasets/).
-2. Add the configuration file of the dataset to [`./config/dataset`](./config/dataset/).
-3. If the dataset is large, consider adding a script to download it in the [`./scripts`](./scripts) directory. Make sure to document how to run the script in the README.
-
-### How to add a new data module?
-
-1. Add the code of the data module to [`./data_module`](./data_module).
-2. Add the configuration file of the data module to [`./config/data_module`](./config/data_module).
-3. If the data module requires specific datasets, make sure to document how to download and prepare them in the README.
-
-### How to add a new loss function?
-
-1. Add the code of the loss function to [`./criterion`](./criterion).
-2. Add the configuration file of the loss function to [`./config/criterion`](./config/criterion).
-
-### How to add a new task?
-
-1. Add the code of the task to [`./tasks`](./tasks).
-2. Add the configuration file of the task to [`./config/task`](./config/task).
-3. If the task requires specific datasets or models, make sure to document how to download and prepare them in the README.
-
-### How to add a new scheduler?
-
-1. Add the code of the scheduler to [`./schedulers`](./schedulers).
-2. Add the configuration file of the scheduler to [`./config/scheduler`](./config/scheduler).
-3. If the scheduler requires specific models or tasks, make sure to document how to use it in the README.
-
-### How to add a new model?
-
-1. Add the code of the model to [`./models`](./models).
-2. Add the configuration file of the model to [`./config/model`](./config/model).
-
-### How to start a new experiment with the added model?
-
-1. Add experiment configuration file to [`./config/experiment`](./config/experiment).
- If you are interested, you may check the [Hydra document about it](https://hydra.cc/docs/patterns/configuring_experiments/).
-2. Override the default configurations in the added experiment configuration file.
-3. Run the experiment with the command:
+Or fine-tune a downloaded checkpoint:
```bash
-python -u run_train.py +experiment=your_experiment_name
+python -u run_train.py +experiment=LUNA_finetune /model=LUNA_base \
+ pretrained_safetensors_path=/absolute/path/to/LUNA_base.safetensors
```
-### Contributing improvements to FEMBA weights
+Choose another `+experiment` from the model registry. Before a long run, review the selected file in [`config/experiment`](config/experiment/) and resolve its `#CHANGEME` values.
-We’re excited to see what you build. Because the weights are **CC BY-ND 4.0**, redistribution of **modified** weights (e.g., LoRA/adapters, deltas, pruned or quantized variants) is **not permitted**.
-If you fine-tune internally and believe your results should become an **official** FEMBA release, please open a PR with:
+## Repository Map
-- exact **configs**, **seeds**, and **training scripts**,
-- **environment** and **hardware** details,
-- **evaluation protocol** (TUAB/TUAR/TUSL), **splits**, and full **metrics** (AUROC/AUPR/BA, FLOPs, memory),
-- training and validation **logs**.
+| Path | Responsibility |
+| --- | --- |
+| [`biofoundation`](biofoundation/) | Shared batch, environment, and model metadata contracts. |
+| [`models`](models/) | Foundation model implementations. |
+| [`tasks`](tasks/) | Lightning pre-training, classification, and regression tasks. |
+| [`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. |
+| [`make_datasets`](make_datasets/) | Raw-data preprocessing and HDF5 conversion. |
+| [`criterion`](criterion/) | Training objectives. |
+| [`tests`](tests/) | Fast repository and refactoring contracts. |
+| [`ARES`](ARES/) | Independent GAP9 and Siracusa deployment toolchain. |
-Maintainers will review; if accepted, we will retrain/validate and publish a new **official** checkpoint on 🤗 under the same license.
+## Documentation
-## General Tips
+- [Documentation index](docs/README.md)
+- [Training and batch contracts](docs/TRAINING.md)
+- [Hydra configuration guide](config/README.md)
+- [Dataset preparation](make_datasets/README.md)
+- [Contribution guide](CONTRIBUTING.md)
+- [Citations](docs/CITATIONS.md)
+- [ARES deployment](ARES/README.md)
-### How to use distributed data parallel?
+Each model page linked from the model zoo contains its input assumptions, architecture, results, Hugging Face download, and fine-tuning example.
-In your experiment configuration file, add the following arguments
+## Development Checks
-```yaml
-trainer:
- accelerator: gpu # Using GPU
- num_nodes: ${num_nodes} # The number of computing nodes
- devices: -1 # Automatically uses all GPUs available
- strategy: ddp # distributed data parallel
-```
+The fast suite checks model metadata, Hydra composition and targets, batch adapters, environment handling, documentation links, and Apache headers:
-### How to save GPU memory?
-
-1. Try fairscale checkpointing first. Check [here](https://fairscale.readthedocs.io/en/stable/api/nn/checkpoint/checkpoint_activations.html) and [here](https://github.com/ofsoundof/GRL-Image-Restoration/blob/main/models/networks/grl.py#L134)
-2. Use sharded training. Check [here](https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html).
-
-## Contact
-
-For questions and support, please open an issue on the GitHub repository.
-
-## Citing this Work
-
-If you find this work useful, please cite the respective papers:
-
-```bibtex
-@misc{tegon2025fembaefficientscalableeeg,
- title={FEMBA: Efficient and Scalable EEG Analysis with a Bidirectional Mamba Foundation Model},
- author={Anna Tegon and Thorir Mar Ingolfsson and Xiaying Wang and Luca Benini and Yawei Li},
- year={2025},
- eprint={2502.06438},
- archivePrefix={arXiv},
- primaryClass={cs.LG},
- url={https://arxiv.org/abs/2502.06438},
-}
-@inproceedings{doner2025luna,
- title={{LUNA}: Efficient and Topology-Agnostic Foundation Model for {EEG} Signal Analysis},
- author={Berkay D{\"o}ner and Thorir Mar Ingolfsson and Luca Benini and Yawei Li},
- booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
- year={2025},
- url={https://openreview.net/forum?id=uazfjnFL0G}
-}
-
-@misc{fasulo2025tinymyotinyfoundationmodel,
- title={TinyMyo: a Tiny Foundation Model for Flexible EMG Signal Processing at the Edge},
- author={Matteo Fasulo and Giusy Spacone and Thorir Mar Ingolfsson and Yawei Li and Luca Benini and Andrea Cossettini},
- year={2025},
- eprint={2512.15729},
- archivePrefix={arXiv},
- primaryClass={eess.SP},
- url={https://arxiv.org/abs/2512.15729},
-}
-@misc{broustail2026lumambalatentunifiedmamba,
- title={LuMamba: Latent Unified Mamba for Electrode Topology-Invariant and Efficient EEG Modeling},
- author={Danaé Broustail and Anna Tegon and Thorir Mar Ingolfsson and Yawei Li and Luca Benini},
- year={2026},
- eprint={2603.19100},
- archivePrefix={arXiv},
- primaryClass={cs.AI},
- url={https://arxiv.org/abs/2603.19100},
-}
-@misc{zelic2026panluna,
- title={PanLUNA: An Efficient and Robust Query-Unified Multimodal Model for Edge Biosignal Intelligence},
- author={Marija Zelic and Anna Tegon and Yawei Li and Thorir Mar Ingolfsson},
- year={2026},
- eprint={2604.04297},
- archivePrefix={arXiv},
- primaryClass={cs.AI},
- url={https://arxiv.org/abs/2604.04297},
-}
+```bash
+python -m unittest discover -s tests -p 'test_*.py' -v
+python -m compileall -q biofoundation run_train.py models tasks datasets data_module \
+ criterion schedulers util make_datasets tests
```
-## License
-
-This project is licensed under the Apache License 2.0. See the [LICENSE](./LICENSE) file for details.
-
-**Note on model weights:** Pretrained weights are hosted at , , , and and licensed under **CC BY-ND 4.0**. You may use and redistribute the **unmodified** weights with attribution. Redistribution of **modified** weights is not permitted. To upstream improvements, please open a PR; accepted changes will be released as **official** checkpoints.
+Numerical changes to models, losses, or datasets should also be tested with representative CPU or GPU batches.
+## Licensing and Support
-## Star History
+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.
-
-
-
-
-
-
-
+For questions and support, open an [issue](https://github.com/pulp-bio/BioFoundation/issues). For changes, start with [CONTRIBUTING.md](CONTRIBUTING.md).
diff --git a/biofoundation/__init__.py b/biofoundation/__init__.py
new file mode 100644
index 0000000..da9f3de
--- /dev/null
+++ b/biofoundation/__init__.py
@@ -0,0 +1,21 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+"""Shared infrastructure for the BioFoundation model zoo."""
+
diff --git a/biofoundation/core/__init__.py b/biofoundation/core/__init__.py
new file mode 100644
index 0000000..57e6448
--- /dev/null
+++ b/biofoundation/core/__init__.py
@@ -0,0 +1,30 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+"""Stable contracts shared by model-specific implementations."""
+
+from biofoundation.core.batch import BatchRequirements, SignalBatch, as_signal_batch, require_batch_fields
+
+__all__ = [
+ "BatchRequirements",
+ "SignalBatch",
+ "as_signal_batch",
+ "require_batch_fields",
+]
+
diff --git a/biofoundation/core/batch.py b/biofoundation/core/batch.py
new file mode 100644
index 0000000..5f44502
--- /dev/null
+++ b/biofoundation/core/batch.py
@@ -0,0 +1,89 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+"""Canonical batch handling across biosignal datasets and model families."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, TypedDict, cast
+
+
+class SignalBatch(TypedDict, total=False):
+ """Common batch fields used by BioFoundation tasks.
+
+ Only ``input`` is universally required. Model adapters can require channel
+ or sensor metadata without forcing simpler models to manufacture it.
+ """
+
+ input: Any
+ label: Any
+ channel_names: Any
+ channel_locations: Any
+ sensor_type: Any
+ metadata: Mapping[str, Any]
+
+
+@dataclass(frozen=True)
+class BatchRequirements:
+ """Metadata fields required by a particular model adapter."""
+
+ label: bool = False
+ channel_names: bool = False
+ channel_locations: bool = False
+ sensor_type: bool = False
+
+
+def as_signal_batch(batch: Any) -> SignalBatch:
+ """Normalize existing tensor, tuple, and mapping batches to one contract."""
+
+ if isinstance(batch, Mapping):
+ normalized = dict(batch)
+ elif isinstance(batch, (tuple, list)):
+ if len(batch) == 1:
+ normalized = {"input": batch[0]}
+ elif len(batch) == 2:
+ normalized = {"input": batch[0], "label": batch[1]}
+ else:
+ raise ValueError(
+ "Sequence batches must contain either input or (input, label); "
+ f"received {len(batch)} values."
+ )
+ else:
+ normalized = {"input": batch}
+
+ if "input" not in normalized:
+ raise ValueError("Signal batches must contain an 'input' field.")
+
+ return cast(SignalBatch, normalized)
+
+
+def require_batch_fields(batch: SignalBatch, requirements: BatchRequirements) -> SignalBatch:
+ """Validate model-specific batch requirements with a useful error message."""
+
+ required_fields = (
+ ("label", requirements.label),
+ ("channel_names", requirements.channel_names),
+ ("channel_locations", requirements.channel_locations),
+ ("sensor_type", requirements.sensor_type),
+ )
+ missing = [name for name, required in required_fields if required and name not in batch]
+ if missing:
+ raise ValueError(f"Signal batch is missing required fields: {', '.join(missing)}")
+ return batch
+
diff --git a/biofoundation/core/environment.py b/biofoundation/core/environment.py
new file mode 100644
index 0000000..304936f
--- /dev/null
+++ b/biofoundation/core/environment.py
@@ -0,0 +1,46 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+"""Environment validation shared by command-line entry points."""
+
+import os
+from collections.abc import Mapping, Sequence
+
+
+def missing_environment_variables(
+ names: Sequence[str],
+ environ: Mapping[str, str] | None = None,
+) -> tuple[str, ...]:
+ """Return required variables that are unset or still placeholders."""
+
+ values = os.environ if environ is None else environ
+ return tuple(name for name in names if not values.get(name) or values.get(name) == "#CHANGEME")
+
+
+def require_environment(
+ names: Sequence[str],
+ environ: Mapping[str, str] | None = None,
+) -> None:
+ """Raise a single actionable error for missing runtime environment values."""
+
+ missing = missing_environment_variables(names, environ)
+ if missing:
+ joined = ", ".join(missing)
+ raise RuntimeError(f"Required environment variables are not set: {joined}")
+
diff --git a/biofoundation/model_registry.py b/biofoundation/model_registry.py
new file mode 100644
index 0000000..7b801f2
--- /dev/null
+++ b/biofoundation/model_registry.py
@@ -0,0 +1,114 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+"""Discoverable metadata for every model family distributed in this repository."""
+
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Mapping
+
+from biofoundation.core.batch import BatchRequirements
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+ """Human- and machine-readable entry for a BioFoundation model family."""
+
+ display_name: str
+ modalities: tuple[str, ...]
+ architecture: str
+ model_target: str
+ pretrain_experiment: str
+ finetune_experiment: str
+ huggingface_url: str
+ paper_url: str
+ batch_requirements: BatchRequirements = BatchRequirements()
+
+
+MODEL_REGISTRY: Mapping[str, ModelSpec] = MappingProxyType(
+ {
+ "femba": ModelSpec(
+ display_name="FEMBA",
+ modalities=("EEG",),
+ architecture="Bidirectional Mamba",
+ model_target="models.FEMBA.FEMBA",
+ pretrain_experiment="FEMBA_pretrain",
+ finetune_experiment="FEMBA_finetune",
+ huggingface_url="https://huggingface.co/PulpBio/FEMBA",
+ paper_url="https://arxiv.org/abs/2502.06438",
+ ),
+ "luna": ModelSpec(
+ display_name="LUNA",
+ modalities=("EEG",),
+ architecture="Query-unified Transformer",
+ model_target="models.LUNA.LUNA",
+ pretrain_experiment="LUNA_pretrain",
+ finetune_experiment="LUNA_finetune",
+ huggingface_url="https://huggingface.co/PulpBio/LUNA",
+ paper_url="https://arxiv.org/abs/2510.22257",
+ batch_requirements=BatchRequirements(channel_locations=True),
+ ),
+ "tinymyo": ModelSpec(
+ display_name="TinyMyo",
+ modalities=("sEMG",),
+ architecture="Rotary Transformer",
+ model_target="models.TinyMyo.TinyMyo",
+ pretrain_experiment="TinyMyo_pretrain",
+ finetune_experiment="TinyMyo_finetune",
+ huggingface_url="https://huggingface.co/PulpBio/TinyMyo",
+ paper_url="https://arxiv.org/abs/2512.15729",
+ ),
+ "lumamba": ModelSpec(
+ display_name="LuMamba",
+ modalities=("EEG",),
+ architecture="Query-unified Mamba",
+ model_target="models.LuMamba.LuMamba",
+ pretrain_experiment="LuMamba_pretrain",
+ finetune_experiment="LuMamba_finetune",
+ huggingface_url="https://huggingface.co/PulpBio/LuMamba",
+ paper_url="https://arxiv.org/abs/2603.19100",
+ batch_requirements=BatchRequirements(channel_locations=True),
+ ),
+ "panluna": ModelSpec(
+ display_name="PanLUNA",
+ modalities=("EEG", "ECG", "PPG"),
+ architecture="Multimodal query-unified Transformer",
+ model_target="models.PanLUNA.PanLUNA",
+ pretrain_experiment="PanLUNA_pretrain",
+ finetune_experiment="PanLUNA_finetune",
+ huggingface_url="https://huggingface.co/PulpBio/PanLUNA",
+ paper_url="https://arxiv.org/abs/2604.04297",
+ batch_requirements=BatchRequirements(
+ channel_locations=True,
+ sensor_type=True,
+ ),
+ ),
+ }
+)
+
+
+def get_model_spec(name: str) -> ModelSpec:
+ """Return a model specification by its case-insensitive registry key."""
+
+ key = name.casefold()
+ try:
+ return MODEL_REGISTRY[key]
+ except KeyError as error:
+ available = ", ".join(sorted(MODEL_REGISTRY))
+ raise KeyError(f"Unknown model '{name}'. Available models: {available}") from error
diff --git a/config/data_module/dataset_types.yaml b/config/data_module/dataset_types.yaml
index 34c5e58..b2c65c5 100644
--- a/config/data_module/dataset_types.yaml
+++ b/config/data_module/dataset_types.yaml
@@ -1,3 +1,22 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
dataset_types: # both pretraining datasets and unimodal finetuning datasets
tueg:
channels: [FP1-F7, F7-T3, T3-T5, T5-O1, FP2-F8, F8-T4, T4-T6, T6-O2, T3-C3, C3-CZ, CZ-C4, C4-T4, FP1-F3, F3-C3, C3-P3, P3-O1, FP2-F4, F4-C4, C4-P4, P4-O2, A1-T3, T4-A2]
@@ -30,4 +49,4 @@ dataset_types: # both pretraining datasets and unimodal finetuning datasets
cpsc2018:
channels: [I, II, III, AVR, AVL, AVF, V1, V2, V3, V4, V5, V6]
location_fn: "ecg"
- sensor_type: 0
\ No newline at end of file
+ sensor_type: 0
diff --git a/config/data_module/finetune_data_module_multimodal_PanLUNA.yaml b/config/data_module/finetune_data_module_multimodal_PanLUNA.yaml
index 1f52129..489e705 100644
--- a/config/data_module/finetune_data_module_multimodal_PanLUNA.yaml
+++ b/config/data_module/finetune_data_module_multimodal_PanLUNA.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
# This is example for SEED-VII dataset
# To adjust for other multimodal datasets replace channel_groups field with appropriate modalities and their channels
# Adjust slicing with channel start and end indexes
@@ -106,4 +106,4 @@ data_module:
test:
_target_: 'datasets.finetuning_multimodal_datasets_PanLUNA.FinetuningMultimodal_Dataset'
hdf5_file: "#CHANGEME"
- <<: *seed_channels
\ No newline at end of file
+ <<: *seed_channels
diff --git a/config/data_module/finetune_data_module_unimodal_PanLUNA.yaml b/config/data_module/finetune_data_module_unimodal_PanLUNA.yaml
index 39ee7ea..d5b5a12 100644
--- a/config/data_module/finetune_data_module_unimodal_PanLUNA.yaml
+++ b/config/data_module/finetune_data_module_unimodal_PanLUNA.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
defaults:
- dataset_types
diff --git a/config/data_module/multiloader_data_module_PanLUNA.yaml b/config/data_module/multiloader_data_module_PanLUNA.yaml
index 2595770..141e382 100644
--- a/config/data_module/multiloader_data_module_PanLUNA.yaml
+++ b/config/data_module/multiloader_data_module_PanLUNA.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
defaults:
- dataset_types
@@ -717,4 +717,4 @@ data_module:
num_channels: 29
channels: ${dataset_types.siena.channels}
location_fn: ${dataset_types.siena.location_fn}
- sensor_type: ${dataset_types.siena.sensor_type}
\ No newline at end of file
+ sensor_type: ${dataset_types.siena.sensor_type}
diff --git a/config/experiment/FEMBA_finetune.yaml b/config/experiment/FEMBA_finetune.yaml
index 7d6de75..1f55c96 100644
--- a/config/experiment/FEMBA_finetune.yaml
+++ b/config/experiment/FEMBA_finetune.yaml
@@ -34,6 +34,7 @@ layerwise_lr_decay: 0.75
scheduler_type: cosine
pretrained_checkpoint_path: null
+pretrained_safetensors_path: null
input_normalization:
normalize: True
@@ -82,4 +83,3 @@ scheduler:
min_lr: 2.5e-6 # minimum LR for the cosine scheduler
warmup_lr_init: 2.5e-7 # initial LR for the warmup phase
warmup_epochs: 10
-
diff --git a/config/experiment/LuMamba_finetune.yaml b/config/experiment/LuMamba_finetune.yaml
index 008e4f4..6be8cf8 100644
--- a/config/experiment/LuMamba_finetune.yaml
+++ b/config/experiment/LuMamba_finetune.yaml
@@ -67,7 +67,7 @@ defaults:
- override /data_module: finetune_data_module
- override /model: LuMamba_tiny
- override /scheduler: cosine
- - override /task: "#CHANGEME" # override with "finetune_task_LUNA" for classification and "finetune_regression_task_LuMamba" for regression tasks
+ - override /task: finetune_task_LUNA # Use finetune_regression_task_LuMamba for regression tasks
- override /criterion: finetune_criterion
# Change _target_ to corresponding dataset wrapper for finetuning dataset
diff --git a/config/experiment/PanLUNA_finetune.yaml b/config/experiment/PanLUNA_finetune.yaml
index 1a54431..b8b89b2 100644
--- a/config/experiment/PanLUNA_finetune.yaml
+++ b/config/experiment/PanLUNA_finetune.yaml
@@ -1,5 +1,24 @@
-
# @package _global_
+#*----------------------------------------------------------------------------*
+#* 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: Marija Zelic *
+#* Author: Thorir Mar Ingolfsson *
+#*----------------------------------------------------------------------------*
+
tag: PanLUNA_finetune
model_size: tiny
gpus: 4
@@ -117,5 +136,3 @@ scheduler:
warmup_lr_init: 2.5e-7 # initial LR for the warmup phase
warmup_epochs: 0
-
-
diff --git a/config/experiment/PanLUNA_pretrain.yaml b/config/experiment/PanLUNA_pretrain.yaml
index 2042d20..586ead6 100644
--- a/config/experiment/PanLUNA_pretrain.yaml
+++ b/config/experiment/PanLUNA_pretrain.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
tag: PanLUNA_pretrain
model_size: tiny
gpus: 4
@@ -94,4 +94,4 @@ optimizer:
optim: 'AdamW'
lr: 1.25e-4
betas: [0.9, 0.98]
- weight_decay: 0.05
\ No newline at end of file
+ weight_decay: 0.05
diff --git a/config/model/PanLUNA_finetune.yaml b/config/model/PanLUNA_finetune.yaml
index bd3ea35..cbcfbe6 100644
--- a/config/model/PanLUNA_finetune.yaml
+++ b/config/model/PanLUNA_finetune.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
model:
_target_: models.PanLUNA.PanLUNA
patch_size: 32
@@ -27,4 +27,4 @@ model:
num_queries: 4
mlp_ratio: 4
drop_path: 0.1
- num_classes: 2
\ No newline at end of file
+ num_classes: 2
diff --git a/config/model/PanLUNA_pretrain.yaml b/config/model/PanLUNA_pretrain.yaml
index 94d7311..ac6f080 100644
--- a/config/model/PanLUNA_pretrain.yaml
+++ b/config/model/PanLUNA_pretrain.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,7 +18,6 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
model:
_target_: models.PanLUNA.PanLUNA
patch_size: 32
@@ -27,4 +27,4 @@ model:
num_queries: 4
mlp_ratio: 4
drop_path: 0.0
- num_classes: 0
\ No newline at end of file
+ num_classes: 0
diff --git a/config/task/finetune_task_PanLUNA.yaml b/config/task/finetune_task_PanLUNA.yaml
index 5869f2e..2a7667e 100644
--- a/config/task/finetune_task_PanLUNA.yaml
+++ b/config/task/finetune_task_PanLUNA.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,6 +18,5 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
task:
- _target_: 'tasks.finetune_task_PanLUNA.FinetuneTask'
\ No newline at end of file
+ _target_: 'tasks.finetune_task_PanLUNA.FinetuneTask'
diff --git a/config/task/pretrain_task_PanLUNA.yaml b/config/task/pretrain_task_PanLUNA.yaml
index cf36213..16a1cc9 100644
--- a/config/task/pretrain_task_PanLUNA.yaml
+++ b/config/task/pretrain_task_PanLUNA.yaml
@@ -1,3 +1,4 @@
+# @package _global_
#*----------------------------------------------------------------------------*
#* Copyright (C) 2026 ETH Zurich, Switzerland *
#* SPDX-License-Identifier: Apache-2.0 *
@@ -17,6 +18,5 @@
#* Author: Marija Zelic *
#* Author: Thorir Mar Ingolfsson *
#*----------------------------------------------------------------------------*
-# @package _global_
task:
- _target_: 'tasks.pretrain_task_PanLUNA.MaskTask'
\ No newline at end of file
+ _target_: 'tasks.pretrain_task_PanLUNA.MaskTask'
diff --git a/data_module/multiloader_data_module.py b/data_module/multiloader_data_module.py
index 1ab8e8d..23fc2ed 100644
--- a/data_module/multiloader_data_module.py
+++ b/data_module/multiloader_data_module.py
@@ -1,3 +1,7 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2025-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 *
@@ -256,4 +260,4 @@ def get_balanced_subset(dataloader, total_samples, num_classes):
for c in collected:
print(f" - Class {c}: {len(collected[c])}")
- return balanced_subset
\ No newline at end of file
+ return balanced_subset
diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md
new file mode 100644
index 0000000..630f1ef
--- /dev/null
+++ b/docs/CITATIONS.md
@@ -0,0 +1,88 @@
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE at the repository root for details.
+
+# Citations
+
+Please cite the paper corresponding to the model used in your work. Published proceedings are preferred over the arXiv preprints when available.
+
+| Model | Publication status |
+| --- | --- |
+| FEMBA | Published at IEEE EMBC 2025 ([DOI](https://doi.org/10.1109/EMBC58623.2025.11252697)). |
+| LUNA | Published at NeurIPS 2025 ([proceedings](https://proceedings.neurips.cc/paper_files/paper/2025/hash/66969a9e6bd7a26dfeccea7227178ca7-Abstract-Conference.html)). |
+| TinyMyo | [arXiv preprint](https://arxiv.org/abs/2512.15729). |
+| LuMamba | Accepted at EUSIPCO 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2603.19100)). |
+| PanLUNA | Accepted at IEEE AICAS 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2604.04297)). |
+
+## FEMBA
+
+```bibtex
+@inproceedings{tegon2025femba,
+ title={{FEMBA}: Efficient and Scalable {EEG} Analysis with a Bidirectional Mamba Foundation Model},
+ author={Tegon, Anna and Ingolfsson, Thorir Mar and Wang, Xiaying and Benini, Luca and Li, Yawei},
+ booktitle={2025 47th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC)},
+ year={2025},
+ pages={1--7},
+ doi={10.1109/EMBC58623.2025.11252697}
+}
+```
+
+## LUNA
+
+```bibtex
+@inproceedings{doner2025luna,
+ title={{LUNA}: Efficient and Topology-Agnostic Foundation Model for {EEG} Signal Analysis},
+ author={D{\"o}ner, Berkay and Ingolfsson, Thorir Mar and Benini, Luca and Li, Yawei},
+ booktitle={Advances in Neural Information Processing Systems},
+ volume={38},
+ year={2025},
+ url={https://proceedings.neurips.cc/paper_files/paper/2025/hash/66969a9e6bd7a26dfeccea7227178ca7-Abstract-Conference.html}
+}
+```
+
+## TinyMyo
+
+```bibtex
+@misc{fasulo2025tinymyo,
+ title={{TinyMyo}: a Tiny Foundation Model for Flexible {EMG} Signal Processing at the Edge},
+ author={Fasulo, Matteo and Spacone, Giusy and Ingolfsson, Thorir Mar and Li, Yawei and Benini, Luca and Cossettini, Andrea},
+ year={2025},
+ eprint={2512.15729},
+ archivePrefix={arXiv},
+ primaryClass={eess.SP},
+ doi={10.48550/arXiv.2512.15729},
+ url={https://arxiv.org/abs/2512.15729}
+}
+```
+
+## LuMamba
+
+```bibtex
+@inproceedings{broustail2026lumamba,
+ title={{LuMamba}: Latent Unified Mamba for Electrode Topology-Invariant and Efficient {EEG} Modeling},
+ author={Broustail, Dana{\'e} and Tegon, Anna and Ingolfsson, Thorir Mar and Li, Yawei and Benini, Luca},
+ booktitle={2026 34th European Signal Processing Conference (EUSIPCO)},
+ year={2026},
+ note={Accepted; proceedings forthcoming},
+ eprint={2603.19100},
+ archivePrefix={arXiv},
+ primaryClass={cs.AI},
+ url={https://arxiv.org/abs/2603.19100}
+}
+```
+
+## PanLUNA
+
+```bibtex
+@inproceedings{zelic2026panluna,
+ title={{PanLUNA}: An Efficient and Robust Query-Unified Multimodal Model for Edge Biosignal Intelligence},
+ author={Zelic, Marija and Tegon, Anna and Li, Yawei and Ingolfsson, Thorir Mar and Benini, Luca},
+ booktitle={2026 8th IEEE International Conference on Artificial Intelligence Circuits and Systems (AICAS)},
+ year={2026},
+ note={Accepted; proceedings forthcoming},
+ eprint={2604.04297},
+ archivePrefix={arXiv},
+ primaryClass={cs.AI},
+ url={https://arxiv.org/abs/2604.04297}
+}
+```
+
+The LuMamba and PanLUNA entries should be updated with their final DOI and page numbers once the EUSIPCO and AICAS proceedings are published.
diff --git a/docs/README.md b/docs/README.md
index 50bb92b..e2f88ea 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,10 +1,10 @@
-Copyright (C) 2025 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE file at the root of the repository for details.
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE file at the root of the repository for details.
# Documentation
-Welcome to the documentation for **BioFoundation**. This directory provides in-depth information about the core components of the repository, including the architecture of our models, the datasets we use, and the training tasks we implement.
+Welcome to the BioFoundation documentation. This directory covers the model families, biosignal datasets, and PyTorch Lightning training tasks used by the repository.
-Good documentation is essential for reproducible and collaborative research. We have aimed to provide clear and comprehensive explanations to help you understand, use, and extend our work.
+Start with the root [`README.md`](../README.md) to choose a model and launch an experiment. Use the pages here when adapting an architecture, dataset, or task.
---
@@ -12,34 +12,35 @@ Good documentation is essential for reproducible and collaborative research. We
This documentation is organized into the following sections:
-### 1. **[Datasets](./datasets/)**
+### 1. [Models](./model/)
-This section provides detailed descriptions of the various EEG datasets used for pre-training and fine-tuning in this project. You will find information on:
-- The **TUH EEG Corpus (TUEG)**, used for our large-scale self-supervised pre-training.
-- The **TUH Abnormal EEG (TUAB)** dataset for normal vs. abnormal classification.
-- The **TUH Artifact (TUAR)** dataset for artifact detection tasks.
-- The **TUH Slowing (TUSL)** dataset for identifying slowing events in EEG signals.
+Architecture and usage notes are available for every published model family:
-Each document covers the dataset's purpose, size, and specific characteristics.
+- [FEMBA](./model/FEMBA.md): bidirectional Mamba for EEG.
+- [LUNA](./model/LUNA.md): query-unified, topology-agnostic EEG.
+- [TinyMyo](./model/TinyMyo.md): compact foundation model for sEMG.
+- [LuMamba](./model/LuMamba.md): query-unified Mamba for EEG.
+- [PanLUNA](./model/PanLUNA.md): sensor-aware modeling across EEG, ECG, and PPG.
-### 2. **[Model](./model/)**
+Each model also has a pretrained release linked from the root model zoo. The canonical machine-readable index is [`biofoundation/model_registry.py`](../biofoundation/model_registry.py).
-This section contains a detailed breakdown of our model architectures, **FEMBA** and **LUNA**. For **FEMBA** The documentation covers:
-- The **architecture overview**, including the tokenizer, encoder, decoder, and classifier heads.
-- The **bidirectional Mamba block** that forms the core of our encoder.
-- Details on **model variants** (e.g., FEMBA-tiny, FEMBA-base, FEMBA-large).
-- The self-supervised learning objective and classification protocols.
-For **LUNA**, the documentation includes:
-- The **architecture overview**, including the path-feature extraction, channel-unification module, patch-wise temporal encoder, decoder, and classifier heads.
-- The self-supervised learning objective and classification protocols.
-- Details on **model variants** (e.g., LUNA-base, LUNA-large, LUNA-huge).
+### 2. [Datasets](./datasets/)
+The dataset pages describe preparation and evaluation protocols for the TUH EEG corpora used in pre-training and downstream tasks, including TUEG, TUAB, TUAR, and TUSL. Additional preprocessing entry points for EEG, EMG, ECG, and PPG datasets live in [`make_datasets`](../make_datasets/).
-### 3. **[Tasks](./tasks/)**
+### 3. [Tasks](./tasks/)
-This section describes the PyTorch Lightning `LightningModule` implementations that define our training pipelines. It explains:
-- The **`PretrainTask`**, which handles the self-supervised masked signal reconstruction.
-- The **`FinetuneTask`**, which manages the supervised training for downstream classification tasks.
-- Details on the data flow, loss computation, metrics, and optimization strategies for each task.
+The task pages describe the Lightning modules for self-supervised pre-training, classification, and regression. Task steps normalize incoming data through [`SignalBatch`](../biofoundation/core/batch.py), which provides a common input contract while preserving model-specific channel and sensor metadata.
-We encourage you to explore these documents to gain a deeper understanding of the components that make up the BioFoundation framework.
+### 4. [Training Guide](./TRAINING.md)
+
+The training guide covers environment variables, Hydra experiment selection, the shared batch contract, distributed training, checkpoints, and fast validation.
+
+### 5. Project References
+
+- [`CONTRIBUTING.md`](../CONTRIBUTING.md) defines extension and pull request expectations.
+- [`CITATIONS.md`](./CITATIONS.md) contains BibTeX for all five model families.
+- [`config/README.md`](../config/README.md) explains Hydra composition and overrides.
+- [`make_datasets/README.md`](../make_datasets/README.md) documents preprocessing and HDF5 conversion.
+
+ARES deployment documentation remains self-contained under [`ARES/docs`](../ARES/docs/) because it uses an independent environment and workflow.
diff --git a/docs/TRAINING.md b/docs/TRAINING.md
new file mode 100644
index 0000000..d3ccd73
--- /dev/null
+++ b/docs/TRAINING.md
@@ -0,0 +1,84 @@
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE at the repository root for details.
+
+# Training Guide
+
+BioFoundation uses Hydra for configuration and PyTorch Lightning for training. Run commands from the repository root.
+
+## Environment
+
+Set the prepared dataset root and experiment-output root before starting Hydra:
+
+```bash
+export DATA_PATH=/absolute/path/to/data
+export CHECKPOINT_DIR=/absolute/path/to/experiments
+```
+
+The CLI validates both variables before Hydra composes a run. Repository modules can still be imported without them.
+
+## Choosing an Experiment
+
+Each published model has matching pre-training and fine-tuning files under `config/experiment/`:
+
+```bash
+python -u run_train.py +experiment=FEMBA_pretrain
+python -u run_train.py +experiment=LUNA_finetune \
+ pretrained_safetensors_path=/absolute/path/to/LUNA_base.safetensors
+```
+
+Hydra applies command-line overrides after experiment, module, and default configs. See [`../config/README.md`](../config/README.md) for the composition model. Search the selected experiment and data module for values that must be supplied locally:
+
+```bash
+rg -n '#CHANGEME' config/experiment config/data_module
+```
+
+## Batch Contract
+
+Tasks normalize batches with `biofoundation.core.batch.as_signal_batch`. The canonical mapping supports:
+
+| Field | Purpose |
+| --- | --- |
+| `input` | Signal tensor required by every task. |
+| `label` | Supervised classification or regression target. |
+| `channel_locations` | Required by topology-aware model families. |
+| `channel_names` | Optional channel identity metadata. |
+| `sensor_type` | Required by PanLUNA for EEG, ECG, and PPG identity. |
+| `metadata` | Optional dataset-specific metadata. |
+
+The model registry records the required metadata for each family.
+
+## Distributed Training
+
+Configure Lightning's trainer in the selected experiment:
+
+```yaml
+trainer:
+ accelerator: gpu
+ num_nodes: ${num_nodes}
+ devices: -1
+ strategy: ddp
+```
+
+Set `find_unused_parameters` in the experiment when a fine-tuning strategy leaves branches of the model inactive.
+
+## Checkpoints
+
+Every model page under [`model/`](model/) links to its PulpBio Hugging Face repository and shows the expected checkpoint override. `CHECKPOINT_DIR` remains the output root; pass input weights through `pretrained_safetensors_path` or `pretrained_checkpoint_path`.
+
+## Memory and Runtime
+
+- Reduce `batch_size`, segment duration, or model size first.
+- Use mixed precision through Lightning where the target hardware supports it.
+- Consider activation checkpointing or sharded training for large models only after establishing a correct single-device run.
+- ARES deployment has a separate environment and workflow under [`../ARES`](../ARES/).
+
+## Fast Validation
+
+The fast checks validate repository contracts without model weights or datasets. Install `hydra-core` to include full experiment composition:
+
+```bash
+python -m unittest discover -s tests -p 'test_*.py' -v
+python -m compileall -q biofoundation run_train.py models tasks datasets data_module \
+ criterion schedulers util make_datasets tests
+```
+
+Run representative CPU or GPU batches for any change to numerical model, loss, or data behavior.
diff --git a/docs/model/FEMBA.md b/docs/model/FEMBA.md
index a21e2ad..7883326 100644
--- a/docs/model/FEMBA.md
+++ b/docs/model/FEMBA.md
@@ -119,3 +119,23 @@ FEMBA supports multiple downstream EEG classification paradigms, particularly fo
- FEMBA-Base: 0.731 AUROC
---
+
+### Pretrained Weights
+
+The [PulpBio/FEMBA Hugging Face repository](https://huggingface.co/PulpBio/FEMBA) provides task checkpoints for TUAB, TUAR, and TUSL. The checkpoint weights are licensed under CC BY-ND 4.0.
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+ repo_id="PulpBio/FEMBA",
+ local_dir="checkpoints/FEMBA",
+)
+```
+
+Run fine-tuning from the repository root:
+
+```bash
+python -u run_train.py +experiment=FEMBA_finetune \
+ pretrained_safetensors_path=/absolute/path/to/checkpoints/FEMBA/TUAR/base.safetensors
+```
diff --git a/docs/model/LUNA.md b/docs/model/LUNA.md
index d9d4698..9917389 100644
--- a/docs/model/LUNA.md
+++ b/docs/model/LUNA.md
@@ -117,4 +117,28 @@ LUNA is available in three sizes, scaled by increasing the depth of the temporal
- LUNA-Huge: 0.802 AUROC (State-of-the-art).
**SEED-V (Emotion Recognition)**
-- LUNA-Large: 39.18% balanced accuracy.
\ No newline at end of file
+- LUNA-Large: 39.18% balanced accuracy.
+
+---
+
+### Pretrained Weights
+
+The [PulpBio/LUNA Hugging Face repository](https://huggingface.co/PulpBio/LUNA) provides Base, Large, and Huge checkpoints matching the model configs in `config/model/`.
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+ repo_id="PulpBio/LUNA",
+ local_dir="checkpoints/LUNA",
+)
+```
+
+Select the matching model size and checkpoint:
+
+```bash
+python -u run_train.py +experiment=LUNA_finetune /model=LUNA_base \
+ pretrained_safetensors_path=/absolute/path/to/checkpoints/LUNA/Base/LUNA_base.safetensors
+```
+
+For TUH datasets, use `finetune_data_module`; for subject-independent datasets such as SEED-V, use `subject_independent_data_module`. Update the experiment's `classification_type`, `model.num_classes`, and `#CHANGEME` paths for the downstream task.
diff --git a/docs/model/LuMamba.md b/docs/model/LuMamba.md
index 44f79dd..8935ec0 100644
--- a/docs/model/LuMamba.md
+++ b/docs/model/LuMamba.md
@@ -149,3 +149,25 @@ Larger model sizes can be attained by increasing the number of bi-Mamba blocks `
- LuMamba-Tiny (reconstruction-only): 0.11 $R^2$, 0.38 Pearson's correlation.
---
+
+### Pretrained Weights
+
+The [PulpBio/LuMamba Hugging Face repository](https://huggingface.co/PulpBio/LuMamba) provides LeJEPA-only, reconstruction-only, and mixed pre-training variants.
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+ repo_id="PulpBio/LuMamba",
+ local_dir="checkpoints/LuMamba",
+)
+```
+
+Run fine-tuning from the repository root:
+
+```bash
+python -u run_train.py +experiment=LuMamba_finetune \
+ pretrained_safetensors_path=/absolute/path/to/checkpoints/LuMamba/LuMamba.safetensors
+```
+
+Use the preprocessing scripts in `make_datasets/` for downstream datasets, then update `config/experiment/LuMamba_finetune.yaml` for the dataset paths, task type, class count, and trainer settings.
diff --git a/docs/model/PanLUNA.md b/docs/model/PanLUNA.md
index 3bc6d4b..9b9e7a2 100644
--- a/docs/model/PanLUNA.md
+++ b/docs/model/PanLUNA.md
@@ -176,3 +176,27 @@ State-of-the-art on PTB-XL Super and CSN. QAT INT8 recovers ≥96% of FP32 AUROC
- In the 4-class setting, PanLUNA achieves the strongest results among PPG-based foundation models under 10M parameters, with especially large gains over Light Pulse-PPG in AUC-PR.
- The PPG+ECG setting does not consistently improve over PPG-only in the 4-class case, suggesting dataset- and modality-specific effects.
+
+---
+
+### Pretrained Weights
+
+The [PulpBio/PanLUNA Hugging Face repository](https://huggingface.co/PulpBio/PanLUNA) provides the pretrained pan-modal checkpoint and fine-tuning guidance.
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+ repo_id="PulpBio/PanLUNA",
+ local_dir="checkpoints/PanLUNA",
+)
+```
+
+Run fine-tuning from the repository root:
+
+```bash
+python -u run_train.py +experiment=PanLUNA_finetune \
+ pretrained_safetensors_path=/absolute/path/to/checkpoints/PanLUNA/PanLUNA.safetensors
+```
+
+Choose `finetune_data_module_unimodal_PanLUNA` or `finetune_data_module_multimodal_PanLUNA` in the experiment config. Set `classification_type`, `model.num_classes`, channel groups, and `finetuning.mode` (`full`, `freeze_encoder`, or `lora`) for the target task.
diff --git a/docs/model/TinyMyo.md b/docs/model/TinyMyo.md
index 35a332d..8212ecb 100644
--- a/docs/model/TinyMyo.md
+++ b/docs/model/TinyMyo.md
@@ -247,3 +247,30 @@ This is the **first demonstration of an EMG FM on a microcontroller**.
* **Speech Recognition:** WER 33.95%
Overall TinyMyo matches or exceeds state-of-the-art while being on par with or smaller than prior EMG foundation models.
+
+---
+
+## Pretrained Weights
+
+The [PulpBio/TinyMyo Hugging Face repository](https://huggingface.co/PulpBio/TinyMyo) provides task checkpoints for DB5, UCI EMG, and EPN612, along with dataset download and preprocessing scripts.
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+ repo_id="PulpBio/TinyMyo",
+ local_dir="checkpoints/TinyMyo",
+)
+```
+
+Run fine-tuning from the repository root:
+
+```bash
+python -u run_train.py +experiment=TinyMyo_finetune \
+ pretrained_safetensors_path=/absolute/path/to/checkpoints/TinyMyo/UCI_EMG/base.safetensors
+```
+
+Related experiments remain in dedicated repositories:
+
+- [Silent Speech](https://github.com/MatteoFasulo/silent_speech)
+- [Generic Neuromotor Interface](https://github.com/MatteoFasulo/generic-neuromotor-interface)
diff --git a/models/README.md b/models/README.md
index 300cafe..81c5673 100644
--- a/models/README.md
+++ b/models/README.md
@@ -1,12 +1,17 @@
-Copyright (C) 2025 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE file at the root of the repository for details.
+Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE file at the root of the repository for details.
# Models
-This directory contains the implementations of the deep learning models used in the **BioFoundation** project. Each model is defined as a PyTorch `nn.Module` and is designed to be configurable and extensible for various research tasks.
+This directory contains the PyTorch `nn.Module` implementations for the BioFoundation model families. Hydra model settings live in [`../config/model`](../config/model/), while the canonical family metadata and batch requirements live in [`../biofoundation/model_registry.py`](../biofoundation/model_registry.py).
## Available Models
-- **FEMBA**: A lightweight EEG model designed for both pretraining and fine-tuning tasks. For a more detailed description of the model check the [documentation](../docs/model/FEMBA.md).
-- **LUNA**: An efficient EEG model specifically designed for handling different types of electrode configurations. For a more detailed description of the model check the [documentation](../docs/model/LUNA.md).
-- **TinyMyo**: A 3.6M-parameter Transformer-based foundation model for surface EMG (sEMG). It is pretrained on >480 GB of EMG data and optimized for ultra-low-power, real-time deployment, including microcontrollers (GAP9) where it achieves an inference time of 0.785 s, energy of 44.91 mJ and power envelope of 57.18 mW. For a more detailed description of the model check the [documentation](../docs/model/TinyMyo.md).
-- **PanLUNA**: A 5.4M-parameter pan-modal biosignal foundation model that extends LUNA’s channel-unification mechanism from topology invariance to cross-modal fusion, jointly processing EEG, ECG, and PPG within a single shared encoder via sensor-type embeddings, with no modality-specific backbones or paired multimodal data required during pretraining; it is pretrained on ~40,000 hours of heterogeneous biosignal data using masked signal reconstruction and achieves strong performance across unimodal and multimodal evaluation settings. For more detailed description of the model check [documentation]../docs/model/PanLUNA.md).
+| Model | Signals | Summary | Resources |
+| --- | --- | --- | --- |
+| FEMBA | EEG | Efficient bidirectional Mamba model for long EEG sequences. | [Documentation](../docs/model/FEMBA.md) / [Hugging Face](https://huggingface.co/PulpBio/FEMBA) |
+| LUNA | EEG | Query-based channel unification for topology-agnostic EEG modeling. | [Documentation](../docs/model/LUNA.md) / [Hugging Face](https://huggingface.co/PulpBio/LUNA) |
+| TinyMyo | sEMG | Compact rotary Transformer designed for flexible EMG processing and edge deployment. | [Documentation](../docs/model/TinyMyo.md) / [Hugging Face](https://huggingface.co/PulpBio/TinyMyo) |
+| LuMamba | EEG | LUNA-style channel unification with efficient Mamba temporal modeling. | [Documentation](../docs/model/LuMamba.md) / [Hugging Face](https://huggingface.co/PulpBio/LuMamba) |
+| PanLUNA | EEG, ECG, PPG | Sensor-aware query unification for unimodal and multimodal biosignal learning. | [Documentation](../docs/model/PanLUNA.md) / [Hugging Face](https://huggingface.co/PulpBio/PanLUNA) |
+
+Use the matching pre-training or fine-tuning experiment in [`../config/experiment`](../config/experiment/) rather than instantiating a model in isolation when starting a reproducible run.
diff --git a/models/TinyMyo.py b/models/TinyMyo.py
index 5e15687..0a417b1 100644
--- a/models/TinyMyo.py
+++ b/models/TinyMyo.py
@@ -1,3 +1,22 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2025-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 *
+#*----------------------------------------------------------------------------*
+
import math
from dataclasses import dataclass
from typing import Literal, Optional, Tuple
diff --git a/models/modules/__init__.py b/models/modules/__init__.py
index e69de29..338ba45 100644
--- a/models/modules/__init__.py
+++ b/models/modules/__init__.py
@@ -0,0 +1,18 @@
+#*----------------------------------------------------------------------------*
+#* Copyright (C) 2025-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 *
+#*----------------------------------------------------------------------------*
diff --git a/pyproject.toml b/pyproject.toml
index 65be4eb..ccf1d2d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,6 +10,7 @@ dependencies = [
"fvcore",
"omegaconf",
"hydra-core",
+ "peft",
"scipy",
"h5py",
"python-dateutil",
@@ -23,7 +24,6 @@ dependencies = [
"torch>=2.4.0, <2.8.0",
"scikit-learn",
"tensorboard",
- "tensorrt",
"torch_optimizer",
"tables",
"timm",
@@ -33,6 +33,7 @@ dependencies = [
"warmup-scheduler",
"safetensors",
"rotary-embedding-torch",
+ "wandb",
]
[tool.uv.sources]
diff --git a/requirements.txt b/requirements.txt
index ab79d32..de7de00 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,7 +17,6 @@ tensorboardX
git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git
scikit-learn
tensorboard
-tensorrt
torch_optimizer
tables
timm
@@ -27,4 +26,4 @@ torchmetrics
warmup-scheduler
safetensors
rotary-embedding-torch
-wandb
\ No newline at end of file
+wandb
diff --git a/run_train.py b/run_train.py
index 0431cfc..6fbe6d6 100644
--- a/run_train.py
+++ b/run_train.py
@@ -35,13 +35,9 @@
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.strategies import DDPStrategy
+from biofoundation.core.environment import require_environment
from util.train_utils import find_last_checkpoint_path
-for env_var in ["DATA_PATH", "CHECKPOINT_DIR"]:
- env_var_value = os.environ.get(env_var)
- if env_var_value is None or env_var_value == "#CHANGEME":
- raise RuntimeError(f"Environment variable {env_var} is not set. Please set it before running the script.")
-
OmegaConf.register_new_resolver("env", lambda key: os.getenv(key))
OmegaConf.register_new_resolver("get_method", hydra.utils.get_method)
@@ -191,6 +187,6 @@ def run(cfg: DictConfig):
if __name__ == "__main__":
- # Ensure environment variables are set before Hydra processes the config
+ require_environment(("DATA_PATH", "CHECKPOINT_DIR"))
os.environ["HYDRA_FULL_ERROR"] = "1"
run()
diff --git a/tasks/finetune_regression_task_LuMamba.py b/tasks/finetune_regression_task_LuMamba.py
index d58855c..e7c3bc5 100644
--- a/tasks/finetune_regression_task_LuMamba.py
+++ b/tasks/finetune_regression_task_LuMamba.py
@@ -33,6 +33,8 @@
from safetensors.torch import load_file
from collections import OrderedDict
+from biofoundation.core.batch import as_signal_batch
+
class FinetuneRegressionTask(pl.LightningModule):
"""
@@ -147,6 +149,7 @@ def _step(self, X, mask, channel_locations):
return y_pred
def training_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
@@ -170,6 +173,7 @@ def training_step(self, batch, batch_idx):
return loss
def validation_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
@@ -194,6 +198,7 @@ def validation_step(self, batch, batch_idx):
return loss
def test_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
@@ -268,4 +273,4 @@ def configure_optimizers(self):
"frequency": 1
}
- return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
\ No newline at end of file
+ return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
diff --git a/tasks/finetune_task.py b/tasks/finetune_task.py
index 903b86b..3196221 100644
--- a/tasks/finetune_task.py
+++ b/tasks/finetune_task.py
@@ -30,6 +30,7 @@
Accuracy, Precision, Recall, AUROC,
AveragePrecision, CohenKappa, F1Score
)
+from biofoundation.core.batch import as_signal_batch
from util.train_utils import RobustQuartileNormalize
class FinetuneTask(pl.LightningModule):
@@ -184,7 +185,8 @@ def _step(self, X, mask):
}
def training_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
@@ -208,7 +210,8 @@ def training_step(self, batch, batch_idx):
return loss
def validation_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
@@ -232,7 +235,8 @@ def validation_step(self, batch, batch_idx):
return loss
def test_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
diff --git a/tasks/finetune_task_EMG.py b/tasks/finetune_task_EMG.py
index 8599443..9b0f0c0 100644
--- a/tasks/finetune_task_EMG.py
+++ b/tasks/finetune_task_EMG.py
@@ -36,6 +36,7 @@
Recall,
)
+from biofoundation.core.batch import as_signal_batch
from util.train_utils import MinMaxNormalization
@@ -208,7 +209,8 @@ def _step(self, X: torch.Tensor, mask: Optional[torch.Tensor] = None) -> dict:
}
def training_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
@@ -231,7 +233,8 @@ def training_step(self, batch, batch_idx):
return loss
def validation_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
@@ -246,7 +249,8 @@ def validation_step(self, batch, batch_idx):
return loss
def test_step(self, batch, batch_idx):
- X, y = batch
+ batch = as_signal_batch(batch)
+ X, y = batch["input"], batch["label"]
if self.normalize:
X = self.normalize_fct(X)
mask = self.generate_fake_mask(X.shape[0], X.shape[1], X.shape[2])
diff --git a/tasks/finetune_task_LUNA.py b/tasks/finetune_task_LUNA.py
index 5190f78..5b9069a 100644
--- a/tasks/finetune_task_LUNA.py
+++ b/tasks/finetune_task_LUNA.py
@@ -29,6 +29,7 @@
Accuracy, Precision, Recall, AUROC,
AveragePrecision, CohenKappa, F1Score
)
+from biofoundation.core.batch import as_signal_batch
from safetensors.torch import load_file
from collections import OrderedDict
@@ -219,6 +220,7 @@ def _step(self, X, mask, channel_locations):
}
def training_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
if self.normalize:
@@ -240,6 +242,7 @@ def training_step(self, batch, batch_idx):
return loss
def validation_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
if self.normalize:
@@ -262,6 +265,7 @@ def validation_step(self, batch, batch_idx):
return loss
def test_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
if self.normalize:
diff --git a/tasks/finetune_task_PanLUNA.py b/tasks/finetune_task_PanLUNA.py
index 4486ddf..b820546 100644
--- a/tasks/finetune_task_PanLUNA.py
+++ b/tasks/finetune_task_PanLUNA.py
@@ -32,6 +32,7 @@
from collections import OrderedDict
from safetensors.torch import load_file
from peft import LoraConfig, get_peft_model, TaskType
+from biofoundation.core.batch import as_signal_batch
class ChannelWiseNormalize:
def __init__(self, eps=1e-8):
@@ -293,6 +294,7 @@ def _step(self, X, mask, channel_locations, channel_names, sensor_type):
}
def training_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
@@ -322,6 +324,7 @@ def training_step(self, batch, batch_idx):
return loss
def validation_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
@@ -349,6 +352,7 @@ def validation_step(self, batch, batch_idx):
return loss
def test_step(self, batch, batch_idx):
+ batch = as_signal_batch(batch)
X, y = batch["input"], batch["label"]
channel_locations = batch["channel_locations"]
diff --git a/tasks/pretrain_task.py b/tasks/pretrain_task.py
index 103f3a7..134f5b9 100644
--- a/tasks/pretrain_task.py
+++ b/tasks/pretrain_task.py
@@ -24,6 +24,7 @@
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
+from biofoundation.core.batch import as_signal_batch
from util.train_utils import RobustQuartileNormalize
@@ -97,7 +98,8 @@ def training_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
- X = batch
+ batch = as_signal_batch(batch)
+ X = batch["input"]
mask = self.generate_mask(X.shape[0], X.shape[1], X.shape[2])
if self.normalize:
@@ -124,7 +126,8 @@ def validation_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
- X = batch
+ batch = as_signal_batch(batch)
+ X = batch["input"]
mask = self.generate_mask(X.shape[0], X.shape[1], X.shape[2])
if self.normalize:
diff --git a/tasks/pretrain_task_EMG.py b/tasks/pretrain_task_EMG.py
index c535075..de507c1 100644
--- a/tasks/pretrain_task_EMG.py
+++ b/tasks/pretrain_task_EMG.py
@@ -24,6 +24,7 @@
import torch_optimizer as torch_optim
from omegaconf import DictConfig
+from biofoundation.core.batch import as_signal_batch
from util.train_utils import MinMaxNormalization
@@ -109,7 +110,8 @@ def training_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
- X = batch
+ batch = as_signal_batch(batch)
+ X = batch["input"]
mask = self.generate_mask(X.shape[0], X.shape[1], X.shape[2])
if self.normalize:
@@ -145,7 +147,8 @@ def validation_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
- X = batch
+ batch = as_signal_batch(batch)
+ X = batch["input"]
mask = self.generate_mask(X.shape[0], X.shape[1], X.shape[2])
if self.normalize:
diff --git a/tasks/pretrain_task_LUNA.py b/tasks/pretrain_task_LUNA.py
index c9fb68f..383d10c 100644
--- a/tasks/pretrain_task_LUNA.py
+++ b/tasks/pretrain_task_LUNA.py
@@ -25,6 +25,7 @@
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
+from biofoundation.core.batch import as_signal_batch
from criterion.query_specialization_criterion import QuerySpecializationCriterion
class ChannelWiseNormalize:
@@ -107,6 +108,7 @@ def training_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
@@ -140,6 +142,7 @@ def validation_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
diff --git a/tasks/pretrain_task_LuMamba.py b/tasks/pretrain_task_LuMamba.py
index 2e60073..496a70f 100644
--- a/tasks/pretrain_task_LuMamba.py
+++ b/tasks/pretrain_task_LuMamba.py
@@ -33,6 +33,8 @@
import psutil, os
import gc
+from biofoundation.core.batch import as_signal_batch
+
def is_dist_initialized():
return dist.is_available() and dist.is_initialized()
@@ -406,6 +408,7 @@ def training_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
@@ -473,6 +476,7 @@ def validation_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
@@ -842,4 +846,4 @@ def log_tSNE_embeddings(self, dataset_name="TUAB", mode="flatten"):
del tsne
gc.collect()
- torch.cuda.empty_cache()
\ No newline at end of file
+ torch.cuda.empty_cache()
diff --git a/tasks/pretrain_task_PanLUNA.py b/tasks/pretrain_task_PanLUNA.py
index cad53c5..68e5cc7 100644
--- a/tasks/pretrain_task_PanLUNA.py
+++ b/tasks/pretrain_task_PanLUNA.py
@@ -26,6 +26,7 @@
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
+from biofoundation.core.batch import as_signal_batch
from criterion.query_specialization_criterion import QuerySpecializationCriterion
class ChannelWiseNormalize:
@@ -111,6 +112,7 @@ def training_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
@@ -151,6 +153,7 @@ def validation_step(self, batch, batch_idx):
Returns:
torch.Tensor: Loss value.
"""
+ batch = as_signal_batch(batch)
X = batch["input"]
channel_locations = batch["channel_locations"]
channel_names = batch.get("channel_names", None)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..21ae2a7
--- /dev/null
+++ b/tests/__init__.py
@@ -0,0 +1,19 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
diff --git a/tests/test_batch.py b/tests/test_batch.py
new file mode 100644
index 0000000..ba995c8
--- /dev/null
+++ b/tests/test_batch.py
@@ -0,0 +1,63 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+import unittest
+
+from biofoundation.core.batch import BatchRequirements, as_signal_batch, require_batch_fields
+
+
+class SignalBatchTest(unittest.TestCase):
+ def test_normalizes_raw_input(self):
+ signal = object()
+ self.assertIs(as_signal_batch(signal)["input"], signal)
+
+ def test_normalizes_input_label_tuple(self):
+ signal = object()
+ label = object()
+ batch = as_signal_batch((signal, label))
+ self.assertIs(batch["input"], signal)
+ self.assertIs(batch["label"], label)
+
+ def test_preserves_mapping_metadata(self):
+ source = {"input": object(), "channel_locations": object(), "sensor_type": object()}
+ batch = as_signal_batch(source)
+ self.assertEqual(batch, source)
+ self.assertIsNot(batch, source)
+
+ def test_rejects_ambiguous_sequences(self):
+ with self.assertRaisesRegex(ValueError, "received 3 values"):
+ as_signal_batch((1, 2, 3))
+
+ def test_rejects_mapping_without_input(self):
+ with self.assertRaisesRegex(ValueError, "'input' field"):
+ as_signal_batch({"label": 1})
+
+ def test_validates_model_specific_metadata(self):
+ requirements = BatchRequirements(channel_locations=True, sensor_type=True)
+ batch = as_signal_batch({"input": object(), "channel_locations": object()})
+ with self.assertRaisesRegex(ValueError, "sensor_type"):
+ require_batch_fields(batch, requirements)
+
+ batch["sensor_type"] = object()
+ self.assertIs(require_batch_fields(batch, requirements), batch)
+
+
+if __name__ == "__main__":
+ unittest.main()
+
diff --git a/tests/test_environment.py b/tests/test_environment.py
new file mode 100644
index 0000000..8896ac2
--- /dev/null
+++ b/tests/test_environment.py
@@ -0,0 +1,43 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+import unittest
+
+from biofoundation.core.environment import missing_environment_variables, require_environment
+
+
+class EnvironmentTest(unittest.TestCase):
+ def test_reports_unset_and_placeholder_values(self):
+ environ = {"DATA_PATH": "/data", "CHECKPOINT_DIR": "#CHANGEME"}
+ self.assertEqual(
+ missing_environment_variables(("DATA_PATH", "CHECKPOINT_DIR", "OUTPUT_PATH"), environ),
+ ("CHECKPOINT_DIR", "OUTPUT_PATH"),
+ )
+
+ def test_accepts_complete_environment(self):
+ require_environment(("DATA_PATH", "CHECKPOINT_DIR"), {"DATA_PATH": "/data", "CHECKPOINT_DIR": "/ckpt"})
+
+ def test_raises_one_actionable_error(self):
+ with self.assertRaisesRegex(RuntimeError, "DATA_PATH, CHECKPOINT_DIR"):
+ require_environment(("DATA_PATH", "CHECKPOINT_DIR"), {})
+
+
+if __name__ == "__main__":
+ unittest.main()
+
diff --git a/tests/test_hydra_composition.py b/tests/test_hydra_composition.py
new file mode 100644
index 0000000..1e754c1
--- /dev/null
+++ b/tests/test_hydra_composition.py
@@ -0,0 +1,64 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+import unittest
+from pathlib import Path
+
+from biofoundation.model_registry import MODEL_REGISTRY
+
+try:
+ from hydra import compose, initialize_config_dir
+ from omegaconf import OmegaConf
+except ImportError:
+ compose = None
+ initialize_config_dir = None
+ OmegaConf = None
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+@unittest.skipIf(compose is None, "hydra-core is not installed")
+class HydraCompositionTest(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ if not OmegaConf.has_resolver("env"):
+ OmegaConf.register_new_resolver("env", lambda _: "/tmp/biofoundation")
+
+ def test_every_registered_experiment_composes_and_resolves(self):
+ experiments = {
+ experiment
+ for spec in MODEL_REGISTRY.values()
+ for experiment in (spec.pretrain_experiment, spec.finetune_experiment)
+ }
+
+ for experiment in sorted(experiments):
+ with self.subTest(experiment=experiment):
+ with initialize_config_dir(version_base="1.1", config_dir=str(ROOT / "config")):
+ config = compose(
+ config_name="defaults",
+ overrides=[f"+experiment={experiment}"],
+ )
+ resolved = OmegaConf.to_container(config, resolve=True)
+ self.assertIsInstance(resolved, dict)
+ self.assertTrue(resolved["tag"])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py
new file mode 100644
index 0000000..fdb750a
--- /dev/null
+++ b/tests/test_model_registry.py
@@ -0,0 +1,100 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+import ast
+import unittest
+from pathlib import Path
+
+from biofoundation.core.batch import BatchRequirements
+from biofoundation.model_registry import MODEL_REGISTRY, get_model_spec
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class ModelRegistryTest(unittest.TestCase):
+ def test_registry_contains_every_published_model(self):
+ self.assertEqual(set(MODEL_REGISTRY), {"femba", "luna", "tinymyo", "lumamba", "panluna"})
+
+ def test_model_targets_and_experiments_exist(self):
+ for spec in MODEL_REGISTRY.values():
+ module_name, class_name = spec.model_target.rsplit(".", 1)
+ module_path = ROOT / f"{module_name.replace('.', '/')}.py"
+ self.assertTrue(module_path.is_file(), module_path)
+
+ tree = ast.parse(module_path.read_text(encoding="utf-8"))
+ classes = {node.name for node in tree.body if isinstance(node, ast.ClassDef)}
+ self.assertIn(class_name, classes, spec.model_target)
+
+ for experiment in (spec.pretrain_experiment, spec.finetune_experiment):
+ path = ROOT / "config" / "experiment" / f"{experiment}.yaml"
+ self.assertTrue(path.is_file(), path)
+
+ finetune_config = ROOT / "config" / "experiment" / f"{spec.finetune_experiment}.yaml"
+ self.assertIn("pretrained_safetensors_path", finetune_config.read_text(encoding="utf-8"))
+
+ def test_batch_requirements_match_model_inputs(self):
+ self.assertEqual(MODEL_REGISTRY["femba"].batch_requirements, BatchRequirements())
+ self.assertEqual(MODEL_REGISTRY["tinymyo"].batch_requirements, BatchRequirements())
+ self.assertEqual(
+ MODEL_REGISTRY["luna"].batch_requirements,
+ BatchRequirements(channel_locations=True),
+ )
+ self.assertEqual(
+ MODEL_REGISTRY["lumamba"].batch_requirements,
+ BatchRequirements(channel_locations=True),
+ )
+ self.assertEqual(
+ MODEL_REGISTRY["panluna"].batch_requirements,
+ BatchRequirements(channel_locations=True, sensor_type=True),
+ )
+
+ def test_huggingface_and_paper_links_are_unique_and_documented(self):
+ readme = (ROOT / "README.md").read_text(encoding="utf-8")
+ hubs = [spec.huggingface_url for spec in MODEL_REGISTRY.values()]
+ self.assertEqual(len(hubs), len(set(hubs)))
+
+ for spec in MODEL_REGISTRY.values():
+ self.assertTrue(spec.huggingface_url.startswith("https://huggingface.co/PulpBio/"))
+ self.assertTrue(spec.paper_url.startswith("https://arxiv.org/"))
+ self.assertIn(spec.huggingface_url, readme)
+ self.assertIn(spec.paper_url, readme)
+
+ def test_citation_guide_covers_models_and_publication_venues(self):
+ citations = (ROOT / "docs" / "CITATIONS.md").read_text(encoding="utf-8")
+ expected_venues = {
+ "FEMBA": "EMBC 2025",
+ "LUNA": "NeurIPS 2025",
+ "TinyMyo": "arXiv preprint",
+ "LuMamba": "EUSIPCO 2026",
+ "PanLUNA": "AICAS 2026",
+ }
+ for model, venue in expected_venues.items():
+ self.assertIn(f"## {model}", citations)
+ self.assertIn(venue, citations)
+
+ panluna_entry = citations.split("## PanLUNA", 1)[1]
+ self.assertIn("Benini, Luca", panluna_entry)
+
+ def test_lookup_is_case_insensitive(self):
+ self.assertEqual(get_model_spec("PanLUNA"), MODEL_REGISTRY["panluna"])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_repository_contracts.py b/tests/test_repository_contracts.py
new file mode 100644
index 0000000..b9f9a30
--- /dev/null
+++ b/tests/test_repository_contracts.py
@@ -0,0 +1,164 @@
+#*----------------------------------------------------------------------------*
+#* 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 *
+#*----------------------------------------------------------------------------*
+
+import ast
+import re
+import unittest
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+LOCAL_TARGET_PREFIXES = ("criterion.", "data_module.", "datasets.", "models.", "schedulers.", "tasks.")
+TASK_FILES = (
+ "finetune_regression_task_LuMamba.py",
+ "finetune_task.py",
+ "finetune_task_EMG.py",
+ "finetune_task_LUNA.py",
+ "finetune_task_PanLUNA.py",
+ "pretrain_task.py",
+ "pretrain_task_EMG.py",
+ "pretrain_task_LUNA.py",
+ "pretrain_task_LuMamba.py",
+ "pretrain_task_PanLUNA.py",
+)
+
+
+class RepositoryContractsTest(unittest.TestCase):
+ def test_local_hydra_targets_resolve(self):
+ target_pattern = re.compile(
+ r"^\s*_target_:\s*(?:'([^']+)'|\"([^\"]+)\"|([^\s#]+))",
+ re.MULTILINE,
+ )
+ for config_path in (ROOT / "config").rglob("*.yaml"):
+ matches = target_pattern.findall(config_path.read_text(encoding="utf-8"))
+ for match in matches:
+ target = next(value for value in match if value)
+ if "CHANGEME" in target or not target.startswith(LOCAL_TARGET_PREFIXES):
+ continue
+
+ module_name, attribute = target.rsplit(".", 1)
+ module_path = ROOT / f"{module_name.replace('.', '/')}.py"
+ self.assertTrue(module_path.is_file(), f"{config_path}: {target}")
+
+ tree = ast.parse(module_path.read_text(encoding="utf-8"))
+ definitions = {
+ node.name
+ for node in tree.body
+ if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef))
+ }
+ self.assertIn(attribute, definitions, f"{config_path}: {target}")
+
+ def test_hydra_package_directives_stay_on_the_first_line(self):
+ misplaced = []
+ for config_path in (ROOT / "config").rglob("*.yaml"):
+ lines = config_path.read_text(encoding="utf-8").splitlines()
+ directive_lines = [index for index, line in enumerate(lines) if "@package" in line]
+ if directive_lines and directive_lines != [0]:
+ misplaced.append(str(config_path.relative_to(ROOT)))
+ self.assertEqual(misplaced, [])
+
+ def test_training_steps_use_the_shared_batch_adapter(self):
+ for filename in TASK_FILES:
+ path = ROOT / "tasks" / filename
+ tree = ast.parse(path.read_text(encoding="utf-8"))
+ methods = {
+ node.name: node
+ for class_node in tree.body
+ if isinstance(class_node, ast.ClassDef)
+ for node in class_node.body
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+ }
+ for method_name in ("training_step", "validation_step", "test_step"):
+ if method_name not in methods:
+ continue
+ calls_adapter = any(
+ isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id == "as_signal_batch"
+ for node in ast.walk(methods[method_name])
+ )
+ self.assertTrue(calls_adapter, f"{filename}:{method_name}")
+
+ def test_cli_validates_environment_before_hydra_starts(self):
+ tree = ast.parse((ROOT / "run_train.py").read_text(encoding="utf-8"))
+ main_guard = next(
+ node
+ for node in tree.body
+ if isinstance(node, ast.If)
+ and isinstance(node.test, ast.Compare)
+ and isinstance(node.test.left, ast.Name)
+ and node.test.left.id == "__name__"
+ )
+ call_names = [
+ node.value.func.id
+ for node in main_guard.body
+ if isinstance(node, ast.Expr)
+ and isinstance(node.value, ast.Call)
+ and isinstance(node.value.func, ast.Name)
+ ]
+ self.assertLess(call_names.index("require_environment"), call_names.index("run"))
+
+ def test_onboarding_docs_do_not_link_to_missing_local_paths(self):
+ link_pattern = re.compile(r"\[[^\]]+\]\(([^)]+)\)")
+ paths = (
+ ROOT / "README.md",
+ ROOT / "CONTRIBUTING.md",
+ ROOT / "models" / "README.md",
+ ROOT / "docs" / "README.md",
+ ROOT / "docs" / "TRAINING.md",
+ ROOT / "docs" / "CITATIONS.md",
+ *(ROOT / "docs" / "model").glob("*.md"),
+ )
+ for path in paths:
+ for link in link_pattern.findall(path.read_text(encoding="utf-8")):
+ if "://" in link or link.startswith(("#", "mailto:")):
+ continue
+ local_path = (path.parent / link.split("#", 1)[0]).resolve()
+ self.assertTrue(local_path.exists(), f"{path.relative_to(ROOT)}: {link}")
+
+ def test_source_and_config_files_keep_the_apache_header(self):
+ source_roots = (
+ ROOT / "ARES",
+ ROOT / "biofoundation",
+ ROOT / "criterion",
+ ROOT / "data_module",
+ ROOT / "datasets",
+ ROOT / "make_datasets",
+ ROOT / "models",
+ ROOT / "schedulers",
+ ROOT / "tasks",
+ ROOT / "tests",
+ ROOT / "util",
+ )
+ paths = list(ROOT.glob("*.py"))
+ for source_root in source_roots:
+ paths.extend(source_root.rglob("*.py"))
+ paths.extend((ROOT / "config").rglob("*.yaml"))
+ paths.extend((ROOT / ".github").rglob("*.yml"))
+
+ missing = [
+ str(path.relative_to(ROOT))
+ for path in paths
+ if "SPDX-License-Identifier: Apache-2.0" not in path.read_text(encoding="utf-8")
+ ]
+ self.assertEqual(missing, [])
+
+
+if __name__ == "__main__":
+ unittest.main()
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-
- 
-
-