Unify BioMod interfaces and centralize checkpoint compatibility

Use the bundled BioModEncoder, BioModLM and BioModLM_GO interfaces throughout. Remove obsolete CLI aliases and the external attention fallback. Centralize legacy configuration and state-dictionary translation in biomod_compat.py and document the shared asset. Model weights, paired configuration, datasets and unrelated files remain unchanged.

BioMod/README.md

@@ -156,6 +156,7 @@ $BIOMOD_ROOT/model_dependencies/              # Shared gene mappings, gene lists
 │   ├── pytorch.py
 │   ├── pytorch_old.py
 │   └── reversible.py
+├── biomod_compat.py                          # Legacy checkpoint and configuration compatibility
 ├── gene2ens.csv
 ├── selected_gene2vec_27k.npy
 └── selected_genes_27k.txt
@@ -351,6 +352,8 @@ files and shows which entries are missing.
 Check `BIOMOD_ROOT` for shell commands or `PROJECT_ROOT` in a notebook, and verify that the BioMod model weights, matching
 `.pkl` configuration, gene2vec file, gene list, gene mapping, and
 `pytorch_cont/` directory exist at the paths listed under Shared Assets.
+Keep `biomod_compat.py` alongside that directory so existing serialized
+configurations and checkpoint field names can be translated during loading.
 
 ### No space left on device
 

BioMod/biomod.py

@@ -65,7 +65,6 @@
     sys.path.insert(0, str(SOURCE_DIR))
 from models_biomod import (
     BioMod,
-    normalize_biomod_decoder_type,
     sanitize_biomod_config,
 )
 from biomod_dataset import BioModSCDataset, load_biomod_data
@@ -115,7 +114,7 @@
                     help='Reconstruct loss aggregation method: gene_mean=masked gene equal weight (current); module_mean=masked module equal weight (MAE style)')
 parser.add_argument("--recon_target_level", type=str, default="gene", choices=["gene", "module_mean_value"],
                     help='Reconstruction target granularity: gene=gene-by-gene reconstruction; module_mean_value=decoder directly outputs the module mean and calculates the loss with the target mean')
