Replace with clean markdown card

README.md

@@ -10,18 +10,16 @@ tags:
   - braindecode
   - foundation-model
   - convolutional
-  - transformer
 ---
 
 # Labram
 
-Labram from Jiang, W B et al (2024) .
+Labram from Jiang, W B et al (2024) [Jiang2024].
 
-> **Architecture-only repository.** This repo documents the
+> **Architecture-only repository.** Documents the
 > `braindecode.models.Labram` class. **No pretrained weights are
-> distributed here** — instantiate the model and train it on your own
-> data, or fine-tune from a published foundation-model checkpoint
-> separately.
+> distributed here.** Instantiate the model and train it on your own
+> data.
 
 ## Quick start
 
@@ -40,241 +38,59 @@ model = Labram(
 )
 ```
 
-The signal-shape arguments above are example defaults — adjust them
-to match your recording.
+The signal-shape arguments above are illustrative defaults — adjust to
+match your recording.
 
 ## Documentation
-
-- Full API reference (parameters, references, architecture figure):
-  <https://braindecode.org/stable/generated/braindecode.models.Labram.html>
-- Interactive browser with live instantiation:
+- Full API reference: <https://braindecode.org/stable/generated/braindecode.models.Labram.html>
+- Interactive browser (live instantiation, parameter counts):
   <https://huggingface.co/spaces/braindecode/model-explorer>
 - Source on GitHub: <https://github.com/braindecode/braindecode/blob/master/braindecode/models/labram.py#L196>
 
-## Architecture description
-
-The block below is the rendered class docstring (parameters,
-references, architecture figure where available).
-
-<div class='bd-doc'><main>
-<p>Labram from Jiang, W B et al (2024) [Jiang2024]_.</p>
-<span style="display:inline-block;padding:2px 8px;border-radius:4px;background:#5cb85c;color:white;font-size:11px;font-weight:600;margin-right:4px;">Convolution</span><span style="display:inline-block;padding:2px 8px;border-radius:4px;background:#d9534f;color:white;font-size:11px;font-weight:600;margin-right:4px;">Foundation Model</span>
-
-
-
- .. figure:: https://arxiv.org/html/2405.18765v1/x1.png
-     :align: center
-     :alt: Labram Architecture.
-
- Large Brain Model for Learning Generic Representations with Tremendous
- EEG Data in BCI from [Jiang2024]_.
-
- This is an **adaptation** of the code [Code2024]_ from the Labram model.
-
- The model is transformer architecture with **strong** inspiration from
- BEiTv2 [BeiTv2]_.
-
- The models can be used in two modes:
-
- - Neural Tokenizer: Design to get an embedding layers (e.g. classification).
- - Neural Decoder: To extract the ampliture and phase outputs with a VQSNP.
-
- The braindecode's modification is to allow the model to be used in
- with an input shape of (batch, n_chans, n_times), if neural tokenizer
- equals True. The original implementation uses (batch, n_chans, n_patches,
- patch_size) as input with static segmentation of the input data.
-
- The models have the following sequence of steps::
-
-     if neural tokenizer:
-         - SegmentPatch: Segment the input data in patches;
-         - TemporalConv: Apply a temporal convolution to the segmented data;
-         - Residual adding cls, temporal and position embeddings (optional);
-         - WindowsAttentionBlock: Apply a windows attention block to the data;
-         - LayerNorm: Apply layer normalization to the data;
-         - Linear: An head linear layer to transformer the data into classes.
-
-     else:
-         - PatchEmbed: Apply a patch embedding to the input data;
-         - Residual adding cls, temporal and position embeddings (optional);
-         - WindowsAttentionBlock: Apply a windows attention block to the data;
-         - LayerNorm: Apply layer normalization to the data;
-         - Linear: An head linear layer to transformer the data into classes.
-
- .. important::
-    **Pre-trained Weights Available**
-
-    This model has pre-trained weights available on the Hugging Face Hub.
-    You can load them using:
-
-    .. code:: python
-        from braindecode.models import Labram
-
-        # Load pre-trained model from Hugging Face Hub
-        model = Labram.from_pretrained("braindecode/labram-pretrained")
-
-    To push your own trained model to the Hub:
-
-    .. code:: python
-        # After training your model
-        model.push_to_hub(
-            repo_id="username/my-labram-model", commit_message="Upload trained Labram model"
-        )
-
-    Requires installing ``braindecode[hug]`` for Hub integration.
-
- .. versionadded:: 0.9
-
-
- Examples
- --------
- Load pre-trained weights::
-
-     >>> import torch
-     >>> from braindecode.models import Labram
-     >>> model = Labram(n_times=1600, n_chans=64, n_outputs=4)
-     >>> url = "https://huggingface.co/braindecode/Labram-Braindecode/blob/main/braindecode_labram_base.pt"
-     >>> state = torch.hub.load_state_dict_from_url(url, progress=True)
-     >>> model.load_state_dict(state)
-
-
- Parameters
- ----------
- patch_size : int
-     The size of the patch to be used in the patch embedding.
- learned_patcher : bool
-     Whether to use a learned patch embedding (via a convolutional layer) or a fixed patch embedding (via rearrangement).
- embed_dim : int
-     The dimension of the embedding.
- conv_in_channels : int
-     The number of convolutional input channels.
- conv_out_channels : int
-     The number of convolutional output channels.
- num_layers :  int (default=12)
-     The number of attention layers of the model.
- num_heads : int (default=10)
-     The number of attention heads.
- mlp_ratio : float (default=4.0)
-     The expansion ratio of the mlp layer
- qkv_bias :  bool (default=False)
-     If True, add a learnable bias to the query, key, and value tensors.
- qk_norm : Pytorch Normalize layer (default=nn.LayerNorm)
-     If not None, apply LayerNorm to the query and key tensors.
-     Default is nn.LayerNorm for better weight transfer from original LaBraM.
-     Set to None to disable Q,K normalization.
- qk_scale : float (default=None)
-     If not None, use this value as the scale factor. If None,
-     use head_dim**-0.5, where head_dim = dim // num_heads.
- drop_prob : float (default=0.0)
-     Dropout rate for the attention weights.
- attn_drop_prob : float (default=0.0)
-     Dropout rate for the attention weights.
- drop_path_prob : float (default=0.0)
-     Dropout rate for the attention weights used on DropPath.
- norm_layer : Pytorch Normalize layer (default=nn.LayerNorm)
-     The normalization layer to be used.
- init_values : float (default=0.1)
-     If not None, use this value to initialize the gamma_1 and gamma_2
-     parameters for residual scaling. Default is 0.1 for better weight
-     transfer from original LaBraM. Set to None to disable.
- use_abs_pos_emb : bool (default=True)
-     If True, use absolute position embedding.
- use_mean_pooling : bool (default=True)
-     If True, use mean pooling.
- init_scale : float (default=0.001)
-     The initial scale to be used in the parameters of the model.
- neural_tokenizer : bool (default=True)
-     The model can be used in two modes: Neural Tokenizer or Neural Decoder.
- attn_head_dim : bool (default=None)
-     The head dimension to be used in the attention layer, to be used only
-     during pre-training.
- activation: nn.Module, default=nn.GELU
-     Activation function class to apply. Should be a PyTorch activation
-     module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.GELU``.
-
- References
- ----------
- .. [Jiang2024] Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024, May.
-    Large Brain Model for Learning Generic Representations with Tremendous
-    EEG Data in BCI. The Twelfth International Conference on Learning
-    Representations, ICLR.
- .. [Code2024] Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024. Labram
-    Large Brain Model for Learning Generic Representations with Tremendous
-    EEG Data in BCI. GitHub https://github.com/935963004/LaBraM
-    (accessed 2024-03-02)
- .. [BeiTv2] Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, Furu Wei. 2024.
-    BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers.
-    arXiv:2208.06366 [cs.CV]
-
- .. rubric:: Hugging Face Hub integration
-
- When the optional ``huggingface_hub`` package is installed, all models
- automatically gain the ability to be pushed to and loaded from the
- Hugging Face Hub. Install with::
-
-     pip install braindecode[hub]
-
- **Pushing a model to the Hub:**
-
- .. code::
-     from braindecode.models import Labram
-
-     # Train your model
-     model = Labram(n_chans=22, n_outputs=4, n_times=1000)
-     # ... training code ...
-
-     # Push to the Hub
-     model.push_to_hub(
-         repo_id="username/my-labram-model",
-         commit_message="Initial model upload",
-     )
-
- **Loading a model from the Hub:**
-
- .. code::
-     from braindecode.models import Labram
-
-     # Load pretrained model
-     model = Labram.from_pretrained("username/my-labram-model")
-
-     # Load with a different number of outputs (head is rebuilt automatically)
-     model = Labram.from_pretrained("username/my-labram-model", n_outputs=4)
-
- **Extracting features and replacing the head:**
 
- .. code::
-     import torch
+## Architecture
 
-     x = torch.randn(1, model.n_chans, model.n_times)
-     # Extract encoder features (consistent dict across all models)
-     out = model(x, return_features=True)
-     features = out["features"]
+![Labram architecture](https://arxiv.org/html/2405.18765v1/x1.png)
 
-     # Replace the classification head
-     model.reset_head(n_outputs=10)
 
- **Saving and restoring full configuration:**
+## Parameters
 
- .. code::
-     import json
+| Parameter | Type | Description |
+|---|---|---|
+| `patch_size` | int | The size of the patch to be used in the patch embedding. |
+| `learned_patcher` | bool | Whether to use a learned patch embedding (via a convolutional layer) or a fixed patch embedding (via rearrangement). |
+| `embed_dim` | int | The dimension of the embedding. |
+| `conv_in_channels` | int | The number of convolutional input channels. |
+| `conv_out_channels` | int | The number of convolutional output channels. |
+| `num_layers` |  int (default=12) | The number of attention layers of the model. |
+| `num_heads` | int (default=10) | The number of attention heads. |
+| `mlp_ratio` | float (default=4.0) | The expansion ratio of the mlp layer |
+| `qkv_bias` |  bool (default=False) | If True, add a learnable bias to the query, key, and value tensors. |
+| `qk_norm` | Pytorch Normalize layer (default=nn.LayerNorm) | If not None, apply LayerNorm to the query and key tensors. Default is nn.LayerNorm for better weight transfer from original LaBraM. Set to None to disable Q,K normalization. |
+| `qk_scale` | float (default=None) | If not None, use this value as the scale factor. If None, use head_dim**-0.5, where head_dim = dim // num_heads. |
+| `drop_prob` | float (default=0.0) | Dropout rate for the attention weights. |
+| `attn_drop_prob` | float (default=0.0) | Dropout rate for the attention weights. |
+| `drop_path_prob` | float (default=0.0) | Dropout rate for the attention weights used on DropPath. |
+| `norm_layer` | Pytorch Normalize layer (default=nn.LayerNorm) | The normalization layer to be used. |
+| `init_values` | float (default=0.1) | If not None, use this value to initialize the gamma_1 and gamma_2 parameters for residual scaling. Default is 0.1 for better weight transfer from original LaBraM. Set to None to disable. |
+| `use_abs_pos_emb` | bool (default=True) | If True, use absolute position embedding. |
+| `use_mean_pooling` | bool (default=True) | If True, use mean pooling. |
+| `init_scale` | float (default=0.001) | The initial scale to be used in the parameters of the model. |
+| `neural_tokenizer` | bool (default=True) | The model can be used in two modes: Neural Tokenizer or Neural Decoder. |
+| `attn_head_dim` | bool (default=None) | The head dimension to be used in the attention layer, to be used only during pre-training. |
+| `activation: nn.Module, default=nn.GELU` | — | Activation function class to apply. Should be a PyTorch activation module class like `nn.ReLU` or `nn.ELU`. Default is `nn.GELU`. |
 
-     config = model.get_config()            # all __init__ params
-     with open("config.json", "w") as f:
-         json.dump(config, f)
 
-     model2 = Labram.from_config(config)    # reconstruct (no weights)
+## References
 
- All model parameters (both EEG-specific and model-specific such as
- dropout rates, activation functions, number of filters) are automatically
- saved to the Hub and restored when loading.
+1. Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024, May. Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. The Twelfth International Conference on Learning Representations, ICLR.
+2. Wei-Bang Jiang, Li-Ming Zhao, Bao-Liang Lu. 2024. Labram Large Brain Model for Learning Generic Representations with Tremendous EEG Data in BCI. GitHub https://github.com/935963004/LaBraM (accessed 2024-03-02)
+3. Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, Furu Wei. 2024. BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers. arXiv:2208.06366 [cs.CV]
 
- See :ref:`load-pretrained-models` for a complete tutorial.</main>
-</div>
 
 ## Citation
 
-Please cite both the original paper for this architecture (see the
-*References* section above) and braindecode:
+Cite the original architecture paper (see *References* above) and braindecode:
 
 ```bibtex
 @article{aristimunha2025braindecode,