initial commit

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "word_embedding_dimension": 4096,
+  "pooling_mode_cls_token": false,
+  "pooling_mode_mean_tokens": true,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false,
+  "pooling_mode_weightedmean_tokens": false,
+  "pooling_mode_lasttoken": false,
+  "include_prompt": false
+}

README.md

@@ -0,0 +1,115 @@
+---
+language:
+- en
+license: cc-by-nc-4.0
+---
+## Introduction
+This is fixed version of `nvidia/NV-Embed-v1` to download model without error.
+
+## How to use
+
+Here is an example of how to encode queries and passages using Huggingface-transformer and Sentence-transformer.
+
+### Usage (HuggingFace Transformers)
+
+```python
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+
+# Each query needs to be accompanied by an corresponding instruction describing the task.
+task_name_to_instruct = {"example": "Given a question, retrieve passages that answer the question",}
+
+query_prefix = "Instruct: "+task_name_to_instruct["example"]+"\nQuery: "
+queries = [
+    'are judo throws allowed in wrestling?', 
+    'how to become a radiology technician in michigan?'
+    ]
+
+# No instruction needed for retrieval passages
+passage_prefix = ""
+passages = [
+    "Since you're reading this, you are probably someone from a judo background or someone who is just wondering how judo techniques can be applied under wrestling rules. So without further ado, let's get to the question. Are Judo throws allowed in wrestling? Yes, judo throws are allowed in freestyle and folkstyle wrestling. You only need to be careful to follow the slam rules when executing judo throws. In wrestling, a slam is lifting and returning an opponent to the mat with unnecessary force.",
+    "Below are the basic steps to becoming a radiologic technologist in Michigan:Earn a high school diploma. As with most careers in health care, a high school education is the first step to finding entry-level employment. Taking classes in math and science, such as anatomy, biology, chemistry, physiology, and physics, can help prepare students for their college studies and future careers.Earn an associate degree. Entry-level radiologic positions typically require at least an Associate of Applied Science. Before enrolling in one of these degree programs, students should make sure it has been properly accredited by the Joint Review Committee on Education in Radiologic Technology (JRCERT).Get licensed or certified in the state of Michigan."
+]
+
+# load model with tokenizer
+model = AutoModel.from_pretrained('bzantium/NV-Embed-v1', trust_remote_code=True)
+
+# get the embeddings
+max_length = 4096
+query_embeddings = model.encode(queries, instruction=query_prefix, max_length=max_length)
+passage_embeddings = model.encode(passages, instruction=passage_prefix, max_length=max_length)
+
+# normalize embeddings
+query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
+passage_embeddings = F.normalize(passage_embeddings, p=2, dim=1)
+
+# get the embeddings with DataLoader (spliting the datasets into multiple mini-batches)
+# batch_size=2
+# query_embeddings = model._do_encode(queries, batch_size=batch_size, instruction=query_prefix, max_length=max_length, num_workers=32)
+# passage_embeddings = model._do_encode(passages, batch_size=batch_size, instruction=passage_prefix, max_length=max_length, num_workers=32)
+
+scores = (query_embeddings @ passage_embeddings.T) * 100
+print(scores.tolist())
+#[[77.9402084350586, 0.4248958230018616], [3.757718086242676, 79.60113525390625]]
+```
+
+
+### Usage (Sentence-Transformers)
+
+```python
+import torch
+from sentence_transformers import SentenceTransformer
+
+# Each query needs to be accompanied by an corresponding instruction describing the task.
+task_name_to_instruct = {"example": "Given a question, retrieve passages that answer the question",}
+
+query_prefix = "Instruct: "+task_name_to_instruct["example"]+"\nQuery: "
+queries = [
+    'are judo throws allowed in wrestling?', 
+    'how to become a radiology technician in michigan?'
+    ]
+
+# No instruction needed for retrieval passages
+passages = [
+    "Since you're reading this, you are probably someone from a judo background or someone who is just wondering how judo techniques can be applied under wrestling rules. So without further ado, let's get to the question. Are Judo throws allowed in wrestling? Yes, judo throws are allowed in freestyle and folkstyle wrestling. You only need to be careful to follow the slam rules when executing judo throws. In wrestling, a slam is lifting and returning an opponent to the mat with unnecessary force.",
+    "Below are the basic steps to becoming a radiologic technologist in Michigan:Earn a high school diploma. As with most careers in health care, a high school education is the first step to finding entry-level employment. Taking classes in math and science, such as anatomy, biology, chemistry, physiology, and physics, can help prepare students for their college studies and future careers.Earn an associate degree. Entry-level radiologic positions typically require at least an Associate of Applied Science. Before enrolling in one of these degree programs, students should make sure it has been properly accredited by the Joint Review Committee on Education in Radiologic Technology (JRCERT).Get licensed or certified in the state of Michigan."
+]
+
+# load model with tokenizer
+model = SentenceTransformer('bzantium/NV-Embed-v1', trust_remote_code=True)
+model.max_seq_length = 4096
+model.tokenizer.padding_side="right"
+
+def add_eos(input_examples):
+  input_examples = [input_example + model.tokenizer.eos_token for input_example in input_examples]
+  return input_examples
+
+# get the embeddings
+batch_size = 2
+query_embeddings = model.encode(add_eos(queries), batch_size=batch_size, prompt=query_prefix, normalize_embeddings=True)
+passage_embeddings = model.encode(add_eos(passages), batch_size=batch_size, normalize_embeddings=True)
+
+scores = (query_embeddings @ passage_embeddings.T) * 100
+print(scores.tolist())
+```
+
+## Correspondence to
+Chankyu Lee ([email protected]), Wei Ping ([email protected])
+
+## Citation
+If you find this code useful in your research, please consider citing:
+
+```bibtex
+@misc{lee2024nvembed,
+      title={NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models}, 
+      author={Chankyu Lee and Rajarshi Roy and Mengyao Xu and Jonathan Raiman and Mohammad Shoeybi and Bryan Catanzaro and Wei Ping},
+      year={2024},
+      eprint={2405.17428},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```
+## License
+This model should not be used for any commercial purpose. Refer the [license](https://spdx.org/licenses/CC-BY-NC-4.0) for the detailed terms.

config.json

@@ -0,0 +1,100 @@
+{
+  "_name_or_path": "bzantium/NV-Embed-v1",
+  "add_eos": true,
+  "add_pad_token": true,
+  "architectures": [
+    "NVEmbedModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_nvembed.NVEmbedConfig",
+    "AutoModel": "modeling_nvembed.NVEmbedModel"
+  },
+  "is_mask_instruction": true,
+  "latent_attention_config": {
+    "model_type": "latent_attention"
+  },
+  "mask_type": "b",
+  "model_type": "nvembed",
+  "padding_side": "right",
+  "text_config": {
+    "_name_or_path": "bzantium/NV-Embed-v1",
+    "add_cross_attention": false,
+    "architectures": [
+      "MistralModel"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 14336,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 32768,
+    "min_length": 0,
+    "model_type": "bidir_mistral",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 10000.0,
+    "sep_token_id": null,
+    "sliding_window": 4096,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 32000
+  },
+  "torch_dtype": "float16",
+  "transformers_version": "4.37.2"
+}

config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.7.0",
+    "transformers": "4.37.2",
+    "pytorch": "2.2.0+cu121"
+  },
+  "prompts": {},
+  "default_prompt_name": null
+}

