Update transformers package

ci/tests/test_helix_mrna/test_helix_mrna_model.py

@@ -11,6 +11,17 @@ def helixmRNA(self):
         config = HelixmRNAConfig(batch_size=1, device=self.device, max_length=100)
         return HelixmRNA(configurer=config)
 
+    @pytest.fixture
+    def helixmRNA_eager(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        config = HelixmRNAConfig(
+            batch_size=1,
+            device=self.device,
+            max_length=100,
+            output_attentions=True,
+        )
+        return HelixmRNA(configurer=config)
+
     @pytest.fixture
     def mock_data(self, helixmRNA):
         input_sequences = [
@@ -43,6 +54,34 @@ def test_helix_mrna_get_embeddings(self, mock_data, helixmRNA):
             embeddings is not None
         ), f"Embeddings should not be None for sequence: {mock_data}"
 
+    def test_get_embeddings_output_attentions_requires_eager_config(
+        self, mock_data, helixmRNA
+    ):
+        with pytest.raises(
+            ValueError,
+            match=r"output_attentions=True requires the model to be loaded with eager",
+        ):
+            helixmRNA.get_embeddings(mock_data, output_attentions=True)
+
+    def test_get_embeddings_output_attentions_returns_attention_tensors(
+        self, helixmRNA_eager
+    ):
+        sequences = ["EACUEGG", "EACUEGG"]
+        dataset = helixmRNA_eager.process_data(sequences)
+
+        result = helixmRNA_eager.get_embeddings(dataset, output_attentions=True)
+
+        assert isinstance(result, tuple), "Expected tuple when output_attentions=True"
+        assert len(result) == 2
+        embeddings, attentions = result
+        assert isinstance(attentions, list)
+        assert len(attentions) >= 1
+        for attn in attentions:
+            assert attn.ndim == 4, (
+                f"Expected attention shape (batch, heads, seq, seq); "
+                f"got shape {attn.shape}"
+            )
+
     @pytest.mark.parametrize(
         "data, raise_exception",
         [

examples/notebooks/Hyena-DNA-Inference.ipynb

@@ -100,16 +100,13 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Filter: 100%|██████████| 461850/461850 [00:00<00:00, 477649.20 examples/s]\n",
-      "Filter: 100%|██████████| 48797/48797 [00:00<00:00, 492263.67 examples/s]\n"
+      "Filter: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 461850/461850 [00:00<00:00, 477649.20 examples/s]\n",
+      "Filter: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 48797/48797 [00:00<00:00, 492263.67 examples/s]\n"
      ]
     }
    ],
    "source": [
-    "from datasets import load_dataset\n",
-    "label = \"promoter_tata\"\n",
-    "\n",
-    "dataset = load_dataset(\"InstaDeepAI/nucleotide_transformer_downstream_tasks\",  trust_remote_code=True).filter(lambda x: x[\"task\"] == \"promoter_tata\")"
+    "from datasets import load_dataset\nlabel = \"promoter_tata\"\n\ndataset = load_dataset(\"InstaDeepAI/nucleotide_transformer_downstream_tasks_revised\", \"promoter_tata\")"
    ]
   },
   {
@@ -209,7 +206,7 @@
      "text": [
       "INFO:helical.models.hyena_dna.model:Succesfully prepared the HyenaDNA Dataset.\n",
       "INFO:helical.models.hyena_dna.model:Started getting embeddings:\n",
-      "Getting embeddings: 100%|██████████| 1102/1102 [00:03<00:00, 337.42it/s]\n",
+      "Getting embeddings: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 1102/1102 [00:03<00:00, 337.42it/s]\n",
       "INFO:helical.models.hyena_dna.model:Finished getting embeddings.\n"
      ]
     }
@@ -452,7 +449,7 @@
       "INFO:helical.models.hyena_dna.model:Processing data for HyenaDNA.\n",
       "INFO:helical.models.hyena_dna.model:Succesfully prepared the HyenaDNA Dataset.\n",
       "INFO:helical.models.hyena_dna.model:Started getting embeddings:\n",
-      "Getting embeddings: 100%|██████████| 125/125 [00:00<00:00, 439.64it/s]\n",
+      "Getting embeddings: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 125/125 [00:00<00:00, 439.64it/s]\n",
       "INFO:helical.models.hyena_dna.model:Finished getting embeddings.\n"
      ]
     }
