initial commit

adapter_config.json

@@ -0,0 +1,254 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "NLLBLLM2Vec",
+    "parent_library": "modeling_nllbllm2vec"
+    },
+  "base_model_name_or_path": "fdschmidt93/NLLBLLM2vec",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "task_type": "FEATURE_EXTRACTION",
+  "target_modules": [
+       "llm2vec.layers.0.self_attn.q_proj",
+       "llm2vec.layers.0.self_attn.k_proj",
+       "llm2vec.layers.0.self_attn.v_proj",
+       "llm2vec.layers.0.self_attn.o_proj",
+       "llm2vec.layers.0.mlp.gate_proj",
+       "llm2vec.layers.0.mlp.up_proj",
+       "llm2vec.layers.0.mlp.down_proj",
+       "llm2vec.layers.1.self_attn.q_proj",
+       "llm2vec.layers.1.self_attn.k_proj",
+       "llm2vec.layers.1.self_attn.v_proj",
+       "llm2vec.layers.1.self_attn.o_proj",
+       "llm2vec.layers.1.mlp.gate_proj",
+       "llm2vec.layers.1.mlp.up_proj",
+       "llm2vec.layers.1.mlp.down_proj",
+       "llm2vec.layers.2.self_attn.q_proj",
+       "llm2vec.layers.2.self_attn.k_proj",
+       "llm2vec.layers.2.self_attn.v_proj",
+       "llm2vec.layers.2.self_attn.o_proj",
+       "llm2vec.layers.2.mlp.gate_proj",
+       "llm2vec.layers.2.mlp.up_proj",
+       "llm2vec.layers.2.mlp.down_proj",
+       "llm2vec.layers.3.self_attn.q_proj",
+       "llm2vec.layers.3.self_attn.k_proj",
+       "llm2vec.layers.3.self_attn.v_proj",
+       "llm2vec.layers.3.self_attn.o_proj",
+       "llm2vec.layers.3.mlp.gate_proj",
+       "llm2vec.layers.3.mlp.up_proj",
+       "llm2vec.layers.3.mlp.down_proj",
+       "llm2vec.layers.4.self_attn.q_proj",
+       "llm2vec.layers.4.self_attn.k_proj",
+       "llm2vec.layers.4.self_attn.v_proj",
+       "llm2vec.layers.4.self_attn.o_proj",
+       "llm2vec.layers.4.mlp.gate_proj",
+       "llm2vec.layers.4.mlp.up_proj",
+       "llm2vec.layers.4.mlp.down_proj",
+       "llm2vec.layers.5.self_attn.q_proj",
+       "llm2vec.layers.5.self_attn.k_proj",
+       "llm2vec.layers.5.self_attn.v_proj",
+       "llm2vec.layers.5.self_attn.o_proj",
+       "llm2vec.layers.5.mlp.gate_proj",
+       "llm2vec.layers.5.mlp.up_proj",
+       "llm2vec.layers.5.mlp.down_proj",
+       "llm2vec.layers.6.self_attn.q_proj",
+       "llm2vec.layers.6.self_attn.k_proj",
+       "llm2vec.layers.6.self_attn.v_proj",
+       "llm2vec.layers.6.self_attn.o_proj",
+       "llm2vec.layers.6.mlp.gate_proj",
+       "llm2vec.layers.6.mlp.up_proj",
+       "llm2vec.layers.6.mlp.down_proj",
+       "llm2vec.layers.7.self_attn.q_proj",
+       "llm2vec.layers.7.self_attn.k_proj",
+       "llm2vec.layers.7.self_attn.v_proj",
+       "llm2vec.layers.7.self_attn.o_proj",
+       "llm2vec.layers.7.mlp.gate_proj",
+       "llm2vec.layers.7.mlp.up_proj",
+       "llm2vec.layers.7.mlp.down_proj",
+       "llm2vec.layers.8.self_attn.q_proj",
+       "llm2vec.layers.8.self_attn.k_proj",
+       "llm2vec.layers.8.self_attn.v_proj",
+       "llm2vec.layers.8.self_attn.o_proj",
+       "llm2vec.layers.8.mlp.gate_proj",
+       "llm2vec.layers.8.mlp.up_proj",
+       "llm2vec.layers.8.mlp.down_proj",
+       "llm2vec.layers.9.self_attn.q_proj",
+       "llm2vec.layers.9.self_attn.k_proj",
+       "llm2vec.layers.9.self_attn.v_proj",
+       "llm2vec.layers.9.self_attn.o_proj",
+       "llm2vec.layers.9.mlp.gate_proj",
+       "llm2vec.layers.9.mlp.up_proj",
+       "llm2vec.layers.9.mlp.down_proj",
+       "llm2vec.layers.10.self_attn.q_proj",
+       "llm2vec.layers.10.self_attn.k_proj",
+       "llm2vec.layers.10.self_attn.v_proj",
+       "llm2vec.layers.10.self_attn.o_proj",
+       "llm2vec.layers.10.mlp.gate_proj",
+       "llm2vec.layers.10.mlp.up_proj",
+       "llm2vec.layers.10.mlp.down_proj",
+       "llm2vec.layers.11.self_attn.q_proj",
+       "llm2vec.layers.11.self_attn.k_proj",
+       "llm2vec.layers.11.self_attn.v_proj",
+       "llm2vec.layers.11.self_attn.o_proj",
+       "llm2vec.layers.11.mlp.gate_proj",
+       "llm2vec.layers.11.mlp.up_proj",
+       "llm2vec.layers.11.mlp.down_proj",
+       "llm2vec.layers.12.self_attn.q_proj",
+       "llm2vec.layers.12.self_attn.k_proj",
+       "llm2vec.layers.12.self_attn.v_proj",
+       "llm2vec.layers.12.self_attn.o_proj",
+       "llm2vec.layers.12.mlp.gate_proj",
+       "llm2vec.layers.12.mlp.up_proj",
+       "llm2vec.layers.12.mlp.down_proj",
+       "llm2vec.layers.13.self_attn.q_proj",
+       "llm2vec.layers.13.self_attn.k_proj",
+       "llm2vec.layers.13.self_attn.v_proj",
+       "llm2vec.layers.13.self_attn.o_proj",
+       "llm2vec.layers.13.mlp.gate_proj",
+       "llm2vec.layers.13.mlp.up_proj",
+       "llm2vec.layers.13.mlp.down_proj",
+       "llm2vec.layers.14.self_attn.q_proj",
+       "llm2vec.layers.14.self_attn.k_proj",
+       "llm2vec.layers.14.self_attn.v_proj",
+       "llm2vec.layers.14.self_attn.o_proj",
+       "llm2vec.layers.14.mlp.gate_proj",
+       "llm2vec.layers.14.mlp.up_proj",
+       "llm2vec.layers.14.mlp.down_proj",
+       "llm2vec.layers.15.self_attn.q_proj",
+       "llm2vec.layers.15.self_attn.k_proj",
+       "llm2vec.layers.15.self_attn.v_proj",
+       "llm2vec.layers.15.self_attn.o_proj",
+       "llm2vec.layers.15.mlp.gate_proj",
+       "llm2vec.layers.15.mlp.up_proj",
+       "llm2vec.layers.15.mlp.down_proj",
+       "llm2vec.layers.16.self_attn.q_proj",
+       "llm2vec.layers.16.self_attn.k_proj",
+       "llm2vec.layers.16.self_attn.v_proj",
+       "llm2vec.layers.16.self_attn.o_proj",
+       "llm2vec.layers.16.mlp.gate_proj",
+       "llm2vec.layers.16.mlp.up_proj",
+       "llm2vec.layers.16.mlp.down_proj",
+       "llm2vec.layers.17.self_attn.q_proj",
+       "llm2vec.layers.17.self_attn.k_proj",
+       "llm2vec.layers.17.self_attn.v_proj",
+       "llm2vec.layers.17.self_attn.o_proj",
+       "llm2vec.layers.17.mlp.gate_proj",
+       "llm2vec.layers.17.mlp.up_proj",
+       "llm2vec.layers.17.mlp.down_proj",
+       "llm2vec.layers.18.self_attn.q_proj",
+       "llm2vec.layers.18.self_attn.k_proj",
+       "llm2vec.layers.18.self_attn.v_proj",
+       "llm2vec.layers.18.self_attn.o_proj",
+       "llm2vec.layers.18.mlp.gate_proj",
+       "llm2vec.layers.18.mlp.up_proj",
+       "llm2vec.layers.18.mlp.down_proj",
+       "llm2vec.layers.19.self_attn.q_proj",
+       "llm2vec.layers.19.self_attn.k_proj",
+       "llm2vec.layers.19.self_attn.v_proj",
+       "llm2vec.layers.19.self_attn.o_proj",
+       "llm2vec.layers.19.mlp.gate_proj",
+       "llm2vec.layers.19.mlp.up_proj",
+       "llm2vec.layers.19.mlp.down_proj",
+       "llm2vec.layers.20.self_attn.q_proj",
+       "llm2vec.layers.20.self_attn.k_proj",
+       "llm2vec.layers.20.self_attn.v_proj",
+       "llm2vec.layers.20.self_attn.o_proj",
+       "llm2vec.layers.20.mlp.gate_proj",
+       "llm2vec.layers.20.mlp.up_proj",
+       "llm2vec.layers.20.mlp.down_proj",
+       "llm2vec.layers.21.self_attn.q_proj",
+       "llm2vec.layers.21.self_attn.k_proj",
+       "llm2vec.layers.21.self_attn.v_proj",
+       "llm2vec.layers.21.self_attn.o_proj",
+       "llm2vec.layers.21.mlp.gate_proj",
+       "llm2vec.layers.21.mlp.up_proj",
+       "llm2vec.layers.21.mlp.down_proj",
+       "llm2vec.layers.22.self_attn.q_proj",
+       "llm2vec.layers.22.self_attn.k_proj",
+       "llm2vec.layers.22.self_attn.v_proj",
+       "llm2vec.layers.22.self_attn.o_proj",
+       "llm2vec.layers.22.mlp.gate_proj",
+       "llm2vec.layers.22.mlp.up_proj",
+       "llm2vec.layers.22.mlp.down_proj",
+       "llm2vec.layers.23.self_attn.q_proj",
+       "llm2vec.layers.23.self_attn.k_proj",
+       "llm2vec.layers.23.self_attn.v_proj",
+       "llm2vec.layers.23.self_attn.o_proj",
+       "llm2vec.layers.23.mlp.gate_proj",
+       "llm2vec.layers.23.mlp.up_proj",
+       "llm2vec.layers.23.mlp.down_proj",
+       "llm2vec.layers.24.self_attn.q_proj",
+       "llm2vec.layers.24.self_attn.k_proj",
+       "llm2vec.layers.24.self_attn.v_proj",
+       "llm2vec.layers.24.self_attn.o_proj",
+       "llm2vec.layers.24.mlp.gate_proj",
+       "llm2vec.layers.24.mlp.up_proj",
+       "llm2vec.layers.24.mlp.down_proj",
+       "llm2vec.layers.25.self_attn.q_proj",
+       "llm2vec.layers.25.self_attn.k_proj",
+       "llm2vec.layers.25.self_attn.v_proj",
+       "llm2vec.layers.25.self_attn.o_proj",
+       "llm2vec.layers.25.mlp.gate_proj",
+       "llm2vec.layers.25.mlp.up_proj",
+       "llm2vec.layers.25.mlp.down_proj",
+       "llm2vec.layers.26.self_attn.q_proj",
+       "llm2vec.layers.26.self_attn.k_proj",
+       "llm2vec.layers.26.self_attn.v_proj",
+       "llm2vec.layers.26.self_attn.o_proj",
+       "llm2vec.layers.26.mlp.gate_proj",
+       "llm2vec.layers.26.mlp.up_proj",
+       "llm2vec.layers.26.mlp.down_proj",
+       "llm2vec.layers.27.self_attn.q_proj",
+       "llm2vec.layers.27.self_attn.k_proj",
+       "llm2vec.layers.27.self_attn.v_proj",
+       "llm2vec.layers.27.self_attn.o_proj",
+       "llm2vec.layers.27.mlp.gate_proj",
+       "llm2vec.layers.27.mlp.up_proj",
+       "llm2vec.layers.27.mlp.down_proj",
+       "llm2vec.layers.28.self_attn.q_proj",
+       "llm2vec.layers.28.self_attn.k_proj",
+       "llm2vec.layers.28.self_attn.v_proj",
+       "llm2vec.layers.28.self_attn.o_proj",
+       "llm2vec.layers.28.mlp.gate_proj",
+       "llm2vec.layers.28.mlp.up_proj",
+       "llm2vec.layers.28.mlp.down_proj",
+       "llm2vec.layers.29.self_attn.q_proj",
+       "llm2vec.layers.29.self_attn.k_proj",
+       "llm2vec.layers.29.self_attn.v_proj",
+       "llm2vec.layers.29.self_attn.o_proj",
+       "llm2vec.layers.29.mlp.gate_proj",
+       "llm2vec.layers.29.mlp.up_proj",
+       "llm2vec.layers.29.mlp.down_proj",
+       "llm2vec.layers.30.self_attn.q_proj",
+       "llm2vec.layers.30.self_attn.k_proj",
+       "llm2vec.layers.30.self_attn.v_proj",
+       "llm2vec.layers.30.self_attn.o_proj",
+       "llm2vec.layers.30.mlp.gate_proj",
+       "llm2vec.layers.30.mlp.up_proj",
+       "llm2vec.layers.30.mlp.down_proj",
+       "llm2vec.layers.31.self_attn.q_proj",
+       "llm2vec.layers.31.self_attn.k_proj",
+       "llm2vec.layers.31.self_attn.v_proj",
+       "llm2vec.layers.31.self_attn.o_proj",
+       "llm2vec.layers.31.mlp.gate_proj",
+       "llm2vec.layers.31.mlp.up_proj",
+       "llm2vec.layers.31.mlp.down_proj"
+  ],
+  "use_dora": false,
+  "use_rslora": false
+}

adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fd0c6b2576840df0e26ccac0f6eb18e27945cefcd33668a3720cb20fec94b66
+size 167833136

config.json

@@ -0,0 +1,41 @@
+{
+  "_name_or_path": "fdschmidt93/NLLBLLM2Vec",
+  "architectures": [
+    "NLLBLLM2Vec"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_nllbllm2vec.NLLBLLM2VecConfig",
+    "AutoModel": "modeling_nllbllm2vec.NLLBLLM2Vec"
+  },
+  "llm2vec_config": {
+    "_name_or_path": "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp",
+    "bos_token_id": 128000,
+    "eos_token_id": 128001,
+    "intermediate_size": 14336,
+    "max_position_embeddings": 8192,
+    "model_type": "llama",
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 500000,
+    "torch_dtype": "bfloat16",
+    "use_cache": false,
+    "vocab_size": 128256
+  },
+  "model_type": "nllb-llm2vec",
+  "nllb_config": {
+    "_name_or_path": "facebook/nllb-200-distilled-600M",
+    "architectures": [
+      "M2M100Encoder"
+    ],
+    "decoder_layerdrop": 0,
+    "encoder_layerdrop": 0,
+    "max_length": 200,
+    "model_type": "m2m_100",
+    "tokenizer_class": "NllbTokenizer",
+    "torch_dtype": "bfloat16",
+    "vocab_size": 256206
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.2"
+}
+

configuration_nllbllm2vec.py

@@ -0,0 +1,84 @@
+from transformers import AutoConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.m2m_100.configuration_m2m_100 import M2M100Config
+
+NLLBLLM2VEC_TYPE = "nllb-llm2vec"
+
+DEFAULT_M2M100_CONFIG = {
+    "activation_dropout": 0.0,
+    "activation_function": "relu",
+    "architectures": ["M2M100Encoder"],
+    "attention_dropout": 0.1,
+    "bos_token_id": 0,
+    "d_model": 1024,
+    "decoder_attention_heads": 16,
+    "decoder_ffn_dim": 4096,
+    "decoder_layerdrop": 0,
+    "decoder_layers": 12,
+    "decoder_start_token_id": 2,
+    "dropout": 0.1,
+    "encoder_attention_heads": 16,
+    "encoder_ffn_dim": 4096,
+    "encoder_layerdrop": 0,
+    "encoder_layers": 12,
+    "eos_token_id": 2,
+    "init_std": 0.02,
+    "is_encoder_decoder": True,
+    "max_position_embeddings": 1024,
+    "model_type": "m2m_100",
+    "num_hidden_layers": 12,
+    "pad_token_id": 1,
+    "scale_embedding": True,
+    "torch_dtype": "float32",
+    "transformers_version": "4.21.0.dev0",
+    "use_cache": True,
+    "vocab_size": 256206,
+    "tokenizer_class": "NllbTokenizer",
+    "max_length": 200,
+}
+
+DEFAULT_LLAMA_CONFIG = {
+    "attention_bias": False,
+    "attention_dropout": 0,
+    "bos_token_id": 128000,
+    "eos_token_id": 128001,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 14336,
+    "max_position_embeddings": 8192,
+    "model_type": "llama",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 8,
+    "pretraining_tp": 1,
+    "rms_norm_eps": 0.00001,
+    "rope_scaling": None,
+    "rope_theta": 500000,
+    "tie_word_embeddings": False,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.40.0.dev0",
+    "use_cache": False,
+    "vocab_size": 128256,
+}
+
+
+class NLLBLLM2VecConfig(PretrainedConfig):
+    model_type = "nllb-llm2vec"
+    is_composition = False
+
+    def __init__(
+        self,
+        nllb_config: dict = DEFAULT_M2M100_CONFIG,
+        llm2vec_config: dict = DEFAULT_LLAMA_CONFIG,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.nllb_config = M2M100Config(**nllb_config)
+        self.llm2vec_config = LlamaConfig(**llm2vec_config)
+
+
+AutoConfig.register(NLLBLLM2VEC_TYPE, NLLBLLM2VecConfig)
+
+NLLBLLM2VecConfig.register_for_auto_class()

model-00001-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2409158d0692d1b6d8ade7342ddf40bbaeda0ed6bb545624d679bd3549d4d894
+size 4795637544

model-00002-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:302a8d954626008cb14b1f74ae39d399ade67fe2d3404a7788ed5cb792b68cfa
+size 4832007544

model-00003-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a77f2fa237b592921ffecb745ac613764da66643fbe98145435532cc1e63707a
+size 4999813200

model-00004-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a22835800d55155cfc48b8ec2f06a173d203eec52af80c0923c5ad057ee977e9
+size 4999813232

model-00005-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1482c1048034b6bbb552b0d6a3969f5e76149e8dd4a7cc45ddb6cfae4d524a92
+size 4832007592

model-00006-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6a5b0755ef742857b9f785cd6b61c9f3f16677581208d105562f53d6094647b
+size 4999813232

model-00007-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2b0493c68cadb801fe6b5c0ff4d6c9d2093eca333c98736665c36ab7c523900
+size 2231487200

model.safetensors.index.json

@@ -0,0 +1,493 @@
+{
+  "metadata": {
+    "total_size": 31690522624
+  },
+  "weight_map": {
+    "llm2vec.embed_tokens.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.input_layernorm.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.mlp.down_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.mlp.gate_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.mlp.up_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.self_attn.o_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.0.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.1.input_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.1.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.1.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.1.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.1.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.1.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.1.self_attn.o_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.1.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.1.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "llm2vec.layers.10.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.10.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.11.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.12.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.12.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.12.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.12.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.12.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.12.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.12.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.12.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.12.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.13.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.13.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.14.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.15.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.16.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.input_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.17.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.18.input_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.18.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.18.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.18.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.18.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.18.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.18.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.18.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.18.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
+    "llm2vec.layers.19.input_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.19.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.2.input_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.2.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.20.input_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.20.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.input_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.21.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.input_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.22.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.23.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.23.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.23.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.23.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
+    "llm2vec.layers.24.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.24.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.25.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.26.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.27.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.input_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.28.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.29.input_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.29.mlp.down_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.29.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.29.mlp.up_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.29.post_attention_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.29.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.29.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.29.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.29.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
+    "llm2vec.layers.3.input_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.3.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.30.input_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.mlp.down_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.mlp.gate_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.mlp.up_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.post_attention_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.self_attn.k_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.self_attn.o_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.self_attn.q_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.30.self_attn.v_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.input_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.mlp.down_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.mlp.gate_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.mlp.up_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.post_attention_layernorm.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.self_attn.k_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.self_attn.o_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.self_attn.q_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.31.self_attn.v_proj.weight": "model-00007-of-00007.safetensors",
+    "llm2vec.layers.4.input_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.4.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.input_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.5.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.6.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.6.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.6.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.6.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
+    "llm2vec.layers.7.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.7.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.8.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.input_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.layers.9.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
+    "llm2vec.norm.weight": "model-00007-of-00007.safetensors",
+    "nllb_encoder.embed_tokens.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.0.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.1.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.10.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.11.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.2.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.3.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.4.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.5.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.6.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.7.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.8.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.fc1.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.fc1.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.fc2.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.fc2.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.final_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.final_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.out_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.out_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn_layer_norm.bias": "model-00001-of-00007.safetensors",
+    "nllb_encoder.layers.9.self_attn_layer_norm.weight": "model-00001-of-00007.safetensors",
+    "up_proj.weight": "model-00007-of-00007.safetensors"
+  }
+}

modeling_llama_encoder.py

@@ -0,0 +1,195 @@
+import importlib.metadata
+
+import torch
+from packaging import version
+from torch import nn
+from transformers import LlamaConfig, LlamaModel, LlamaPreTrainedModel
+from transformers.cache_utils import Cache, StaticCache
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaDecoderLayer,
+    LlamaFlashAttention2,
+    LlamaMLP,
+    LlamaRMSNorm,
+    LlamaRotaryEmbedding,
+    LlamaSdpaAttention,
+)
+from transformers.utils import logging
+from transformers.utils.import_utils import _is_package_available
+
+logger = logging.get_logger(__name__)
+
+
+def is_transformers_attn_greater_or_equal_4_43_1():
+    if not _is_package_available("transformers"):
+        return False
+
+    return version.parse(importlib.metadata.version("transformers")) >= version.parse(
+        "4.43.1"
+    )
+
+
+class ModifiedLlamaAttention(LlamaAttention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False
+
+
+class ModifiedLlamaFlashAttention2(LlamaFlashAttention2):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False
+
+
+class ModifiedLlamaSdpaAttention(LlamaSdpaAttention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False
+
+
+LLAMA_ATTENTION_CLASSES = {
+    "eager": ModifiedLlamaAttention,
+    "flash_attention_2": ModifiedLlamaFlashAttention2,
+    "sdpa": ModifiedLlamaSdpaAttention,
+}
+
+
+class ModifiedLlamaDecoderLayer(LlamaDecoderLayer):
+    def __init__(self, config: LlamaConfig, layer_idx: int):
+        nn.Module.__init__(self)
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](
+            config=config, layer_idx=layer_idx
+        )
+
+        self.mlp = LlamaMLP(config)
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+
+class LlamaEncoderModel(LlamaModel):
+    _no_split_modules = ["ModifiedLlamaDecoderLayer"]
+
+    def __init__(self, config: LlamaConfig):
+        if not is_transformers_attn_greater_or_equal_4_43_1():
+            raise ValueError(
+                "The current implementation of LlamaEncoderModel follows modeling_llama.py of transformers version >= 4.43.1"
+            )
+        LlamaPreTrainedModel.__init__(self, config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                ModifiedLlamaDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def _update_causal_mask(
+        self,
+        attention_mask,
+        input_tensor,
+        cache_position,
+        past_key_values: Cache,
+        output_attentions: bool,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = (
+            past_key_values.get_seq_length() if past_key_values is not None else 0
+        )
+        using_static_cache = isinstance(past_key_values, StaticCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        # if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
+        #     if AttentionMaskConverter._ignore_causal_mask_sdpa(
+        #         attention_mask,
+        #         inputs_embeds=input_tensor,
+        #         past_key_values_length=past_seen_tokens,
+        #         is_training=self.training,
+        #     ):
+        #         return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        if using_static_cache:
+            target_length = past_key_values.get_max_length()
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        causal_mask = torch.zeros(
+            (sequence_length, target_length), dtype=dtype, device=device
+        )  # in original implementation - torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
+        # Commenting out next 2 lines to disable causal masking
+        # if sequence_length != 1:
+        #     causal_mask = torch.triu(causal_mask, diagonal=1)
+        causal_mask *= torch.arange(
+            target_length, device=device
+        ) > cache_position.reshape(-1, 1)
+        causal_mask = causal_mask[None, None, :, :].expand(
+            input_tensor.shape[0], 1, -1, -1
+        )
+        if attention_mask is not None:
+            causal_mask = (
+                causal_mask.clone()
+            )  # copy to contiguous memory for in-place edit
+            if attention_mask.dim() == 2:
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[
+                    :, None, None, :
+                ].eq(0.0)
+                causal_mask[..., :mask_length] = causal_mask[
+                    ..., :mask_length
+                ].masked_fill(padding_mask, min_dtype)
+            elif attention_mask.dim() == 4:
+                # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
+                # cache. In that case, the 4D attention mask attends to the newest tokens only.
+                if attention_mask.shape[-2] < cache_position[0] + sequence_length:
+                    offset = cache_position[0]
+                else:
+                    offset = 0
+                mask_shape = attention_mask.shape
+                mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
+                causal_mask[
+                    : mask_shape[0],
+                    : mask_shape[1],
+                    offset : mask_shape[2] + offset,
+                    : mask_shape[3],
+                ] = mask_slice
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type == "cuda"
+            and not output_attentions
+        ):
+            causal_mask = AttentionMaskConverter._unmask_unattended(
+                causal_mask, min_dtype
+            )
+
+        return causal_mask

modeling_nllbllm2vec.py

@@ -0,0 +1,481 @@
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.auto import AutoModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.m2m_100.modeling_m2m_100 import M2M100Encoder
+
+from .configuration_nllbllm2vec import NLLBLLM2VecConfig
+from .modeling_llama_encoder import LlamaEncoderModel
+
+
+class NLLBLLM2Vec(PreTrainedModel):
+    config_class = NLLBLLM2VecConfig
+    """
+    NLLBLLM2Vec model combining NLLB and LLama encoders.
+
+    Args:
+        config (Optional[NLLBLLM2VecConfig]): Configuration object.
+        nllb_encoder (Optional[M2M100Encoder]): Pre-initialized NLLB encoder.
+        llm2vec (Optional[LlamaEncoderModel]): Pre-initialized LLama encoder.
+        *inputs: Additional positional arguments.
+        **kwargs: Additional keyword arguments.
+    """
+
+    model_type = "nllb-llm2vec"
+
+    def __init__(
+        self,
+        config: Optional[NLLBLLM2VecConfig] = None,
+        nllb_encoder: Optional[M2M100Encoder] = None,
+        llm2vec: Optional[LlamaEncoderModel] = None,
+        *inputs,
+        **kwargs,
+    ):
+        # Ensure that either config is not None or both encoders are provided
+        if config is None and (nllb_encoder is None or llm2vec is None):
+            raise ValueError(
+                "Either `config` must be provided, or both `nllb_encoder` and `llm2vec` must be specified."
+            )
+
+        if config is not None:
+            super().__init__(config, *inputs, **kwargs)
+            self.nllb_encoder = nllb_encoder or M2M100Encoder(config.nllb_config)
+            self.llm2vec = llm2vec or LlamaEncoderModel(config.llm2vec_config)
+            self.config = config
+        else:
+            # Both encoders are provided
+            self.nllb_encoder = cast(M2M100Encoder, nllb_encoder)
+            self.llm2vec = cast(LlamaEncoderModel, llm2vec)
+            self.config = NLLBLLM2VecConfig(
+                nllb_config=self.nllb_encoder.config,  # type: ignore
+                llm2vec_config=self.llm2vec.config,  # type: ignore
+            )
+            super().__init__(self.config, *inputs, **kwargs)
+
+        self.up_proj = nn.Linear(
+            self.nllb_encoder.config.d_model,
+            self.llm2vec.config.hidden_size,
+            bias=False,
+        )
+        # Additional initialization logic can go here
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        indices: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        *args,
+        **kwargs,
+    ) -> BaseModelOutputWithPooling:
+        """
+        Forward pass of the model.
+
+        Args:
+            input_ids (torch.Tensor): Input token IDs.
+            attention_mask (torch.Tensor): Attention mask.
+            indices (Optional[Tuple[torch.Tensor, torch.Tensor]]): Precomputed input indices and offsets.
+
+        Returns:
+            BaseModelOutputWithPooling: Model outputs with last hidden state and pooled output.
+        """
+        # Compute input indices and offsets if not provided
+        if indices is None:
+            seq_indices, seq_offsets = self._get_input_offsets(attention_mask)
+        else:
+            seq_indices, seq_offsets = indices
+
+        with torch.inference_mode():
+            nllb_outputs = self.nllb_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            nllb_last_hidden_state = nllb_outputs.last_hidden_state
+            nllb_last_hidden_state = self.up_proj(nllb_last_hidden_state)
+        if self.training:
+            # Inference mode otherwise yields embeddings that cannot be trained with
+            nllb_last_hidden_state = nllb_last_hidden_state.detach().clone()
+        outputs = self.llm2vec(
+            inputs_embeds=nllb_last_hidden_state,
+            attention_mask=attention_mask,
+        )
+        pooler_output = self._mean_embedding(
+            hidden_states=outputs.last_hidden_state,
+            input_indices=seq_indices,
+            offsets=seq_offsets,
+        )
+        return BaseModelOutputWithPooling(
+            last_hidden_state=outputs.last_hidden_state,
+            pooler_output=pooler_output,
+        )
+
+    @property
+    def tokenizer(self):
+        """
+        Get the tokenizer associated with the model.
+
+        Returns:
+            PreTrainedTokenizer: The tokenizer instance.
+        """
+        if not hasattr(self, "_tokenizer"):
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/nllb-200-distilled-600M", padding_side="right"
+            )
+        return self._tokenizer
+
+    def encode(
+        self,
+        inputs: List[str],
+        src_lang: str = "eng_Latn",
+        tokenize_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> torch.Tensor:
+        """
+        Encode input texts into embeddings.
+
+        Args:
+            inputs (List[str]): List of input texts.
+            src_lang (str): Source language code.
+            tokenize_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments for the tokenizer.
+                Defaults to:
+                >>    tokenize_kwargs = {
+                >>        "padding": True,
+                >>        "truncation": True,
+                >>        "max_length": 512,
+                >>        "return_tensors": "pt",
+                >>    }
+
+        Returns:
+            torch.Tensor: Mean-pooled sequence embeddings of the inputs.
+        """
+        if tokenize_kwargs is None:
+            tokenize_kwargs = {
+                "padding": True,
+                "truncation": True,
+                "max_length": 512,
+                "return_tensors": "pt",
+            }
+
+        tokenizer = self.tokenizer
+        tokenizer.src_lang = src_lang
+        device = next(self.parameters()).device
+        batch = tokenizer(inputs, **tokenize_kwargs).to(device)
+        device_type = device.type  # e.g., 'cuda' or 'cpu'
+
+        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+            return self(**batch).pooler_output
+
+    @staticmethod
+    def _get_input_offsets(
+        attention_mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute indices and offsets for mean pooling using EmbeddingBag.
+
+        Args:
+            attention_mask (torch.Tensor): Attention mask of shape (batch_size, seq_len).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+                - input_indices: Indices of non-padded tokens in the flattened input.
+                - offsets: Offsets indicating the start index of each sequence in the flattened input.
+        """
+        # Find the indices of non-padded tokens in flattened hidden_states
+        input_indices = attention_mask.view(-1).nonzero(as_tuple=False).squeeze()
+
+        # Compute the offsets: for each sequence, where it starts in the flattened input
+        non_padded_lengths = attention_mask.sum(
+            dim=1
+        )  # Count non-padded tokens per sequence
+        offsets = torch.cat(
+            [
+                torch.tensor([0], device=attention_mask.device),
+                non_padded_lengths.cumsum(dim=0)[:-1],
+            ]
+        )
+        return input_indices, offsets
+
+    @staticmethod
+    def _mean_embedding(
+        hidden_states: torch.Tensor,
+        input_indices: torch.Tensor,
+        offsets: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Compute the mean of non-padded embeddings using `embedding_bag`,
+        properly handling padding with offsets.
+
+        Args:
+            hidden_states (torch.Tensor): Hidden states of shape (batch_size, seq_len, embed_dim).
+            input_indices (torch.Tensor): Indices of non-padded tokens in flattened form.
+            offsets (torch.Tensor): Offsets specifying the start of each sequence.
+
+        Returns:
+            torch.Tensor: Pooled mean embeddings of shape (batch_size, embed_dim).
+        """
+        # Flatten hidden_states to 2D: shape (batch_size * seq_len, embedding_dim)
+        batch_size, seq_len, embed_dim = hidden_states.shape
+        token_embeds = hidden_states.view(-1, embed_dim)
+
+        # Use embedding_bag with mode 'mean' and appropriate indices
+        return F.embedding_bag(
+            input=input_indices,  # Indices of non-padded tokens in flattened form
+            weight=token_embeds,  # The flattened hidden states as embedding matrix
+            offsets=offsets,  # Offsets specifying start of each sequence
+            mode="mean",  # Aggregation mode
+        )
+
+
+AutoModel.register(NLLBLLM2VecConfig, NLLBLLM2Vec)
+
+
+def repl():
+    cfg = NLLBLLM2VecConfig()
+    model = NLLBLLM2Vec(cfg)
+
+    from peft.mapping import get_peft_model
+    from peft.tuners.lora.config import LoraConfig
+
+    lora_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.0,
+        bias="none",
+        task_type="FEATURE_EXTRACTION",
+        target_modules=[
+            "llm2vec.layers.0.self_attn.q_proj",
+            "llm2vec.layers.0.self_attn.k_proj",
+            "llm2vec.layers.0.self_attn.v_proj",
+            "llm2vec.layers.0.self_attn.o_proj",
+            "llm2vec.layers.0.mlp.gate_proj",
+            "llm2vec.layers.0.mlp.up_proj",
+            "llm2vec.layers.0.mlp.down_proj",
+            "llm2vec.layers.1.self_attn.q_proj",
+            "llm2vec.layers.1.self_attn.k_proj",
+            "llm2vec.layers.1.self_attn.v_proj",
+            "llm2vec.layers.1.self_attn.o_proj",
+            "llm2vec.layers.1.mlp.gate_proj",
+            "llm2vec.layers.1.mlp.up_proj",
+            "llm2vec.layers.1.mlp.down_proj",
+            "llm2vec.layers.2.self_attn.q_proj",
+            "llm2vec.layers.2.self_attn.k_proj",
+            "llm2vec.layers.2.self_attn.v_proj",
+            "llm2vec.layers.2.self_attn.o_proj",
+            "llm2vec.layers.2.mlp.gate_proj",
+            "llm2vec.layers.2.mlp.up_proj",
+            "llm2vec.layers.2.mlp.down_proj",
+            "llm2vec.layers.3.self_attn.q_proj",
+            "llm2vec.layers.3.self_attn.k_proj",
+            "llm2vec.layers.3.self_attn.v_proj",
+            "llm2vec.layers.3.self_attn.o_proj",
+            "llm2vec.layers.3.mlp.gate_proj",
+            "llm2vec.layers.3.mlp.up_proj",
+            "llm2vec.layers.3.mlp.down_proj",
+            "llm2vec.layers.4.self_attn.q_proj",
+            "llm2vec.layers.4.self_attn.k_proj",
+            "llm2vec.layers.4.self_attn.v_proj",
+            "llm2vec.layers.4.self_attn.o_proj",
+            "llm2vec.layers.4.mlp.gate_proj",
+            "llm2vec.layers.4.mlp.up_proj",
+            "llm2vec.layers.4.mlp.down_proj",
+            "llm2vec.layers.5.self_attn.q_proj",
+            "llm2vec.layers.5.self_attn.k_proj",
+            "llm2vec.layers.5.self_attn.v_proj",
+            "llm2vec.layers.5.self_attn.o_proj",
+            "llm2vec.layers.5.mlp.gate_proj",
+            "llm2vec.layers.5.mlp.up_proj",
+            "llm2vec.layers.5.mlp.down_proj",
+            "llm2vec.layers.6.self_attn.q_proj",
+            "llm2vec.layers.6.self_attn.k_proj",
+            "llm2vec.layers.6.self_attn.v_proj",
+            "llm2vec.layers.6.self_attn.o_proj",
+            "llm2vec.layers.6.mlp.gate_proj",
+            "llm2vec.layers.6.mlp.up_proj",
+            "llm2vec.layers.6.mlp.down_proj",
+            "llm2vec.layers.7.self_attn.q_proj",
+            "llm2vec.layers.7.self_attn.k_proj",
+            "llm2vec.layers.7.self_attn.v_proj",
+            "llm2vec.layers.7.self_attn.o_proj",
+            "llm2vec.layers.7.mlp.gate_proj",
+            "llm2vec.layers.7.mlp.up_proj",
+            "llm2vec.layers.7.mlp.down_proj",
+            "llm2vec.layers.8.self_attn.q_proj",
+            "llm2vec.layers.8.self_attn.k_proj",
+            "llm2vec.layers.8.self_attn.v_proj",
+            "llm2vec.layers.8.self_attn.o_proj",
+            "llm2vec.layers.8.mlp.gate_proj",
+            "llm2vec.layers.8.mlp.up_proj",
+            "llm2vec.layers.8.mlp.down_proj",
+            "llm2vec.layers.9.self_attn.q_proj",
+            "llm2vec.layers.9.self_attn.k_proj",
+            "llm2vec.layers.9.self_attn.v_proj",
+            "llm2vec.layers.9.self_attn.o_proj",
+            "llm2vec.layers.9.mlp.gate_proj",
+            "llm2vec.layers.9.mlp.up_proj",
+            "llm2vec.layers.9.mlp.down_proj",
+            "llm2vec.layers.10.self_attn.q_proj",
+            "llm2vec.layers.10.self_attn.k_proj",
+            "llm2vec.layers.10.self_attn.v_proj",
+            "llm2vec.layers.10.self_attn.o_proj",
+            "llm2vec.layers.10.mlp.gate_proj",
+            "llm2vec.layers.10.mlp.up_proj",
+            "llm2vec.layers.10.mlp.down_proj",
+            "llm2vec.layers.11.self_attn.q_proj",
+            "llm2vec.layers.11.self_attn.k_proj",
+            "llm2vec.layers.11.self_attn.v_proj",
+            "llm2vec.layers.11.self_attn.o_proj",
+            "llm2vec.layers.11.mlp.gate_proj",
+            "llm2vec.layers.11.mlp.up_proj",
+            "llm2vec.layers.11.mlp.down_proj",
+            "llm2vec.layers.12.self_attn.q_proj",
+            "llm2vec.layers.12.self_attn.k_proj",
+            "llm2vec.layers.12.self_attn.v_proj",
+            "llm2vec.layers.12.self_attn.o_proj",
+            "llm2vec.layers.12.mlp.gate_proj",
+            "llm2vec.layers.12.mlp.up_proj",
+            "llm2vec.layers.12.mlp.down_proj",
+            "llm2vec.layers.13.self_attn.q_proj",
+            "llm2vec.layers.13.self_attn.k_proj",
+            "llm2vec.layers.13.self_attn.v_proj",
+            "llm2vec.layers.13.self_attn.o_proj",
+            "llm2vec.layers.13.mlp.gate_proj",
+            "llm2vec.layers.13.mlp.up_proj",
+            "llm2vec.layers.13.mlp.down_proj",
+            "llm2vec.layers.14.self_attn.q_proj",
+            "llm2vec.layers.14.self_attn.k_proj",
+            "llm2vec.layers.14.self_attn.v_proj",
+            "llm2vec.layers.14.self_attn.o_proj",
+            "llm2vec.layers.14.mlp.gate_proj",
+            "llm2vec.layers.14.mlp.up_proj",
+            "llm2vec.layers.14.mlp.down_proj",
+            "llm2vec.layers.15.self_attn.q_proj",
+            "llm2vec.layers.15.self_attn.k_proj",
+            "llm2vec.layers.15.self_attn.v_proj",
+            "llm2vec.layers.15.self_attn.o_proj",
+            "llm2vec.layers.15.mlp.gate_proj",
+            "llm2vec.layers.15.mlp.up_proj",
+            "llm2vec.layers.15.mlp.down_proj",
+            "llm2vec.layers.16.self_attn.q_proj",
+            "llm2vec.layers.16.self_attn.k_proj",
+            "llm2vec.layers.16.self_attn.v_proj",
+            "llm2vec.layers.16.self_attn.o_proj",
+            "llm2vec.layers.16.mlp.gate_proj",
+            "llm2vec.layers.16.mlp.up_proj",
+            "llm2vec.layers.16.mlp.down_proj",
+            "llm2vec.layers.17.self_attn.q_proj",
+            "llm2vec.layers.17.self_attn.k_proj",
+            "llm2vec.layers.17.self_attn.v_proj",
+            "llm2vec.layers.17.self_attn.o_proj",
+            "llm2vec.layers.17.mlp.gate_proj",
+            "llm2vec.layers.17.mlp.up_proj",
+            "llm2vec.layers.17.mlp.down_proj",
+            "llm2vec.layers.18.self_attn.q_proj",
+            "llm2vec.layers.18.self_attn.k_proj",
+            "llm2vec.layers.18.self_attn.v_proj",
+            "llm2vec.layers.18.self_attn.o_proj",
+            "llm2vec.layers.18.mlp.gate_proj",
+            "llm2vec.layers.18.mlp.up_proj",
+            "llm2vec.layers.18.mlp.down_proj",
+            "llm2vec.layers.19.self_attn.q_proj",
+            "llm2vec.layers.19.self_attn.k_proj",
+            "llm2vec.layers.19.self_attn.v_proj",
+            "llm2vec.layers.19.self_attn.o_proj",
+            "llm2vec.layers.19.mlp.gate_proj",
+            "llm2vec.layers.19.mlp.up_proj",
+            "llm2vec.layers.19.mlp.down_proj",
+            "llm2vec.layers.20.self_attn.q_proj",
+            "llm2vec.layers.20.self_attn.k_proj",
+            "llm2vec.layers.20.self_attn.v_proj",
+            "llm2vec.layers.20.self_attn.o_proj",
+            "llm2vec.layers.20.mlp.gate_proj",
+            "llm2vec.layers.20.mlp.up_proj",
+            "llm2vec.layers.20.mlp.down_proj",
+            "llm2vec.layers.21.self_attn.q_proj",
+            "llm2vec.layers.21.self_attn.k_proj",
+            "llm2vec.layers.21.self_attn.v_proj",
+            "llm2vec.layers.21.self_attn.o_proj",
+            "llm2vec.layers.21.mlp.gate_proj",
+            "llm2vec.layers.21.mlp.up_proj",
+            "llm2vec.layers.21.mlp.down_proj",
+            "llm2vec.layers.22.self_attn.q_proj",
+            "llm2vec.layers.22.self_attn.k_proj",
+            "llm2vec.layers.22.self_attn.v_proj",
+            "llm2vec.layers.22.self_attn.o_proj",
+            "llm2vec.layers.22.mlp.gate_proj",
+            "llm2vec.layers.22.mlp.up_proj",
+            "llm2vec.layers.22.mlp.down_proj",
+            "llm2vec.layers.23.self_attn.q_proj",
+            "llm2vec.layers.23.self_attn.k_proj",
+            "llm2vec.layers.23.self_attn.v_proj",
+            "llm2vec.layers.23.self_attn.o_proj",
+            "llm2vec.layers.23.mlp.gate_proj",
+            "llm2vec.layers.23.mlp.up_proj",
+            "llm2vec.layers.23.mlp.down_proj",
+            "llm2vec.layers.24.self_attn.q_proj",
+            "llm2vec.layers.24.self_attn.k_proj",
+            "llm2vec.layers.24.self_attn.v_proj",
+            "llm2vec.layers.24.self_attn.o_proj",
+            "llm2vec.layers.24.mlp.gate_proj",
+            "llm2vec.layers.24.mlp.up_proj",
+            "llm2vec.layers.24.mlp.down_proj",
+            "llm2vec.layers.25.self_attn.q_proj",
+            "llm2vec.layers.25.self_attn.k_proj",
+            "llm2vec.layers.25.self_attn.v_proj",
+            "llm2vec.layers.25.self_attn.o_proj",
+            "llm2vec.layers.25.mlp.gate_proj",
+            "llm2vec.layers.25.mlp.up_proj",
+            "llm2vec.layers.25.mlp.down_proj",
+            "llm2vec.layers.26.self_attn.q_proj",
+            "llm2vec.layers.26.self_attn.k_proj",
+            "llm2vec.layers.26.self_attn.v_proj",
+            "llm2vec.layers.26.self_attn.o_proj",
+            "llm2vec.layers.26.mlp.gate_proj",
+            "llm2vec.layers.26.mlp.up_proj",
+            "llm2vec.layers.26.mlp.down_proj",
+            "llm2vec.layers.27.self_attn.q_proj",
+            "llm2vec.layers.27.self_attn.k_proj",
+            "llm2vec.layers.27.self_attn.v_proj",
+            "llm2vec.layers.27.self_attn.o_proj",
+            "llm2vec.layers.27.mlp.gate_proj",
+            "llm2vec.layers.27.mlp.up_proj",
+            "llm2vec.layers.27.mlp.down_proj",
+            "llm2vec.layers.28.self_attn.q_proj",
+            "llm2vec.layers.28.self_attn.k_proj",
+            "llm2vec.layers.28.self_attn.v_proj",
+            "llm2vec.layers.28.self_attn.o_proj",
+            "llm2vec.layers.28.mlp.gate_proj",
+            "llm2vec.layers.28.mlp.up_proj",
+            "llm2vec.layers.28.mlp.down_proj",
+            "llm2vec.layers.29.self_attn.q_proj",
+            "llm2vec.layers.29.self_attn.k_proj",
+            "llm2vec.layers.29.self_attn.v_proj",
+            "llm2vec.layers.29.self_attn.o_proj",
+            "llm2vec.layers.29.mlp.gate_proj",
+            "llm2vec.layers.29.mlp.up_proj",
+            "llm2vec.layers.29.mlp.down_proj",
+            "llm2vec.layers.30.self_attn.q_proj",
+            "llm2vec.layers.30.self_attn.k_proj",
+            "llm2vec.layers.30.self_attn.v_proj",
+            "llm2vec.layers.30.self_attn.o_proj",
+            "llm2vec.layers.30.mlp.gate_proj",
+            "llm2vec.layers.30.mlp.up_proj",
+            "llm2vec.layers.30.mlp.down_proj",
+            "llm2vec.layers.31.self_attn.q_proj",
+            "llm2vec.layers.31.self_attn.k_proj",
+            "llm2vec.layers.31.self_attn.v_proj",
+            "llm2vec.layers.31.self_attn.o_proj",
+            "llm2vec.layers.31.mlp.gate_proj",
+            "llm2vec.layers.31.mlp.up_proj",
+            "llm2vec.layers.31.mlp.down_proj",
+        ],
+    )
+    peft_model = get_peft_model(model, lora_config)
+    peft_model.save_pretrained("../nllb-llm2vec-saved")
+    import json
+
+    with open("./model.safetensors.index.json", "r") as f:
+        print(json.load(f))