configuration_nvembed.py

@@ -0,0 +1,88 @@
+
+from typing import Literal
+from transformers import AutoConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+from transformers.models.mistral import MistralConfig
+
+NVEMBED_TYPE = "nvembed"
+LATENT_ATTENTION_TYPE = "latent_attention"
+BIDIR_MISTRAL_TYPE = "bidir_mistral"
+
+class NVEmbedConfig(PretrainedConfig):
+    model_type = "nvembed"
+    is_composition = False
+
+    def __init__(
+        self,
+        hidden_size=4096,
+        latent_attention_config=None,
+        text_config=None,
+        padding_side: Literal["right", "left"]="right",
+        add_pad_token: bool=True,
+        is_mask_instruction: bool = True,
+        add_eos: bool=True,
+        mask_type: str="b",
+        **kwargs,
+    ):
+        if isinstance(latent_attention_config, dict):
+            latent_attention_config["model_type"] = (
+                latent_attention_config["model_type"] if "model_type" in latent_attention_config else LATENT_ATTENTION_TYPE
+            )
+            latent_attention_config = CONFIG_MAPPING[latent_attention_config["model_type"]](**latent_attention_config)
+        elif latent_attention_config is None:
+            latent_attention_config = CONFIG_MAPPING[LATENT_ATTENTION_TYPE]()
+
+        self.latent_attention_config = latent_attention_config
+
+        if isinstance(text_config, dict):
+            text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama"
+            text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
+        elif text_config is None:
+            text_config = None
+
+        self.hidden_size = hidden_size
+        self.text_config = text_config
+        self.padding_side = padding_side
+        self.is_mask_instruction = is_mask_instruction
+        self.add_pad_token = add_pad_token
+        self.add_eos = add_eos
+        self.mask_type = mask_type
+
+        super().__init__(**kwargs)
+
+
+class LatentAttentionConfig(PretrainedConfig):
+    model_type = LATENT_ATTENTION_TYPE
+    is_composition = False
+    _name_or_path = "latent_attention"
+
+    def __init__(
+        self,
+        num_latents_value: int=512,
+        num_cross_heads: int=8,
+        output_normalize: bool=True,
+        hidden_dim: int=4096,
+        latent_dim: int=4096,
+        cross_dim_head: int=4096,
+        **kwargs,
+    ):
+        self.num_latents_value = num_latents_value
+        self.num_cross_heads = num_cross_heads
+        self.output_normalize = output_normalize
+        self.hidden_dim = hidden_dim
+        self.latent_dim = latent_dim
+        self.cross_dim_head = cross_dim_head
+
+
+class BidirectionalMistralConfig(MistralConfig):
+    model_type = BIDIR_MISTRAL_TYPE
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+AutoConfig.register(NVEMBED_TYPE, NVEmbedConfig)
+AutoConfig.register(LATENT_ATTENTION_TYPE, LatentAttentionConfig)
+AutoConfig.register(BIDIR_MISTRAL_TYPE, BidirectionalMistralConfig)
+
+NVEmbedConfig.register_for_auto_class()
+LatentAttentionConfig.register_for_auto_class()
+BidirectionalMistralConfig.register_for_auto_class()

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,311 @@
+{
+  "metadata": {
+    "total_size": 15702032384
+  },
+  "weight_map": {
+    "embedding_model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.19.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.20.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.30.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.30.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.30.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.30.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.30.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "embedding_model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.31.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "embedding_model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "embedding_model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.8.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "embedding_model.norm.weight": "model-00004-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.fn.to_kv.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.fn.to_out.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.fn.to_q.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm.bias": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm_context.bias": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm_context.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.0.bias": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.0.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.2.bias": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.2.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.norm.bias": "model-00001-of-00004.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.norm.weight": "model-00001-of-00004.safetensors",
+    "latent_attention_model.latents": "model-00001-of-00004.safetensors"
+  }
+}