@@ -683,7 +680,7 @@
        "\n",
        "#sk-container-id-1 label.sk-toggleable__label-arrow:before {\n",
        "  /* Arrow on the left of the label */\n",
-       "  content: \"▸\";\n",
+       "  content: \"\u25b8\";\n",
        "  float: left;\n",
        "  margin-right: 0.25em;\n",
        "  color: var(--sklearn-color-icon);\n",
@@ -728,7 +725,7 @@
        "}\n",
        "\n",
        "#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {\n",
-       "  content: \"▾\";\n",
+       "  content: \"\u25be\";\n",
        "}\n",
        "\n",
        "/* Pipeline/ColumnTransformer-specific style */\n",
@@ -1296,4 +1293,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

examples/notebooks/HyenaDNA-Fine-Tuning.ipynb

@@ -95,11 +95,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "label = \"promoter_tata\"\n",
-    "\n",
-    "dataset = load_dataset(\"InstaDeepAI/nucleotide_transformer_downstream_tasks\",  trust_remote_code=True).filter(lambda x: x[\"task\"] == \"promoter_tata\")\n",
-    "dataset_train = dataset[\"train\"]\n",
-    "dataset_test = dataset[\"test\"]"
+    "label = \"promoter_tata\"\n\ndataset = load_dataset(\"InstaDeepAI/nucleotide_transformer_downstream_tasks_revised\", \"promoter_tata\")\ndataset_train = dataset[\"train\"]\ndataset_test = dataset[\"test\"]"
    ]
   },
   {
@@ -184,8 +180,8 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Processing sequences: 100%|██████████| 5062/5062 [00:00<00:00, 10775.16it/s]\n",
-      "Processing sequences: 100%|██████████| 212/212 [00:00<00:00, 9641.03it/s]\n"
+      "Processing sequences: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 5062/5062 [00:00<00:00, 10775.16it/s]\n",
+      "Processing sequences: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 212/212 [00:00<00:00, 9641.03it/s]\n"
      ]
     }
    ],
@@ -217,26 +213,26 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Fine-Tuning: epoch 1/10: 100%|██████████| 507/507 [00:02<00:00, 191.25it/s, loss=0.654]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 729.64it/s, val_loss=0.602]\n",
-      "Fine-Tuning: epoch 2/10: 100%|██████████| 507/507 [00:02<00:00, 213.34it/s, loss=0.441]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 717.45it/s, val_loss=0.796]\n",
-      "Fine-Tuning: epoch 3/10: 100%|██████████| 507/507 [00:02<00:00, 211.52it/s, loss=0.453]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 716.82it/s, val_loss=0.597]\n",
-      "Fine-Tuning: epoch 4/10: 100%|██████████| 507/507 [00:02<00:00, 206.29it/s, loss=0.444]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 730.15it/s, val_loss=0.492]\n",
-      "Fine-Tuning: epoch 5/10: 100%|██████████| 507/507 [00:02<00:00, 232.99it/s, loss=0.439]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 724.66it/s, val_loss=0.424]\n",
-      "Fine-Tuning: epoch 6/10: 100%|██████████| 507/507 [00:01<00:00, 256.32it/s, loss=0.438]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 724.62it/s, val_loss=0.382]\n",
-      "Fine-Tuning: epoch 7/10: 100%|██████████| 507/507 [00:01<00:00, 258.73it/s, loss=0.436]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 715.60it/s, val_loss=0.351]\n",
-      "Fine-Tuning: epoch 8/10: 100%|██████████| 507/507 [00:02<00:00, 233.66it/s, loss=0.434]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 709.78it/s, val_loss=0.336]\n",
-      "Fine-Tuning: epoch 9/10: 100%|██████████| 507/507 [00:02<00:00, 240.10it/s, loss=0.432]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 714.25it/s, val_loss=0.328]\n",
-      "Fine-Tuning: epoch 10/10: 100%|██████████| 507/507 [00:02<00:00, 234.96it/s, loss=0.428]\n",
-      "Fine-Tuning Validation: 100%|██████████| 22/22 [00:00<00:00, 716.93it/s, val_loss=0.324]\n"
+      "Fine-Tuning: epoch 1/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 191.25it/s, loss=0.654]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 729.64it/s, val_loss=0.602]\n",
+      "Fine-Tuning: epoch 2/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 213.34it/s, loss=0.441]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 717.45it/s, val_loss=0.796]\n",
+      "Fine-Tuning: epoch 3/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 211.52it/s, loss=0.453]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 716.82it/s, val_loss=0.597]\n",
+      "Fine-Tuning: epoch 4/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 206.29it/s, loss=0.444]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 730.15it/s, val_loss=0.492]\n",
+      "Fine-Tuning: epoch 5/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 232.99it/s, loss=0.439]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 724.66it/s, val_loss=0.424]\n",
+      "Fine-Tuning: epoch 6/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:01<00:00, 256.32it/s, loss=0.438]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 724.62it/s, val_loss=0.382]\n",
+      "Fine-Tuning: epoch 7/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:01<00:00, 258.73it/s, loss=0.436]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 715.60it/s, val_loss=0.351]\n",
+      "Fine-Tuning: epoch 8/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 233.66it/s, loss=0.434]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 709.78it/s, val_loss=0.336]\n",
+      "Fine-Tuning: epoch 9/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 240.10it/s, loss=0.432]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 714.25it/s, val_loss=0.328]\n",
+      "Fine-Tuning: epoch 10/10: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 507/507 [00:02<00:00, 234.96it/s, loss=0.428]\n",
+      "Fine-Tuning Validation: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 716.93it/s, val_loss=0.324]\n"
      ]
     }
    ],
