Upload model

chatNT_config.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+
+from transformers import PretrainedConfig
+
+from genomics_research.biobrain_p1.porting_to_pytorch.configs.esm_config import (
+    ESMTransformerConfig,
+)
+from genomics_research.biobrain_p1.porting_to_pytorch.configs.gpt_config import (
+    GptConfig,
+)
+from genomics_research.biobrain_p1.porting_to_pytorch.configs.perceiver_resampler_config import (  # noqa
+    PerceiverResamplerConfig,
+)
+
+
+@dataclass
+class ChatNTConfig(PretrainedConfig):
+    model_type = "ChatNT"
+
+    def __init__(self, **kwargs):  # type: ignore
+        self.gpt_config: GptConfig = kwargs.get("gpt_config", GptConfig(32000, 3))
+        self.esm_config: ESMTransformerConfig = kwargs.get(
+            "esm_config", ESMTransformerConfig(4000, 1, 4)
+        )
+        self.perceiver_resampler_config: PerceiverResamplerConfig = kwargs.get(
+            "perceiver_resampler_config", PerceiverResamplerConfig()
+        )
+        self.seq_token_id: int = kwargs.get("seq_token_id", 32000)
+        self.bio_pad_token_id: int = kwargs.get("bio_pad_token_id", 1)
+        self.english_pad_token_id: int = kwargs.get("english_pad_token_id", 2)
+        super().__init__(**kwargs)
+
+    def to_dict(self):  # type: ignore
+        print("(debug) Going into ChatNTConfig to_dict")
+        output = super().to_dict()
+
+        def serialize(obj):  # type: ignore
+            return obj.to_dict() if hasattr(obj, "to_dict") else vars(obj)
+
+        print("(debug) Before serialize gpt_config to_dict")
+        output["gpt_config"] = serialize(self.gpt_config)  # type: ignore
+        print("(debug) Before serialize esm_config to_dict")
+        output["esm_config"] = serialize(self.esm_config)  # type: ignore
+        print("(debug) Before serialize perceiver_resampler_config to_dict")
+        output["perceiver_resampler_config"] = serialize(  # type: ignore
+            self.perceiver_resampler_config
+        )
+        print("(debug) after serializing all ")
+        print("(debug) output : ", output)
+        return output

config.json

@@ -0,0 +1,85 @@
+{
+  "architectures": [
+    "TorchMultiOmicsModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "chatNT_config.ChatNTConfig",
+    "AutoModel": "multi_omics_model.TorchMultiOmicsModel"
+  },
+  "bio_pad_token_id": 1,
+  "english_pad_token_id": 2,
+  "esm_config": {
+    "add_bias_ffn": false,
+    "add_bias_kv": false,
+    "alphabet_size": 4107,
+    "attention_heads": 16,
+    "attention_maps_to_save": [],
+    "bias_word_embedding": false,
+    "emb_layer_norm_before": false,
+    "embed_dim": 1024,
+    "embed_scale": 1.0,
+    "embeddings_layers_to_save": [
+      21
+    ],
+    "ffn_activation_name": "swish",
+    "ffn_embed_dim": 4096,
+    "key_size": 64,
+    "layer_norm_eps": 1e-05,
+    "lm_head": "roberta",
+    "mask_before_attention": false,
+    "mask_token_id": 2,
+    "masking_prob": 0.0,
+    "masking_ratio": 0.0,
+    "max_positions": 2048,
+    "num_layers": 29,
+    "pad_token_id": 1,
+    "positional_embedding": null,
+    "pre_layer_norm": true,
+    "rescaling_factor": null,
+    "token_dropout": false,
+    "use_glu_in_ffn": true,
+    "use_gradient_checkpointing": false,
+    "use_rotary_embedding": true
+  },
+  "gpt_config": {
+    "add_bias_attn": false,
+    "add_bias_ffn": false,
+    "add_bias_lm_head": false,
+    "embed_dim": 4096,
+    "eos_token_id": 2,
+    "ffn_activation_name": "silu",
+    "ffn_embed_dim": 11008,
+    "norm_type": "RMS_norm",
+    "num_heads": 32,
+    "num_kv_heads": 32,
+    "num_layers": 32,
+    "parallel_attention_ff": false,
+    "rms_norm_eps": 1e-06,
+    "rope_config": {
+      "dim": 128,
+      "max_seq_len": 2048,
+      "theta": 10000.0
+    },
+    "use_glu_in_ffn": true,
+    "use_gradient_checkpointing": false,
+    "vocab_size": 32000
+  },
+  "model_type": "ChatNT",
+  "perceiver_resampler_config": {
+    "add_bias_ffn": true,
+    "add_bias_kv": false,
+    "attention_heads": 32,
+    "emb_layer_norm_before": false,
+    "embed_dim": 4096,
+    "ffn_activation_name": "gelu-no-approx",
+    "ffn_embed_dim": 11008,
+    "key_size": 128,
+    "num_layers": 3,
+    "resampled_length": 64,
+    "use_glu_in_ffn": false,
+    "use_gradient_checkpointing": false
+  },
+  "seq_token_id": 32000,
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.1"
+}