modeling_nvembed.py

@@ -0,0 +1,436 @@
+from typing import List, Union, Dict, Mapping, Optional, Tuple, TypedDict
+import torch
+import os
+import json
+import numpy as np
+from functools import partial
+from contextlib import nullcontext
+from transformers import AutoModel, PreTrainedTokenizerFast, BatchEncoding, DataCollatorWithPadding
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.auto import AutoTokenizer
+from transformers.models.mistral.modeling_mistral import MISTRAL_INPUTS_DOCSTRING
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
+from transformers import MistralModel, MistralConfig
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    logging,
+)
+from einops import rearrange, repeat
+from tqdm.auto import tqdm
+from datasets import Dataset
+from torch.utils.data import DataLoader
+from .configuration_nvembed import NVEmbedConfig, LatentAttentionConfig, BidirectionalMistralConfig
+
+logger = logging.get_logger(__name__)
+
+class NVEmbedFeatures(TypedDict):
+    input_dict: torch.Tensor
+    attention_mask: torch.Tensor
+    pool_mask: torch.Tensor
+
+class BidirectionalMistralModel(MistralModel):
+    config_class = BidirectionalMistralConfig
+    
+    def __init__(self, config: MistralConfig):
+        super().__init__(config)
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+        self._attn_implementation = "eager"
+
+    @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
+            is_padding_right = attention_mask[:, -1].sum().item() != batch_size
+            if is_padding_right:
+                raise ValueError(
+                    "You are attempting to perform batched generation with padding_side='right'"
+                    " this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
+                    " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                )
+
+        if self._attn_implementation == "flash_attention_2":
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._attn_implementation == "sdpa" and not output_attentions:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                attention_mask, inputs_embeds.dtype
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_attention_mask(
+                attention_mask, inputs_embeds.dtype,
+            )
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+    
+def _move_to_device(maybe_tensor, device: torch.device):
+    if torch.is_tensor(maybe_tensor):
+        return maybe_tensor.to(device, non_blocking=device.type == "cuda")
+    elif isinstance(maybe_tensor, dict):
+        return {key: _move_to_device(value, device) for key, value in maybe_tensor.items()}
+    elif isinstance(maybe_tensor, list):
+        return [_move_to_device(x, device) for x in maybe_tensor]
+    elif isinstance(maybe_tensor, tuple):
+        return tuple([_move_to_device(x, device) for x in maybe_tensor])
+    elif isinstance(maybe_tensor, Mapping):
+        return type(maybe_tensor)({k: _move_to_device(v, device) for k, v in maybe_tensor.items()})
+    else:
+        return maybe_tensor
+
+def move_to_device(sample, device: torch.device):
+    if device.type == "cpu":
+        return sample
+    
+    if len(sample) == 0:
+        return {}
+    return _move_to_device(sample, device)
+
+
+def input_transform_func(
+    tokenizer: PreTrainedTokenizerFast,
+    examples: Dict[str, List],
+    always_add_eos: bool,
+    max_length: int,
+    instruction: str,
+) -> BatchEncoding:
+    if always_add_eos:
+        examples['input_texts'] = [instruction + input_example + tokenizer.eos_token for input_example in examples['input_texts']]
+    batch_dict = tokenizer(
+        examples['input_texts'],
+        max_length=max_length,
+        padding=True,
+        return_token_type_ids=False,
+        return_tensors="pt",
+        truncation=True)
+    return batch_dict
+
+
+class PreNorm(torch.nn.Module):
+    def __init__(self, dim, fn, context_dim = None):
+        super().__init__()
+        self.fn = fn
+        self.norm = torch.nn.LayerNorm(dim)
+        self.norm_context = torch.nn.LayerNorm(context_dim) if exists(context_dim) else None
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        if exists(self.norm_context):
+            context = kwargs['context']
+            normed_context = self.norm_context(context)
+            kwargs.update(context = normed_context)
+        return self.fn(x, **kwargs)
+
+class GEGLU(torch.nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim = -1)
+        return x * torch.nn.functional.gelu(gates)
+
+class FeedForward(torch.nn.Module):
+    def __init__(self, dim, mult = 4):
+        super().__init__()
+        self.net = torch.nn.Sequential(torch.nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            torch.nn.Linear(dim * mult, dim))
+
+    def forward(self, x):
+        return self.net(x)
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+class Attention(torch.nn.Module):
+    def __init__(self, query_dim, context_dim = None, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+
+        self.to_q = torch.nn.Linear(query_dim, inner_dim, bias = False)
+        self.to_kv = torch.nn.Linear(context_dim, inner_dim * 2, bias = False)
+        self.to_out = torch.nn.Linear(inner_dim, query_dim, bias = False)
+
+    def forward(self, x, context = None, mask = None):
+        h = self.heads
+        q = self.to_q(x)
+        context = default(context, x)
+        k, v = self.to_kv(context).chunk(2, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True):
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+        out = rearrange(out, '(b h) n d -> b n (h d)', h = h)
+        return self.to_out(out)
+
+
+class LatentAttentionModel(PreTrainedModel):
+    config_class = LatentAttentionConfig
+
+    def __init__(self, config: LatentAttentionConfig):
+        super().__init__(config)
+        ## cross-attention block
+        num_latents, latent_dim, cross_heads, cross_dim_head = config.num_latents_value, config.latent_dim, config.num_cross_heads, config.cross_dim_head
+        dim = config.hidden_dim
+        # init latent_attention and latents
+        self.cross_attend_blocks = torch.nn.ModuleList([
+            PreNorm(latent_dim, Attention(latent_dim, dim, heads = cross_heads, dim_head = cross_dim_head),
+                    context_dim = dim),
+            PreNorm(latent_dim, FeedForward(latent_dim)),
+        ])
+        self.output_normalize = config.output_normalize
+        self.register_parameter("latents", torch.nn.Parameter(torch.randn(num_latents, latent_dim)))
+
+    def forward(self, hiddens, attention_mask: torch.Tensor=None):
+        ## cross-attention block
+        cross_attn, cross_ff = self.cross_attend_blocks
+        b, *_, device = *hiddens.shape, hiddens.device
+        x = repeat(self.latents, 'n d -> b n d', b = b)
+        hiddens = cross_attn(hiddens, context = x, mask = None) + hiddens
+        hiddens = cross_ff(hiddens) + hiddens
+        if attention_mask !=None:
+            s = torch.sum(hiddens * attention_mask.unsqueeze(-1).float(), dim=1)
+            d = attention_mask.sum(dim=1, keepdim=True).float()
+            hiddens = s / d
+            if self.output_normalize:
+                hiddens = torch.nn.functional.normalize(hiddens, p=2, dim=-1)
+        return hiddens
+    
+class NVEmbedModel(PreTrainedModel):
+    config_class = NVEmbedConfig
+
+    def __init__(self, config: NVEmbedConfig):
+        super().__init__(config)
+        self.latent_attention_model = AutoModel.from_config(config.latent_attention_config)
+        self.embedding_model = AutoModel.from_config(
+            config.text_config,
+        ) if config.text_config is not None else None
+        self.tokenizer = AutoTokenizer.from_pretrained(config.text_config._name_or_path, trust_remote_code=True) if config.text_config is not None else None
+        self.padding_side = config.padding_side
+        self.is_mask_instruction = config.is_mask_instruction
+        self.add_eos = config.add_eos
+        self.mask_type = config.mask_type
+        if config.add_pad_token and self.tokenizer is not None:
+            self.add_pad_token()
+
+    def add_pad_token(self):
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = self.padding_side
+    
+    def prepare_kwargs_from_batch(self, batch_dict: dict, instruction_lens: int, device: torch.device):
+        batch_dict = move_to_device(batch_dict, device)
+        attention_mask = batch_dict['attention_mask'].clone() if 'attention_mask' in batch_dict else None
+        if (attention_mask is not None and
+            self.padding_side == "right" and
+            self.is_mask_instruction == True and
+            instruction_lens > 0):
+            # Mask out the instruction tokens for mean-pooling
+            attention_mask[:, :instruction_lens] = 0
+        features: NVEmbedFeatures = {
+            'input_ids': torch.tensor(batch_dict.get('input_ids').to(batch_dict.get('input_ids')).long()),
+            'attention_mask': batch_dict['attention_mask'],
+            'pool_mask': attention_mask,
+        }
+        return features
+
+    @torch.no_grad()
+    def _do_encode(self,
+        prompts: List[str],
+        batch_size,
+        instruction: str,
+        max_length: int=4096,
+        num_workers: int=32,
+    ) -> np.ndarray:
+        dataset: Dataset = Dataset.from_dict({'input_texts': prompts})
+        dataset.set_transform(partial(input_transform_func,
+                                      self.tokenizer,
+                                      always_add_eos=True,
+                                      max_length=max_length,
+                                      instruction=instruction))
+
+        data_collator = DataCollatorWithPadding(self.tokenizer)
+        data_loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            drop_last=False,
+            num_workers=num_workers,
+            collate_fn=data_collator,
+            pin_memory=True)
+
+        if self.padding_side == "right" and self.is_mask_instruction == True and len(instruction) > 0:
+            instruction_lens = len(self.tokenizer.tokenize(instruction))
+        else:
+            instruction_lens = 0
+
+        encoded_embeds = []
+        device = next(self.embedding_model.parameters()).device
+        for batch_dict in tqdm(data_loader, desc='encoding', mininterval=10):
+            features = self.prepare_kwargs_from_batch(batch_dict, instruction_lens, device=device)
+            embeds=self(**features)["sentence_embeddings"].squeeze(1)
+            encoded_embeds.append(embeds.cpu().numpy())
+        return np.concatenate(encoded_embeds, axis=0)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pool_mask: Optional[torch.Tensor]=None, return_dict: bool=True):
+        autocast_ctx = torch.cuda.amp.autocast if torch.cuda.is_available() else nullcontext
+        with autocast_ctx():
+            ## decoder only layer
+            outputs = self.embedding_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            ## latent attention layer
+            embeds = self.latent_attention_model(
+                outputs.last_hidden_state,
+                pool_mask,
+            )
+        if not return_dict:
+            return (embeds,)
+        return {"sentence_embeddings": embeds}
+        
+    
+    @torch.no_grad()
+    def encode(self, prompts: List[str], instruction: str="", max_length: int=4096):
+        if self.padding_side == "right" and self.is_mask_instruction == True and len(instruction) > 0:
+            instruction_lens = len(self.tokenizer.tokenize(instruction))
+        else:
+            instruction_lens = 0
+        
+        device = next(self.embedding_model.parameters()).device
+        batch_dict = input_transform_func(self.tokenizer,
+                                          {"input_texts": [prompt for prompt in prompts]},
+                                          always_add_eos=True,
+                                          max_length=max_length,
+                                          instruction=instruction)
+
+        features: NVEmbedFeatures = self.prepare_kwargs_from_batch(batch_dict, instruction_lens, device=device)
+        return self(**features)["sentence_embeddings"].squeeze(1)
+
+
+## AutoModel Register
+AutoModel.register(NVEmbedConfig, NVEmbedModel)
+AutoModel.register(LatentAttentionConfig, LatentAttentionModel)
+AutoModel.register(BidirectionalMistralConfig, BidirectionalMistralModel)
+
+## Register for auto class
+NVEmbedModel.register_for_auto_class("AutoModel")
+LatentAttentionModel.register_for_auto_class("AutoModel")
+BidirectionalMistralModel.register_for_auto_class("AutoModel")

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 4096,
+  "do_lower_case": false
+}

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,42 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}