Upload 15 files

1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+    "word_embedding_dimension": 384,
+    "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": true
+}

2_Dense/config.json

@@ -0,0 +1,6 @@
+{
+    "in_features": 384,
+    "out_features": 1024,
+    "bias": true,
+    "activation_function": "torch.nn.modules.linear.Identity"
+}

2_Dense/model.safetensors

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

README.md

@@ -1,3 +1,191 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0  
+base_model: microsoft/MiniLM-L6-v2  
+tags:  
+- transformers  
+- sentence-transformers  
+- sentence-similarity  
+- feature-extraction  
+- text-embeddings-inference  
+- information-retrieval  
+- knowledge-distillation  
+language:
+- en
+---
+<div style="display: flex; justify-content: center;">      
+    <div style="display: flex; align-items: center; gap: 10px;">      
+        <img src="logo.webp" alt="MongoDB Logo" style="height: 36px; width: auto; border-radius: 4px;">      
+        <span style="font-size: 32px; font-weight: bold">MongoDB/mdbr-leaf-mt</span>      
+    </div>      
+</div>  
+
+# Introduction
+
+`mdbr-leaf-mt` is a compact high-performance text embedding model designed for classification, clustering, semantic sentence similarity and summarization tasks. 
+
+To enable even greater efficiency, `mdbr-leaf-mt` supports [flexible asymmetric architectures](#asymmetric-retrieval-setup) and is robust to [vector quantization](#vector-quantization) and [MRL truncation](#mrl).
+
+If you are looking to perform semantic search / information retrieval (e.g. for RAGs), please check out our [`mdbr-leaf-ir`](https://huggingface.co/MongoDB/mdbr-leaf-ir) model, which is specifically trained for these tasks.
+
+> [!Note]  
+> **Note**: this model has been developed by the ML team of MongoDB Research. At the time of writing it is not used in any of MongoDB's commercial product or service offerings.
+
+# Technical Report
+
+A technical report detailing our proposed `LEAF` training procedure is [available here (TBD)](http://FILL_HERE_ARXIV_LINK).
+
+# Highlights  
+
+* **State-of-the-Art Performance**: `mdbr-leaf-mt` achieves new state-of-the-art results for compact embedding models, ranking <span style="color:red">#TBD</span> on the [public MTEB v2 (Eng) benchmark leaderboard](https://huggingface.co/spaces/mteb/leaderboard) for models <30M parameters with an average score of <span style="color:red">[TBD HERE]</span>.
+* **Flexible Architecture Support**: `mdbr-leaf-mt` supports asymmetric retrieval architectures enabling even greater retrieval results. [See below](#asymmetric-retrieval-setup) for more information.
+* **MRL and Quantization Support**: embedding vectors generated by `mdbr-leaf-mt` compress well when truncated (MRL) and/or can be stored using more efficient types like `int8` and `binary`.  [See below](#mrl) for more information.
+
+# Quickstart  
+  
+## Sentence Transformers  
+  
+```python  
+from sentence_transformers import SentenceTransformer  
+  
+# Load the model  
+model = SentenceTransformer("MongoDB/mdbr-leaf-mt")  
+  
+# Example queries and documents  
+queries = [
+    "What is machine learning?",  
+    "How does neural network training work?"  
+]  
+  
+documents = [  
+    "Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.",  
+    "Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors."  
+]  
+  
+# Encode queries and documents  
+query_embeddings = model.encode(queries, prompt_name="query")  
+document_embeddings = model.encode(documents)  
+  
+# Compute similarity scores  
+scores = model.similarity(query_embeddings, document_embeddings)  
+
+# Print results
+for i, query in enumerate(queries):
+    print(f"Query: {query}")
+    for j, doc in enumerate(documents):
+        print(f" Similarity: {scores[i, j]:.4f} | Document {j}: {doc[:80]}...")
+
+# Query: What is machine learning?
+#  Similarity: 0.9063 | Document 0: Machine learning is a subset of ...
+#  Similarity: 0.7287 | Document 1: Neural networks are trained ...
+# 
+# Query: How does neural network training work?
+#  Similarity: 0.6725 | Document 0: Machine learning is a subset of ...
+#  Similarity: 0.8287 | Document 1: Neural networks are trained ...
+```
+  
+## Transformers Usage  
+
+See [here](https://huggingface.co/MongoDB/mdbr-leaf-mt/blob/main/transformers_example_mt.ipynb).
+  
+## Asymmetric Retrieval Setup  
+  
+`mdbr-leaf-mt` is *aligned* to [`mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1), the model it has been distilled from, making the asymmetric system below possible: 
+  
+```python  
+# Use mdbr-leaf-mt for query encoding (real-time, low latency)  
+query_model = SentenceTransformer("MongoDB/mdbr-leaf-mt")  
+query_embeddings = query_model.encode(queries, prompt_name="query")  
+
+# Use a larger model for document encoding (one-time, at index time)  
+doc_model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")  
+document_embeddings = doc_model.encode(documents)  
+  
+# Compute similarities  
+scores = query_model.similarity(query_embeddings, document_embeddings)  
+```
+Retrieval results from asymmetric mode are usually superior to the [standard mode above](#sentence-transformers).
+
+## MRL Truncation
+
+Embeddings have been trained via [MRL](https://arxiv.org/abs/2205.13147) and can be truncated for more efficient storage:
+```python
+from torch.nn import functional as F
+
+query_embeds = model.encode(queries, prompt_name="query", convert_to_tensor=True)
+doc_embeds = model.encode(documents, convert_to_tensor=True)
+
+# Truncate and normalize according to MRL
+query_embeds = F.normalize(query_embeds[:, :256], dim=-1)
+doc_embeds = F.normalize(doc_embeds[:, :256], dim=-1)
+
+similarities = model.similarity(query_embeds, doc_embeds)
+
+print('After MRL:')
+print(f"* Embeddings dimension: {query_embeds.shape[1]}")
+print(f"* Similarities:\n\t{similarities}")
+
+# After MRL:
+# * Embeddings dimension: 256
+# * Similarities:
+# 	  tensor([[0.9164, 0.7219],
+#             [0.6682, 0.8393]], device='cuda:0')
+```
+
+## Vector Quantization
+Vector quantization, for example to `int8` or `binary`, can be performed as follows:
+
+**Note**: For vector quantization to types other than binary, we suggest performing a calibration to determine the optimal ranges, [see here](https://sbert.net/examples/sentence_transformer/applications/embedding-quantization/README.html#scalar-int8-quantization). 
+Good initial values are -1.0 and +1.0.
+```python
+from sentence_transformers.quantization import quantize_embeddings
+import torch
+
+query_embeds = model.encode(queries, prompt_name="query")
+doc_embeds = model.encode(documents)
+
+# Quantize embeddings to int8 using -1.0 and +1.0
+ranges = torch.tensor([[-1.0], [+1.0]]).expand(2, query_embeds.shape[1]).cpu().numpy()
+query_embeds = quantize_embeddings(query_embeds, "int8", ranges=ranges)
+doc_embeds = quantize_embeddings(doc_embeds, "int8", ranges=ranges)
+
+# Calculate similarities; cast to int64 to avoid under/overflow
+similarities = query_embeds.astype(int) @ doc_embeds.astype(int).T
+
+print('After quantization:')
+print(f"* Embeddings type: {query_embeds.dtype}")
+print(f"* Similarities:\n{similarities}")
+
+# After quantization:
+# * Embeddings type: int8
+# * Similarities:
+#   [[2202032 1422868]
+#    [1421197 1845580]]
+```
+
+# Evaluation
+
+The checkpoint used to produce the scores presented in the paper [is here](https://huggingface.co/MongoDB/mdbr-leaf-mt/commit/ea98995e96beac21b820aa8ad9afaa6fd29b243d).
+
+# Citation  
+  
+If you use this model in your work, please cite:  
+  
+```bibtex  
+@article{mdb_leaf,  
+  title         = {LEAF: Lightweight Embedding Alignment Knowledge Distillation Framework},  
+  author        = {Robin Vujanic and Thomas Rueckstiess},  
+  year          = {2025}
+  eprint        = {TBD},
+  archiveprefix = {arXiv},
+  primaryclass  = {FILL HERE},
+  url           = {FILL HERE}
+}  
+```  
+  
+# License  
+  
+This model is released under Apache 2.0 License.  
+  
+# Contact  
+  
+For questions or issues, please open an issue or pull request. You can also contact the MongoDB ML Research team at [email protected].  

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.52.4",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "SentenceTransformer",
+  "__version__": {
+    "sentence_transformers": "5.1.0",
+    "transformers": "4.52.4",
+    "pytorch": "2.6.0+cu126"
+  },
+  "prompts": {
+    "query": "Represent this sentence for searching relevant passages: ",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

model.safetensors

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

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_Dense",
+    "type": "sentence_transformers.models.Dense"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,7 @@
+{
+    "max_seq_length": 512,
+    "do_lower_case": false,
+    "model_args": {
+        "add_pooling_layer": false
+    }
+}

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,65 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 128,
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}

transformers_example_mt.ipynb

@@ -0,0 +1,167 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "2a12a2b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from safetensors import safe_open\n",
+    "import torch\n",
+    "from torch.nn import functional as F\n",
+    "from transformers import AutoModel, AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "148ce181",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# First clone the model locally\n",
+    "!git clone https://huggingface.co/MongoDB/mdbr-leaf-mt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "ba9ec6c7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some weights of BertModel were not initialized from the model checkpoint at mdbr-leaf-mt and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n",
+      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Then load it\n",
+    "MODEL = \"mdbr-leaf-mt\"\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
+    "model = AutoModel.from_pretrained(MODEL)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ebaf1a76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tensors = {}\n",
+    "with safe_open(MODEL + \"/2_Dense/model.safetensors\", framework=\"pt\") as f:\n",
+    "    for k in f.keys():\n",
+    "        tensors[k] = f.get_tensor(k)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "03ffcd9c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Similarities:\n",
+      "tensor([[0.9063, 0.7287],\n",
+      "        [0.6725, 0.8287]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "if 'linear.bias' in tensors:\n",
+    "    W_out = torch.nn.Linear(in_features=384, out_features=1024, bias=True)\n",
+    "    W_out.load_state_dict({\n",
+    "        \"weight\": tensors[\"linear.weight\"], \n",
+    "        \"bias\": tensors[\"linear.bias\"]\n",
+    "    })\n",
+    "else:\n",
+    "    W_out = torch.nn.Linear(in_features=384, out_features=1024, bias=False)\n",
+    "    W_out.load_state_dict({\n",
+    "        \"weight\": tensors[\"linear.weight\"]\n",
+    "    })\n",
+    "\n",
+    "_ = model.eval()\n",
+    "_ = W_out.eval()\n",
+    "\n",
+    "# Example queries and documents  \n",
+    "queries = [\n",
+    "    \"What is machine learning?\",  \n",
+    "    \"How does neural network training work?\"  \n",
+    "]  \n",
+    "  \n",
+    "documents = [  \n",
+    "    \"Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.\",  \n",
+    "    \"Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors.\"  \n",
+    "]\n",
+    "\n",
+    "# Tokenize\n",
+    "QUERY_PREFIX = 'Represent this sentence for searching relevant passages: '\n",
+    "queries_with_prefix = [QUERY_PREFIX + query for query in queries]\n",
+    "\n",
+    "query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)\n",
+    "document_tokens =  tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=512)\n",
+    "\n",
+    "# Perform Inference\n",
+    "with torch.inference_mode():\n",
+    "    y_queries = model(**query_tokens).last_hidden_state\n",
+    "    y_docs = model(**document_tokens).last_hidden_state\n",
+    "\n",
+    "    # perform pooling\n",
+    "    y_queries = y_queries * query_tokens.attention_mask.unsqueeze(-1)\n",
+    "    y_queries_pooled = y_queries.sum(dim=1) / query_tokens.attention_mask.sum(dim=1, keepdim=True)\n",
+    "\n",
+    "    y_docs = y_docs * document_tokens.attention_mask.unsqueeze(-1)\n",
+    "    y_docs_pooled = y_docs.sum(dim=1) / document_tokens.attention_mask.sum(dim=1, keepdim=True)\n",
+    "\n",
+    "    # map to desired output dimension\n",
+    "    query_embeddings = W_out(y_queries_pooled)\n",
+    "    document_embeddings = W_out(y_docs_pooled)\n",
+    "\n",
+    "similarities = F.cosine_similarity(query_embeddings.unsqueeze(0), document_embeddings.unsqueeze(1), dim=-1).T\n",
+    "print(f\"Similarities:\\n{similarities}\")\n",
+    "\n",
+    "# Similarities:\n",
+    "#  tensor([[0.6857, 0.4598],\n",
+    "#          [0.4238, 0.5723]])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a2b0244",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "alexis",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}