model.safetensors.index.json

@@ -0,0 +1,763 @@
+{
+  "metadata": {
+    "total_size": 32174289964
+  },
+  "weight_map": {
+    "biobrain_decoder.gpt_model.final_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.attn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.ffn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.self_attn.key_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.self_attn.out_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.self_attn.query_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.0.self_attn.value_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.attn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.ffn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.self_attn.key_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.self_attn.out_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.self_attn.query_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.1.self_attn.value_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.10.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.11.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.12.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.13.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.14.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.self_attn.key_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.self_attn.query_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.15.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.16.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.17.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.18.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.19.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.attn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.ffn_norm.scale": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.self_attn.key_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.self_attn.out_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.self_attn.query_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.2.self_attn.value_linear.weight": "model-00001-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.fc1.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.fc2.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.20.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.attn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.ffn_norm.scale": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.self_attn.key_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.self_attn.out_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.self_attn.query_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.21.self_attn.value_linear.weight": "model-00004-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.22.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.23.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.24.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.25.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.fc1.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.fc2.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.26.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.attn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.fc1.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.fc2.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.ffn_norm.scale": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.self_attn.key_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.self_attn.out_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.self_attn.query_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.27.self_attn.value_linear.weight": "model-00005-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.attn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.fc1.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.fc2.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.ffn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.self_attn.key_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.self_attn.out_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.self_attn.query_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.28.self_attn.value_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.attn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.fc1.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.fc2.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.ffn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.self_attn.key_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.self_attn.out_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.self_attn.query_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.29.self_attn.value_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.3.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.attn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.fc1.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.fc2.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.ffn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.self_attn.key_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.self_attn.out_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.self_attn.query_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.30.self_attn.value_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.attn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.fc1.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.fc2.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.ffn_norm.scale": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.self_attn.key_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.self_attn.out_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.self_attn.query_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.31.self_attn.value_linear.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.4.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.5.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.6.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.7.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.attn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.fc1.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.fc2.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.ffn_norm.scale": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.self_attn.out_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.8.self_attn.value_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.attn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.fc1.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.fc2.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.ffn_norm.scale": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.self_attn.key_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.self_attn.out_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.self_attn.query_linear.weight": "model-00002-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.layers.9.self_attn.value_linear.weight": "model-00003-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.lm_head.fc.weight": "model-00006-of-00007.safetensors",
+    "biobrain_decoder.gpt_model.token_embed.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.0.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.1.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.10.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.11.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.12.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.13.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.14.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.15.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.16.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.17.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.18.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.19.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.2.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.20.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.21.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.22.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.23.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.24.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.25.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.26.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.27.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.28.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.3.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.4.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.5.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.6.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.7.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.8.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.fc2.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.layer_norm_mlp.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.layer_norm_mlp.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.layer_norm_self_attention.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.layer_norm_self_attention.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.output.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.output.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_k.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_k.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_q.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_q.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_v.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.attention_blocks.9.mha.w_v.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.embed_layer.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._fc1.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._fc1.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._final_fc.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._final_fc.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._first_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._first_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._second_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "biobrain_encoder.esm_model.lm_head._second_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "projection_model.bio_projection.bias": "model-00006-of-00007.safetensors",
+    "projection_model.bio_projection.weight": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.latent_queries": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_k.bias": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_k.weight": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_q.bias": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_q.weight": "model-00006-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_1.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_k.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_k.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_q.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_q.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.cross_attention_2.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.fc1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.fc1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.fc2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.fc2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_cross_attention_1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_cross_attention_1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_cross_attention_2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_cross_attention_2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_mlp.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.0.norm_mlp.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_k.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_k.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_q.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_q.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_1.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_k.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_k.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_q.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_q.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.cross_attention_2.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.fc1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.fc1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.fc2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.fc2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_cross_attention_1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_cross_attention_1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_cross_attention_2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_cross_attention_2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_mlp.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.1.norm_mlp.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_k.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_k.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_q.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_q.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_1.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.output.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.output.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_k.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_k.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_q.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_q.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_v.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.cross_attention_2.w_v.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.fc1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.fc1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.fc2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.fc2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_cross_attention_1.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_cross_attention_1.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_cross_attention_2.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_cross_attention_2.weight": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_mlp.bias": "model-00007-of-00007.safetensors",
+    "projection_model.perceiver_resampler.layers.2.norm_mlp.weight": "model-00007-of-00007.safetensors",
+    "projection_model.token_embedding.weight": "model-00006-of-00007.safetensors"
+  }
+}

multi_omics_model.py

