Upload 16 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

1_Pooling/config.json

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

README.md

@@ -0,0 +1,44 @@
+# Model Card for GeoEmbedding
+The GeoEmbedding model is a geoscience-specific text embedding model built upon a high-performance large language model and fine-tuned on both general-purpose and in-domain geoscientific datasets. It produces accurate, context-aware vector representations of geoscientific texts, forming the backbone of vector-based retrieval in the RAG pipeline.
+
+## Quick Start
+To load the GeoEmbedding model with Transformer, use the following snippet:
+```python
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+task_description = 'Given a web search query, retrieve relevant passages that answer the query'
+def get_detailed_instruct(task_description: str, query: str) -> str:
+    return f'Instruct: {task_description}\nQuery: {query}'
+
+model_name_or_path = 'GeoGPT/GeoEmbedding'
+
+model = SentenceTransformer(model_name_or_path, device="cuda", trust_remote_code=True)
+
+queries = [
+    "What is the main cause of earthquakes?",
+    "How do sedimentary rocks form?",
+]
+
+passages = [
+    "Earthquakes occur due to the sudden release of energy in the Earth's crust, often caused by tectonic plate movements along fault lines.",
+    "Sedimentary rocks form through the deposition and compaction of mineral and organic particles over time, typically in water bodies.",
+]
+
+queries = [get_detailed_instruct(task_description, query) for query in queries]
+
+q_vecs = model.encode(queries, normalize_embeddings=True)
+p_vecs = model.encode(passages, normalize_embeddings=True)
+
+print(np.dot(q_vecs, p_vecs.T)) 
+#[[0.6369  0.2092 ]
+# [0.2499  0.8411 ]]
+```
+
+## License and Uses
+GeoEmbedding is liscensed under Apache License 2.0, and is trained on the foundation of [Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1), which is also licensed under the Apache License 2.0. It is your responsibility to ensure that your use of GeoEmbedding adheres to the terms of both the GeoEmbedding model and its upstream dependency, [Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1).
+
+The model is not intended for use in any manner that violates applicable laws or regulations, nor for any activities prohibited by the license agreement. Additionally, it should not be used in languages other than those explicitly supported, as outlined in this model card.
+
+## Limitations
+GeoEmbedding is trained on English datasets, and performance may be suboptimal for other languages.

config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "MistralDualModel"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_mistral_dual.MistralDualConfig",
+    "AutoModel": "modeling_mistral_dual.MistralDualModel"
+  },
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 32768,
+  "model_type": "mistraldual",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "pad_token_id": 2,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "use_cache": false,
+  "vocab_size": 32000
+}

config_sentence_transformers.json

@@ -0,0 +1,10 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.3.1",
+    "transformers": "4.51.3",
+    "pytorch": "2.5.1+cu124"
+  },
+  "prompts": {},
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

configuration_mistral_dual.py

@@ -0,0 +1,13 @@
+from transformers import MistralConfig, AutoConfig
+
+class MistralDualConfig(MistralConfig):
+    model_type = "mistraldual"
+    def __init__(
+        self,
+        use_cache=False,
+        **kwargs,
+    ):
+        super().__init__(use_cache=use_cache, **kwargs)
+
+AutoConfig.register("mistraldual", MistralDualConfig)
+MistralDualConfig.register_for_auto_class()

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

model.safetensors.index.json

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

modeling_mistral_dual.py

@@ -0,0 +1,176 @@
+from typing import Optional, Tuple, Union
+from functools import partial
+
+import torch
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.utils import logging
+from transformers import AutoModel
+from transformers.models.mistral.configuration_mistral import MistralConfig
+from transformers.models.mistral.modeling_mistral import MistralModel
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
+
+from .configuration_mistral_dual import MistralDualConfig
+
+logger = logging.get_logger(__name__)
+
+class MistralDualModel(MistralModel):
+    config_class = MistralDualConfig
+
+    def __init__(self, config: MistralDualConfig):
+        super().__init__(config)
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = 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,
+        cache_position: Optional[torch.LongTensor] = None,
+        is_causal = False,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ) -> 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
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+        )
+
+        # print(causal_mask)
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    partial(decoder_layer.__call__, is_causal=is_causal),
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    position_embeddings=position_embeddings,
+                    is_causal=is_causal,
+                    **flash_attn_kwargs,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            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,)
+
+        output = BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+        return output if return_dict else output.to_tuple()
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: MistralConfig,
+        past_key_values: Cache,
+    ):
+        """
+        Creates a bidirectional 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`,
+        where all tokens can attend to all others.
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            return attention_mask  # Already in correct shape
+
+        min_dtype = torch.finfo(dtype).min
+        # Create a full attention mask allowing all tokens to attend to all others
+        bidirectional_mask = torch.zeros((sequence_length, target_length), dtype=dtype, device=device)
+        bidirectional_mask = bidirectional_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+
+        if attention_mask is not None:
+            bidirectional_mask = bidirectional_mask.clone()  # Ensure contiguous memory for in-place edit
+            if attention_mask.shape[-1] > target_length:
+                attention_mask = attention_mask[:, :target_length]
+            mask_length = attention_mask.shape[-1]
+            padding_mask = bidirectional_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
+            padding_mask = padding_mask == 0
+            bidirectional_mask[:, :, :, :mask_length] = bidirectional_mask[:, :, :, :mask_length].masked_fill(
+                padding_mask, min_dtype
+            )
+        
+        return bidirectional_mask
+
+
+AutoModel.register(MistralDualConfig, MistralDualModel)
+MistralDualModel.register_for_auto_class()

modules.json

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

sentence_bert_config.json

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

special_tokens_map.json

@@ -0,0 +1,35 @@
+{
+  "additional_special_tokens": [
+    "<unk>",
+    "<s>",
+    "</s>"
+  ],
+  "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.json

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

tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": true,
+  "add_prefix_space": null,
+  "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": [
+    "<unk>",
+    "<s>",
+    "</s>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "max_length": 512,
+  "model_max_length": 1024,
+  "pad_to_multiple_of": null,
+  "pad_token": "</s>",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "stride": 0,
+  "tokenizer_class": "LlamaTokenizerFast",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}