-parser.add_argument("--decoder_type", type=normalize_biomod_decoder_type, default="biomod", choices=["none", "simple", "biomod", "cross_attn"],
+parser.add_argument("--decoder_type", type=str, default="biomod", choices=["none", "simple", "biomod", "cross_attn"],
                     help='decoder type: none=directly use encoder to output predictions; simple=no token mixing; biomod=decoder self-attention; cross_attn=only masked->visible cross-attention')
 parser.add_argument("--decoder_gene_agg", action='store_true',
                     help='The reconstruction head is changed to module->gene aggregation and then gene-level reconstruction (off by default to facilitate comparison with the old module_head)')

BioMod/cell_annotation/train.py

@@ -1173,9 +1173,8 @@ def parse_args():
     parser.add_argument("--max_cells", type=int, default=None)
 
     parser.add_argument("--config_pkl", type=str, default=str(WEIGHTS_DIR / "BioMod.pkl"))
-    # Retain the legacy CLI spelling as a compatibility alias.
     parser.add_argument(
-        "--biomod_checkpoint", "--performer_checkpoint",
+        "--biomod_checkpoint",
         type=str,
         default=str(WEIGHTS_DIR / "BioMod.pth"),
         dest='biomod_checkpoint',

BioMod/models_biomod.py

@@ -15,25 +15,12 @@
 if os.path.isdir(model_dependencies_dir) and model_dependencies_dir not in sys.path:
     sys.path.insert(0, model_dependencies_dir)
 
-try:
-    from pytorch_cont.ding_models import (
-        Performer as BioModEncoder,
-        Always,
-        cast_tuple,
-    )
-except ImportError:
-    try:
-        from pytorch_cont.pytorch import (
-            Performer as BioModEncoder,
-            Always,
-            cast_tuple,
-        )
-    except ImportError:
-        from performer_pytorch.performer_pytorch import (
-            Performer as BioModEncoder,
-            Always,
-            cast_tuple,
-        )
+from biomod_compat import (
+    migrate_legacy_biomod_config,
+    migrate_legacy_biomod_state_dict,
+    normalize_biomod_decoder_type,
+)
+from pytorch_cont.pytorch import BioModEncoder, Always, cast_tuple
 import torch
 from torch import nn
 import torch.nn.functional as F
@@ -41,82 +28,6 @@
 import numpy as np
 
 
-def normalize_biomod_decoder_type(decoder_type: str) -> str:
-    """Normalize the legacy decoder option without changing the attention backend."""
-    if decoder_type == "performer":  # Legacy CLI and serialized-config alias.
-        return "biomod"
-    return decoder_type
-
-
-def _migrate_legacy_biomod_name(value: str) -> str:
-    """Translate serialized pre-rename names without changing their semantics."""
-    replacements = (
-        ("PATHCH", "MODULE"),
-        ("Pathch", "Module"),
-        ("pathch", "module"),
-        ("PATCHES", "MODULES"),
-        ("Patches", "Modules"),
-        ("patches", "modules"),
-        ("PATCH", "MODULE"),
-        ("Patch", "Module"),
-        ("patch", "module"),
-    )
-    for old, new in replacements:
-        value = value.replace(old, new)
-    return value
-
-
-def migrate_legacy_biomod_config(config: Dict[str, object]) -> Dict[str, object]:
-    """Return a BioMod-compatible copy of an older serialized configuration."""
-    def translate(value):
-        if isinstance(value, dict):
-            migrated = {}
-            for key, item in value.items():
-                new_key = _migrate_legacy_biomod_name(key) if isinstance(key, str) else key
-                if new_key in migrated:
-                    raise ValueError(
-                        f"Legacy configuration migration would duplicate key {new_key!r}"
-                    )
-                migrated[new_key] = translate(item)
-            return migrated
-        if isinstance(value, list):
-            return [translate(item) for item in value]
-        if isinstance(value, tuple):
-            return tuple(translate(item) for item in value)
-        if isinstance(value, str):
-            return _migrate_legacy_biomod_name(value)
-        return value
-
-    migrated = translate(dict(config))
-    if "decoder_type" in migrated:
-        migrated["decoder_type"] = normalize_biomod_decoder_type(migrated["decoder_type"])
-    return migrated
-
-
-def migrate_legacy_biomod_state_dict(state_dict):
-    """Translate old state-dict key names before PyTorch performs strict loading."""
-    try:
-        migrated = state_dict.__class__()
-    except TypeError:
-        migrated = {}
-
-    for key, value in state_dict.items():
-        new_key = _migrate_legacy_biomod_name(key) if isinstance(key, str) else key
-        if new_key in migrated:
-            # If both spellings are supplied, prefer the explicit new spelling.
-            if key == new_key:
-                migrated[new_key] = value
-            continue
-        migrated[new_key] = value
-
-    if hasattr(state_dict, "_metadata"):
-        migrated._metadata = {
-            _migrate_legacy_biomod_name(key): value
-            for key, value in state_dict._metadata.items()
-        }
-    return migrated
-
-
 def normalize_fusion_method(fusion_method: str) -> str:
     if fusion_method == "expression_only_mlp":
         return "add"

README.md

@@ -81,6 +81,9 @@ architecture and gene-module definitions.
 
 The model components are bundled under
 [`model_dependencies/pytorch_cont/`](model_dependencies/pytorch_cont/).
+Legacy serialized names are handled by
+[`biomod_compat.py`](model_dependencies/biomod_compat.py) when loading existing
+checkpoints; the original checkpoint files remain unchanged.
 Optional [Hallmark and Reactome source definitions](model_dependencies/biomod_data_sources/)
 are kept in `biomod_data_sources/`. They are separate from the default checkpoint
 runtime definitions in `biomod_data_all/`.
@@ -275,6 +278,7 @@ BioMod/                                                        # Project root: s
 │   │   ├── pytorch.py
 │   │   ├── pytorch_old.py
 │   │   └── reversible.py
+│   ├── biomod_compat.py                                      # Legacy checkpoint and configuration compatibility
 │   ├── gene2ens.csv
 │   ├── selected_gene2vec_27k.npy
 │   └── selected_genes_27k.txt

model_dependencies/biomod_compat.py

@@ -0,0 +1,111 @@
+"""Read historical BioMod serialized names without rewriting checkpoint files.
+
+Current interfaces use BioMod names. Legacy spellings are kept only here so
+existing configuration dictionaries and tensor state dictionaries still load.
+This module intentionally has no model or tensor-library imports.
+"""
+
+from typing import Dict
+
+
+def normalize_biomod_decoder_type(decoder_type: str) -> str:
+    """Normalize the legacy decoder option without changing the attention backend."""
+    if decoder_type == "performer":  # Legacy serialized-config alias.
+        return "biomod"
+    return decoder_type
+
+
+def _migrate_legacy_biomod_name(value: str) -> str:
+    """Translate serialized pre-rename names without changing their semantics."""
+    replacements = (
+        ("PATHCH", "MODULE"),
+        ("Pathch", "Module"),
+        ("pathch", "module"),
+        ("PATCHES", "MODULES"),
+        ("Patches", "Modules"),
+        ("patches", "modules"),
+        ("PATCH", "MODULE"),
+        ("Patch", "Module"),
+        ("patch", "module"),
+    )
+    for old, new in replacements:
+        value = value.replace(old, new)
+    return value
+
+
+def migrate_legacy_biomod_config(config: Dict[str, object]) -> Dict[str, object]:
+    """Return a BioMod-compatible copy of an older serialized configuration."""
+    def translate(value):
+        if isinstance(value, dict):
+            migrated = {}
+            for key, item in value.items():
+                new_key = _migrate_legacy_biomod_name(key) if isinstance(key, str) else key
+                if new_key in migrated:
+                    raise ValueError(
+                        f"Legacy configuration migration would duplicate key {new_key!r}"
+                    )
+                migrated[new_key] = translate(item)
+            return migrated
+        if isinstance(value, list):
+            return [translate(item) for item in value]
+        if isinstance(value, tuple):
+            return tuple(translate(item) for item in value)
+        if isinstance(value, str):
+            return _migrate_legacy_biomod_name(value)
+        return value
+
+    migrated = translate(dict(config))
+    if "decoder_type" in migrated:
+        migrated["decoder_type"] = normalize_biomod_decoder_type(migrated["decoder_type"])
+    return migrated
+
+
+def migrate_legacy_biomod_state_dict(state_dict):
+    """Translate old state-dict key names before PyTorch performs strict loading."""
+    try:
+        migrated = state_dict.__class__()
+    except TypeError:
+        migrated = {}
+
+    for key, value in state_dict.items():
+        new_key = _migrate_legacy_biomod_name(key) if isinstance(key, str) else key
+        if new_key in migrated:
+            # If both spellings are supplied, prefer the explicit new spelling.
+            if key == new_key:
+                migrated[new_key] = value
+            continue
+        migrated[new_key] = value
+
+    if hasattr(state_dict, "_metadata"):
+        migrated._metadata = {
+            _migrate_legacy_biomod_name(key): value
+            for key, value in state_dict._metadata.items()
+        }
+    return migrated
+
+
+def migrate_legacy_encoder_prefix(state_dict, prefix: str) -> None:
+    """Translate one legacy language-model subtree in place before loading.
+
+    Explicit current keys take precedence if both spellings are present.
+    The caller supplies the full parent prefix, including its trailing dot.
+    """
+    old_root = prefix + "performer"
+    new_root = prefix + "biomod"
+    old_prefix = old_root + "."
+    new_prefix = new_root + "."
+    for key in list(state_dict):
+        if isinstance(key, str) and key.startswith(old_prefix):
+            new_key = new_prefix + key[len(old_prefix):]
+            if new_key not in state_dict:
+                state_dict[new_key] = state_dict[key]
+            del state_dict[key]
+
+    metadata = getattr(state_dict, "_metadata", None)
+    if metadata is not None:
+        for key in list(metadata):
+            if key == old_root or key.startswith(old_prefix):
+                new_key = new_root + key[len(old_root):]
+                if new_key not in metadata:
+                    metadata[new_key] = metadata[key]
+                del metadata[key]

model_dependencies/pytorch_cont/__init__.py

@@ -1,2 +1,2 @@
-# Export the bundled Performer attention implementation.
-from pytorch_cont.pytorch import PerformerLM, Performer, FastAttention, SelfAttention
+# Export the bundled BioMod attention implementation.
+from pytorch_cont.pytorch import BioModLM, BioModEncoder, FastAttention, SelfAttention

model_dependencies/pytorch_cont/ding_models.py

@@ -1,3 +1,4 @@
+from biomod_compat import migrate_legacy_encoder_prefix
 from pytorch_cont.pytorch import *
 import torch.nn as nn
 from torch_geometric.nn import SGConv
@@ -125,7 +126,7 @@ def get_similarity_network(gene_list, data_path, num_similar_genes_go_graph):
 
 
 
-class PerformerLM_GO(nn.Module):
+class BioModLM_GO(nn.Module):
     def __init__(
         self,
         *,
@@ -189,17 +190,28 @@ class PerformerLM_GO(nn.Module):
                                          gene2vec_file = gene2vec_file)
         print('dropout')
         self.dropout = nn.Dropout(emb_dropout)
-        print('performer')
-        self.performer = Performer(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
+        print('BioMod')
+        self.biomod = BioModEncoder(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
         self.norm = nn.LayerNorm(dim)
         # self.to_out = nn.Linear(dim, num_tokens) if not tie_embed else None
         self.to_out = nn.Linear(dim, 1) if not tie_embed else None
 
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict,
+        missing_keys, unexpected_keys, error_msgs,
+    ):
+        """Normalize historical encoder prefixes before loading parameters."""
+        migrate_legacy_encoder_prefix(state_dict, prefix)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs,
+        )
+
     def check_redraw_projections(self):
-        self.performer.check_redraw_projections()
+        self.biomod.check_redraw_projections()
 
     def fix_projection_matrices_(self):
-        self.performer.fix_projection_matrices_()
+        self.biomod.fix_projection_matrices_()
 
     def forward(self, x, return_encodings = False, output_attentions = False, **kwargs):
         b, n, device = *x.shape, x.device
@@ -213,11 +225,11 @@ class PerformerLM_GO(nn.Module):
         x += self.go_conv(x)
         x = self.dropout(x)
 
-        # performer layers
+        # BioMod encoder layers
         layer_pos_emb = self.layer_pos_emb(x)
 
         if output_attentions:
-            x, attn_weights = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x, attn_weights = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
             # norm and to logits
             x = self.norm(x)
             if return_encodings:
@@ -228,7 +240,7 @@ class PerformerLM_GO(nn.Module):
 
             return (x @ self.token_emb.weight.t()), attn_weights
         else:
-            x = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
 
             # norm and to logits
             x = self.norm(x)
@@ -318,17 +330,17 @@ class DualEncoderSCFM(nn.Module):
                                          gene2vec_file = gene2vec_file)
         self.dropout = nn.Dropout(emb_dropout)
 
-        self.mini_encoder = Performer(base_dim, mini_enc_depth, mini_enc_heads, mini_enc_dim_head, 
+        self.mini_encoder = BioModEncoder(base_dim, mini_enc_depth, mini_enc_heads, mini_enc_dim_head, 
                                       local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
 
         self.base_to_large = nn.Linear(base_dim, large_dim)
         self.large_in_norm = nn.LayerNorm(large_dim)
-        self.large_encoder = Performer(large_dim, large_enc_depth, large_enc_heads, large_enc_dim_head, 
+        self.large_encoder = BioModEncoder(large_dim, large_enc_depth, large_enc_heads, large_enc_dim_head, 
                                        local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
         self.large_to_base = nn.Linear(large_dim, base_dim)
         self.large_to_base_out_norm = nn.LayerNorm(base_dim)
         
-        self.decoder = Performer(base_dim, dec_depth, dec_heads, dec_dim_head, 
+        self.decoder = BioModEncoder(base_dim, dec_depth, dec_heads, dec_dim_head, 
                                        local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
         self.decode_norm = nn.LayerNorm(base_dim)
         

model_dependencies/pytorch_cont/pytorch.py

@@ -1,3 +1,4 @@
+from biomod_compat import migrate_legacy_encoder_prefix
 import math
 from pathlib import Path
 import numpy as np
@@ -238,7 +239,7 @@ class FastAttention(nn.Module):
                 import fast_transformers.causal_product.causal_product_cuda
                 self.causal_linear_fn = partial(causal_linear_attention)
             except ImportError:
-                print('unable to import cuda code for auto-regressive Performer. will default to the memory inefficient non-cuda version')
+                print('unable to import cuda code for autoregressive BioMod attention. will default to the memory inefficient non-cuda version')
                 self.causal_linear_fn = causal_linear_attention_noncuda
 
     @torch.no_grad()
@@ -472,16 +473,16 @@ class Gene2VecPositionalEmbedding(nn.Module):
         t = torch.arange(x.shape[1], device=x.device)
         return self.emb(t)
 
-# performer
+# BioMod encoder
 
-class Performer(nn.Module):
+class BioModEncoder(nn.Module):
     def __init__(
         self,
         dim,                                # dimension
         depth,                              # layers
         heads,                              # heads
         dim_head,                           # dim of head
-        local_attn_heads = 0,               # num of local attention heads, (heads - local_attn_heads) is num of global performers
+        local_attn_heads = 0,               # num of local attention heads, (heads - local_attn_heads) is number of global attention heads
         local_window_size = 256,            # window size of local attention
         causal = False,                     # autoregressive or not
         ff_mult = 4,                        # dim of intermediate features after attention / dim of input features
@@ -576,7 +577,7 @@ class TwoLayerMLP(nn.Module):
         x = self.fc2(x)
         return x
 
-class PerformerLM(nn.Module):
+class BioModLM(nn.Module):
     def __init__(
         self,
         *,
@@ -626,17 +627,28 @@ class PerformerLM(nn.Module):
             self.layer_pos_emb = Always(None)
         print('dropout')
         self.dropout = nn.Dropout(emb_dropout)
-        print('performer')
-        self.performer = Performer(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
+        print('BioMod')
+        self.biomod = BioModEncoder(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
         self.norm = nn.LayerNorm(dim)
         # self.to_out = nn.Linear(dim, num_tokens) if not tie_embed else None
         self.to_out = nn.Linear(dim, 1) if not tie_embed else None
 
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict,
+        missing_keys, unexpected_keys, error_msgs,
+    ):
+        """Normalize historical encoder prefixes before loading parameters."""
+        migrate_legacy_encoder_prefix(state_dict, prefix)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs,
+        )
+
     def check_redraw_projections(self):
-        self.performer.check_redraw_projections()
+        self.biomod.check_redraw_projections()
 
     def fix_projection_matrices_(self):
-        self.performer.fix_projection_matrices_()
+        self.biomod.fix_projection_matrices_()
 
     def forward(self, x, return_encodings = False, output_attentions = False, **kwargs):
         b, n, device = *x.shape, x.device
@@ -649,11 +661,11 @@ class PerformerLM(nn.Module):
         x += self.pos_emb(x)
         x = self.dropout(x)
 
-        # performer layers
+        # BioMod encoder layers
         layer_pos_emb = self.layer_pos_emb(x)
 
         if output_attentions:
-            x, attn_weights = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x, attn_weights = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
             # norm and to logits
             x = self.norm(x)
             if return_encodings:
@@ -664,7 +676,7 @@ class PerformerLM(nn.Module):
 
             return (x @ self.token_emb.weight.t()), attn_weights
         else:
-            x = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
 
             # norm and to logits
             x = self.norm(x)

model_dependencies/pytorch_cont/pytorch_old.py

@@ -1,3 +1,4 @@
+from biomod_compat import migrate_legacy_encoder_prefix
 import math
 from pathlib import Path
 import numpy as np
@@ -238,7 +239,7 @@ class FastAttention(nn.Module):
                 import fast_transformers.causal_product.causal_product_cuda
                 self.causal_linear_fn = partial(causal_linear_attention)
             except ImportError:
-                print('unable to import cuda code for auto-regressive Performer. will default to the memory inefficient non-cuda version')
+                print('unable to import cuda code for autoregressive BioMod attention. will default to the memory inefficient non-cuda version')
                 self.causal_linear_fn = causal_linear_attention_noncuda
 
     @torch.no_grad()
@@ -462,16 +463,16 @@ class Gene2VecPositionalEmbedding(nn.Module):
         t = torch.arange(x.shape[1], device=x.device)
         return self.emb(t)
 
-# performer
+# BioMod encoder
 
-class Performer(nn.Module):
+class BioModEncoder(nn.Module):
     def __init__(
         self,
         dim,                                # dimension
         depth,                              # layers
         heads,                              # heads
         dim_head,                           # dim of head
-        local_attn_heads = 0,               # num of local attention heads, (heads - local_attn_heads) is num of global performers
+        local_attn_heads = 0,               # num of local attention heads, (heads - local_attn_heads) is number of global attention heads
         local_window_size = 256,            # window size of local attention
         causal = False,                     # autoregressive or not
         ff_mult = 4,                        # dim of intermediate features after attention / dim of input features
@@ -566,7 +567,7 @@ class TwoLayerMLP(nn.Module):
         x = self.fc2(x)
         return x
 
-class PerformerLM(nn.Module):
+class BioModLM(nn.Module):
     def __init__(
         self,
         *,
@@ -616,17 +617,28 @@ class PerformerLM(nn.Module):
             self.layer_pos_emb = Always(None)
         print('dropout')
         self.dropout = nn.Dropout(emb_dropout)
-        print('performer')
-        self.performer = Performer(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
+        print('BioMod')
+        self.biomod = BioModEncoder(dim, depth, heads, dim_head, local_attn_heads, local_window_size, causal, ff_mult, nb_features, feature_redraw_interval, reversible, ff_chunks, generalized_attention, kernel_fn, use_scalenorm, use_rezero, ff_glu, ff_dropout, attn_dropout, cross_attend, no_projection, auto_check_redraw, qkv_bias)
         self.norm = nn.LayerNorm(dim)
         # self.to_out = nn.Linear(dim, num_tokens) if not tie_embed else None
         self.to_out = nn.Linear(dim, 1) if not tie_embed else None
 
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict,
+        missing_keys, unexpected_keys, error_msgs,
+    ):
+        """Normalize historical encoder prefixes before loading parameters."""
+        migrate_legacy_encoder_prefix(state_dict, prefix)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs,
+        )
+
     def check_redraw_projections(self):
-        self.performer.check_redraw_projections()
+        self.biomod.check_redraw_projections()
 
     def fix_projection_matrices_(self):
-        self.performer.fix_projection_matrices_()
+        self.biomod.fix_projection_matrices_()
 
     def forward(self, x, return_encodings = False, output_attentions = False, **kwargs):
         b, n, device = *x.shape, x.device
@@ -639,11 +651,11 @@ class PerformerLM(nn.Module):
         x += self.pos_emb(x)
         x = self.dropout(x)
 
-        # performer layers
+        # BioMod encoder layers
         layer_pos_emb = self.layer_pos_emb(x)
 
         if output_attentions:
-            x, attn_weights = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x, attn_weights = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
             # norm and to logits
             x = self.norm(x)
             if return_encodings:
@@ -654,7 +666,7 @@ class PerformerLM(nn.Module):
 
             return (x @ self.token_emb.weight.t()), attn_weights
         else:
-            x = self.performer(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
+            x = self.biomod(x, pos_emb = layer_pos_emb, output_attentions = output_attentions, **kwargs)
 
             # norm and to logits
             x = self.norm(x)