@@ -0,0 +1,127 @@
+import torch
+from transformers import PreTrainedModel
+
+from genomics_research.biobrain_p1.porting_to_pytorch.configs.chatNT_config import (
+    ChatNTConfig,
+)
+from genomics_research.biobrain_p1.porting_to_pytorch.models.biobrain_decoder import (
+    TorchBioBrainDecoder,
+)
+from genomics_research.biobrain_p1.porting_to_pytorch.models.biobrain_encoder import (
+    TorchBioBrainEncoder,
+)
+from genomics_research.biobrain_p1.porting_to_pytorch.models.perceiver_resampler_projection import (  # noqa
+    TorchMultiModalPerceiverResamplerProjection,
+)
+
+
+class TorchMultiOmicsModel(PreTrainedModel):
+    config_class = ChatNTConfig
+
+    def __init__(self, config: ChatNTConfig) -> None:
+        super().__init__(config=config)
+        self.gpt_config = config.gpt_config
+        self.esm_config = config.esm_config
+        self.perceiver_resampler_config = config.perceiver_resampler_config
+        self.seq_token_id = config.seq_token_id
+        self.bio_pad_token_id = config.bio_pad_token_id
+        self.english_pad_token_id = config.english_pad_token_id
+
+        # Correct seq_token_id
+        self.seq_token_id -= 1
+
+        self.biobrain_encoder = TorchBioBrainEncoder(esm_config=self.esm_config)
+        self.biobrain_decoder = TorchBioBrainDecoder(
+            gpt_config=self.gpt_config, seq_token_id=self.seq_token_id
+        )
+        self.projection_model = TorchMultiModalPerceiverResamplerProjection(
+            perceiver_resampler_config=self.perceiver_resampler_config,
+            input_embed_dim=self.esm_config.embed_dim,
+            embed_dim=self.gpt_config.embed_dim,
+            english_vocab_size=self.gpt_config.vocab_size,
+            bio_pad_token_id=self.bio_pad_token_id,
+            english_pad_token_id=self.english_pad_token_id,
+        )
+
+    def forward(
+        self,
+        multi_omics_tokens_ids: tuple[torch.Tensor, torch.Tensor],
+        projection_english_tokens_ids: torch.Tensor,
+        projected_bio_embeddings: torch.Tensor = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+
+        Args:
+            multi_omics_tokens_ids (Tuple[torch.Tensor, torch.Tensor]):
+                english_tokens_ids: Represents the prompt tokens (english tokens)
+                    Shape (batch_size, num_english_tokens)
+
+                bio_tokens_ids: Represents the bio sequences tokens
+                    Shape (batch_size, num_bio_sequences, num_bio_tokens)
+
+            projection_english_tokens_ids (torch.Tensor):
+                Shape (batch_size, num_english_tokens)
+
+            projected_bio_embeddings (projected_bio_embeddings, optional):
+                Shape (batch_size, num_bio_sequencse, ?, embed_dim).
+                Defaults to None.
+
+        Returns:
+            dict[str, torch.Tensor] containing:
+                - logits:
+                    Shape (batch_size, num_tokens, vocab_size)
+
+                - projected_bio_embeddings:
+                    Shape (batch_size, num_bio_sequences, ?, embed_dim)
+        """
+        english_token_ids, bio_token_ids = multi_omics_tokens_ids
+
+        # Replace config.vocab_size value in english tokens
+        # We do this because the default vocab size (32000) doesn't match with the
+        # number of tokens because of seq_token_id(=32000) that was added
+        # Therefore, we will put seq_token_id to 31999
+        # (I will also put token n°31999 to 0, which is for unknown token)
+        # This is a workaround to avoid having to change the vocab size in the config
+        vocab_size = self.gpt_config.vocab_size
+        # Replace vocab
+        english_token_ids[english_token_ids == vocab_size - 1] = 0
+        projection_english_tokens_ids[
+            projection_english_tokens_ids == vocab_size - 1
+        ] = 0
+        english_token_ids[english_token_ids == vocab_size] = vocab_size - 1
+        projection_english_tokens_ids[projection_english_tokens_ids == vocab_size] = (
+            vocab_size - 1
+        )
+
+        if bio_token_ids is None:
+            projected_bio_embeddings = None
+        else:
+            num_bio_sequences = bio_token_ids.shape[1]
+
+            if projected_bio_embeddings is None:
+                # Compute bio sequences embeddings
+                bio_embeddings_list = [
+                    self.biobrain_encoder(bio_token_ids=bio_token_ids[:, bio_seq_num])
+                    for bio_seq_num in range(num_bio_sequences)
+                ]
+
+                # Project these embeddings
+                projected_bio_embeddings = [
+                    self.projection_model(
+                        bio_token_ids=bio_token_ids[:, bio_seq_num],
+                        bio_embeddings=bio_embeddings,
+                        english_token_ids=projection_english_tokens_ids,
+                    )
+                    for bio_seq_num, bio_embeddings in enumerate(bio_embeddings_list)
+                ]
+                projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
+
+        # decode
+        logits = self.biobrain_decoder(
+            english_token_ids=english_token_ids,
+            projected_bio_embeddings=projected_bio_embeddings,
+        )
+
+        outs = {"logits": logits, "projected_bio_embeddings": projected_bio_embeddings}
+
+        return outs