@@ -267,7 +263,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 22/22 [00:00<00:00, 782.08it/s]\n"
+      "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22/22 [00:00<00:00, 782.08it/s]\n"
      ]
     }
    ],
@@ -363,4 +359,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

examples/run_models/configs/geneformer_attention_config.yaml

@@ -0,0 +1,6 @@
+model_name: "gf-12L-40M-i2048"
+batch_size: 24
+emb_layer: -1
+emb_mode: "cell"
+device: "cpu"
+output_attentions: true

examples/run_models/configs/geneformer_config.yaml

@@ -2,4 +2,5 @@ model_name: "gf-12L-40M-i2048"
 batch_size: 24
 emb_layer: -1
 emb_mode: "cell"
-device: "cpu"
+device: "cpu"
+output_attentions: false

examples/run_models/run_geneformer.py

@@ -30,12 +30,6 @@ def run(cfg: DictConfig):
     embeddings = geneformer.get_embeddings(dataset, output_genes=True)
 
     print(embeddings)
-    embeddings, attention_weights = geneformer.get_embeddings(
-        dataset, output_attentions=True
-    )
-
-    print(embeddings)
-    print(attention_weights)
 
 
 if __name__ == "__main__":

examples/run_models/run_geneformer_attention.py

@@ -0,0 +1,28 @@
+from helical.models.geneformer import Geneformer, GeneformerConfig
+import hydra
+from omegaconf import DictConfig
+import anndata as ad
+
+
+@hydra.main(
+    version_base=None,
+    config_path="configs",
+    config_name="geneformer_attention_config",
+)
+def run(cfg: DictConfig):
+    geneformer_config = GeneformerConfig(**cfg)
+    geneformer = Geneformer(configurer=geneformer_config)
+
+    ann_data = ad.read_h5ad("./yolksac_human.h5ad")
+
+    dataset = geneformer.process_data(ann_data[:10, :100])
+    embeddings, attention_weights = geneformer.get_embeddings(
+        dataset, output_attentions=True
+    )
+
+    print(embeddings)
+    print(attention_weights)
+
+
+if __name__ == "__main__":
+    run()

helical/models/__init__.py

@@ -1 +1,13 @@
+# mamba_ssm (used by Caduceus and HelixmRNA) still imports the greedy/sample decoder
+# output names that transformers deleted; re-alias them here before any model loads.
+import transformers.generation as _tg
+
+if not hasattr(_tg, "GreedySearchDecoderOnlyOutput"):
+    from transformers.generation import GenerateDecoderOnlyOutput as _GDO
+
+    _tg.GreedySearchDecoderOnlyOutput = _GDO
+    _tg.SampleDecoderOnlyOutput = _GDO
+    del _GDO
+del _tg
+
 from .fine_tune.fine_tuning_heads import ClassificationHead, RegressionHead

