change_idf

README.md

@@ -372,7 +372,7 @@ This is a [Asymmetric Inference-free SPLADE Sparse Encoder](https://www.sbert.ne
 ```
 SparseEncoder(
   (0): Router(
-    (query_0_IDF): IDF ({'frozen': False}, dim:30522, tokenizer: BertTokenizerFast)
+    (query_0_SparseStaticEmbedding): SparseStaticEmbedding ({'frozen': False}, dim:30522, tokenizer: BertTokenizerFast)
     (document_0_MLMTransformer): MLMTransformer({'max_seq_length': 512, 'do_lower_case': False}) with MLMTransformer model: BertForMaskedLM 
     (document_1_SpladePooling): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': 30522})
   )

config.json

@@ -1,19 +0,0 @@
-{
-  "types": {
-    "query_0_IDF": "sentence_transformers.sparse_encoder.models.IDF",
-    "document_0_MLMTransformer": "sentence_transformers.sparse_encoder.models.MLMTransformer",
-    "document_1_SpladePooling": "sentence_transformers.sparse_encoder.models.SpladePooling"
-  },
-  "structure": {
-    "query": [
-      "query_0_IDF"
-    ],
-    "document": [
-      "document_0_MLMTransformer",
-      "document_1_SpladePooling"
-    ]
-  },
-  "parameters": {
-    "allow_empty_key": true
-  }
-}

modules.json

@@ -3,6 +3,6 @@
     "idx": 0,
     "name": "0",
     "path": "",
-    "type": "sentence_transformers.models.Asym"
+    "type": "sentence_transformers.models.Router"
   }
 ]

query_0_IDF/special_tokens_map.json

@@ -1,37 +0,0 @@
-{
-  "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
-  }
-}

{query_0_IDF → query_0_SparseStaticEmbedding}/config.json

@@ -1,3 +1,3 @@
-{
-    "frozen": false
+{
+    "frozen": false
 }

{query_0_IDF → query_0_SparseStaticEmbedding}/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cd58ea1d4475f9557440a6feabafd8bfe9dc44e8cc6817ac931b8f8692b12c72
+oid sha256:82ca927cba4916c595053947bfea48b28692538d894ab9f94cceca63f57a845a
 size 122168

query_0_SparseStaticEmbedding/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

{query_0_IDF → query_0_SparseStaticEmbedding}/tokenizer.json

@@ -1,19 +1,7 @@
 {
   "version": "1.0",
-  "truncation": {
-    "direction": "Right",
-    "max_length": 512,
-    "strategy": "LongestFirst",
-    "stride": 0
-  },
-  "padding": {
-    "strategy": "BatchLongest",
-    "direction": "Right",
-    "pad_to_multiple_of": null,
-    "pad_id": 0,
-    "pad_type_id": 0,
-    "pad_token": "[PAD]"
-  },
+  "truncation": null,
+  "padding": null,
   "added_tokens": [
     {
       "id": 0,

{query_0_IDF → query_0_SparseStaticEmbedding}/tokenizer_config.json

@@ -1,65 +1,58 @@
-{
-  "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": true,
-  "cls_token": "[CLS]",
-  "do_basic_tokenize": true,
-  "do_lower_case": true,
-  "extra_special_tokens": {},
-  "mask_token": "[MASK]",
-  "max_length": 512,
-  "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]"
-}
+{
+  "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": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

router_config.json

@@ -0,0 +1,20 @@
+{
+    "types": {
+        "query_0_SparseStaticEmbedding": "sentence_transformers.sparse_encoder.models.SparseStaticEmbedding.SparseStaticEmbedding",
+        "document_0_MLMTransformer": "sentence_transformers.sparse_encoder.models.MLMTransformer.MLMTransformer",
+        "document_1_SpladePooling": "sentence_transformers.sparse_encoder.models.SpladePooling.SpladePooling"
+    },
+    "structure": {
+        "query": [
+            "query_0_SparseStaticEmbedding"
+        ],
+        "document": [
+            "document_0_MLMTransformer",
+            "document_1_SpladePooling"
+        ]
+    },
+    "parameters": {
+        "default_route": "query",
+        "allow_empty_key": true
+    }
+}

train_script.py

@@ -8,8 +8,8 @@ from sentence_transformers import (
 from sentence_transformers.sparse_encoder.losses import SpladeLoss, SparseMultipleNegativesRankingLoss
 from sentence_transformers.training_args import BatchSamplers
 from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
-from sentence_transformers.sparse_encoder.models import SpladePooling, MLMTransformer, IDF
-from sentence_transformers.models import Asym
+from sentence_transformers.sparse_encoder.models import SpladePooling, MLMTransformer, SparseStaticEmbedding
+from sentence_transformers.models import Router
 
 import logging
 
@@ -17,20 +17,24 @@ logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:
 
 # 1. Load a model to finetune with 2. (Optional) model card data
 mlm_transformer = MLMTransformer("prajjwal1/bert-tiny", tokenizer_args={"model_max_length": 512})
-splade_pooling = SpladePooling(pooling_strategy="max", word_embedding_dimension=mlm_transformer.get_sentence_embedding_dimension())
+splade_pooling = SpladePooling(
+    pooling_strategy="max", word_embedding_dimension=mlm_transformer.get_sentence_embedding_dimension()
+)
 
-asym = Asym({
-    "query": [IDF(tokenizer=mlm_transformer.tokenizer, frozen=False)],
-    "document": [mlm_transformer, splade_pooling],
-})
+router = Router(
+    {
+        "query": [SparseStaticEmbedding(tokenizer=mlm_transformer.tokenizer, frozen=False)],
+        "document": [mlm_transformer, splade_pooling],
+    }
+)
 
 model = SparseEncoder(
-    modules=[asym],
+    modules=[router],
     model_card_data=SparseEncoderModelCardData(
         language="en",
         license="apache-2.0",
         model_name="Inference-free SPLADE BERT-tiny trained on Natural-Questions tuples",
-    )
+    ),
 )
 
 # 3. Load a dataset to finetune on
@@ -97,4 +101,4 @@ dev_evaluator(model)
 model.save_pretrained(f"models/{run_name}/final")
 
 # 10. (Optional) Push it to the Hugging Face Hub
-model.push_to_hub(run_name)
+model.push_to_hub(run_name)