From 563f43b997c8fecda337ed67dd8ea7986e415321 Mon Sep 17 00:00:00 2001 From: Santiago Castro Dau Date: Wed, 3 Dec 2025 18:47:56 +0100 Subject: [PATCH 1/5] Add initial script for protein language modeling module --- protein_lm/modeling/scripts/__initi__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 protein_lm/modeling/scripts/__initi__.py diff --git a/protein_lm/modeling/scripts/__initi__.py b/protein_lm/modeling/scripts/__initi__.py new file mode 100644 index 0000000..e69de29 From 531ca7e5d1a7f554aa7c84207312d732073755c8 Mon Sep 17 00:00:00 2001 From: Santiago Castro Dau Date: Tue, 9 Dec 2025 13:39:17 +0100 Subject: [PATCH 2/5] Refactor import statements in lm.py --- protein_lm/modeling/models/mamba/lm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/protein_lm/modeling/models/mamba/lm.py b/protein_lm/modeling/models/mamba/lm.py index b9bc35d..2f7b984 100644 --- a/protein_lm/modeling/models/mamba/lm.py +++ b/protein_lm/modeling/models/mamba/lm.py @@ -8,7 +8,8 @@ import torch import torch.nn as nn from dataclasses import dataclass, field -from mamba_ssm.modules.mamba_simple import Mamba, Block +from mamba_ssm.modules.mamba_simple import Mamba +from mamba_ssm.modules.block import Block from mamba_ssm.utils.generation import GenerationMixin from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf From 78737b331fea2ff2f26ef04272d8f630e8c8a3a9 Mon Sep 17 00:00:00 2001 From: Santiago Castro Dau Date: Tue, 9 Dec 2025 14:09:02 +0100 Subject: [PATCH 3/5] Update load_ckpt function to include weights_only parameter in torch.load --- protein_lm/modeling/scripts/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/protein_lm/modeling/scripts/train.py b/protein_lm/modeling/scripts/train.py index c2af0da..304a181 100644 --- a/protein_lm/modeling/scripts/train.py +++ b/protein_lm/modeling/scripts/train.py @@ -141,7 +141,7 @@ def main(config_dict: DictConfig): ) def load_ckpt(ckpt_path, tokenizer, device): - ckpt = torch.load(ckpt_path) + ckpt = torch.load(ckpt_path, weights_only=False) model_state_dict = ckpt["model"] model_config = ckpt["config"] model_config.vocab_size = tokenizer.get_vocab_size() From 8223dc42d907bb2771b8681e9e77414228629e39 Mon Sep 17 00:00:00 2001 From: Santiago Castro Dau Date: Tue, 9 Dec 2025 14:22:03 +0100 Subject: [PATCH 4/5] Fix import path for RMSNorm and layer normalization functions --- protein_lm/modeling/models/mamba/lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/protein_lm/modeling/models/mamba/lm.py b/protein_lm/modeling/models/mamba/lm.py index 2f7b984..dff6972 100644 --- a/protein_lm/modeling/models/mamba/lm.py +++ b/protein_lm/modeling/models/mamba/lm.py @@ -14,7 +14,7 @@ from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf try: - from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn + from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn except ImportError: RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None From 83ea4398c3ed0c9b3d73cb9b9b3c97804cde1a2e Mon Sep 17 00:00:00 2001 From: Santiago Castro Dau Date: Tue, 9 Dec 2025 14:33:23 +0100 Subject: [PATCH 5/5] Add nn.Identity() to create_block function in lm.py --- protein_lm/modeling/models/mamba/lm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/protein_lm/modeling/models/mamba/lm.py b/protein_lm/modeling/models/mamba/lm.py index dff6972..ffc8d92 100644 --- a/protein_lm/modeling/models/mamba/lm.py +++ b/protein_lm/modeling/models/mamba/lm.py @@ -67,6 +67,7 @@ def create_block( block = Block( d_model, mixer_cls, + nn.Identity(), norm_cls=norm_cls, fused_add_norm=fused_add_norm, residual_in_fp32=residual_in_fp32,