helical/models/base_models.py

@@ -289,19 +289,34 @@ def save_model(self, path: str):
         path : str
             The path to save the model to.
         """
-        torch.save(self.model, path)
+        torch.save(self.model.state_dict(), path)
         LOGGER.info(f"Model saved to {path}")
 
     def load_model(self, path: str):
         """
         Load the model from a file.
 
+        Accepts both current state-dict checkpoints and legacy pickle checkpoints
+        (saved with the old ``torch.save(model, path)`` API).  The secure state-dict
+        load is attempted first; on failure we fall back to unpickling the legacy
+        full-model object and extracting its state dict.
+
         Parameters
         ----------
         path : str
             The path to load the model from.
         """
-        self.model = torch.load(path, weights_only=False)
+        try:
+            state_dict = torch.load(path, weights_only=True)
+        except Exception:
+            LOGGER.warning(
+                f"State-dict load failed for {path}; "
+                f"attempting to load as a legacy pickle checkpoint."
+            )
+            legacy = torch.load(path, weights_only=False)
+            state_dict = legacy.state_dict() if not isinstance(legacy, dict) else legacy
+
+        self.model.load_state_dict(state_dict)
         self.model.eval()
         self.fine_tuning_head.eval()
         LOGGER.info(f"Model loaded from {path}")

helical/models/caduceus/modeling_caduceus.py

@@ -312,6 +312,8 @@ class CaduceusPreTrainedModel(PreTrainedModel):
     base_model_prefix = "caduceus"
     supports_gradient_checkpointing = False
     _no_split_modules = ["BiMambaWrapper"]
+    # Caduceus handles weight tying internally (bidirectional_weight_tie), not via HF.
+    all_tied_weights_keys = {}
 
     def _init_weights(
         self,

helical/models/caduceus/pretrained_config.py

@@ -47,6 +47,8 @@ class CaduceusPretrainedConfig(PretrainedConfig):
         The complement map.
     """
 
+    model_type = "caduceus"
+
     def __init__(
         self,
         # From original MambaConfig

helical/models/geneformer/geneformer_config.py

@@ -32,10 +32,16 @@ class GeneformerConfig:
         The embedding mode to use. "cls" is only available for Geneformer v2 models, returning the embeddings of the cls token.
         For cell level embeddings, a mean across all embeddings excluding the cls token is returned.
         For gene level embeddings, each gene token embedding is returned along with the corresponding ensembl ID.
-    device : Literal["cpu", "cuda"], optional, default="cpu"
-        The device to use. Either use "cuda" or "cpu".
+    device : str, optional, default="cpu"
+        The device to use. Accepts any string torch.device accepts, e.g. "cpu",
+        "cuda", "cuda:0".
     nproc: int, optional, default=1
         Number of processes to use for data processing.
+    output_attentions : bool, optional, default=False
+        Whether to return attention weights from get_embeddings. Must be set at construction time:
+        True forces eager attention (required for attention output; flash_attention_2 and sdpa do
+        not support it), False uses flash_attention_2 when available, else sdpa. Note: eager
+        attention materialises the full O(seq²) matrix and may OOM on long sequences or large batches.
     custom_attr_name_dict : dict, optional, default=None
         A dictionary that contains the names of the custom attributes to be added to the dataset.
         The keys of the dictionary are the names of the custom attributes, and the values are the names of the columns in adata.obs.
@@ -71,8 +77,9 @@ def __init__(
         batch_size: int = 24,
         emb_layer: int = -1,
         emb_mode: Literal["cls", "cell", "gene"] = "cell",
-        device: Literal["cpu", "cuda"] = "cpu",
+        device: str = "cpu",
         nproc: int = 1,
+        output_attentions: bool = False,
         custom_attr_name_dict: Optional[dict] = None,
     ):
 
@@ -212,5 +219,6 @@ def __init__(
             "special_token": self.model_map[model_name]["special_token"],
             "embsize": self.model_map[model_name]["embsize"],
             "nproc": nproc,
+            "output_attentions": output_attentions,
             "custom_attr_name_dict": custom_attr_name_dict,
         }