Add support for Sentence Transformer

1_SpladePooling/config.json

@@ -0,0 +1,5 @@
+{
+    "pooling_strategy": "max",
+    "activation_function": "relu",
+    "word_embedding_dimension": null
+}

README.md

@@ -10,6 +10,12 @@ tags:
 - query-expansion
 - document-expansion
 - bag-of-words
+- sentence-transformers
+- sparse-encoder
+- sparse
+- splade
+pipeline_tag: feature-extraction
+library_name: sentence-transformers
 ---
 
 # opensearch-neural-sparse-encoding-v2-distill
@@ -37,6 +43,91 @@ The training datasets includes MS MARCO, eli5_question_answer, squad_pairs, Wiki
 
 OpenSearch neural sparse feature supports learned sparse retrieval with lucene inverted index. Link: https://opensearch.org/docs/latest/query-dsl/specialized/neural-sparse/. The indexing and search can be performed with OpenSearch high-level API.
 
+## Usage (Sentence Transformers)
+
+First install the Sentence Transformers library:
+
+```bash
+pip install -U sentence-transformers
+```
+
+Then you can load this model and run inference.
+
+```python
+from sentence_transformers.sparse_encoder import SparseEncoder
+
+# Download from the 🤗 Hub
+model = SparseEncoder("opensearch-project/opensearch-neural-sparse-encoding-v2-distill")
+
+query = "What's the weather in ny now?"
+document = "Currently New York is rainy."
+
+query_embed = model.encode_query(query)
+document_embed = model.encode_document(document)
+
+sim = model.similarity(query_embed, document_embed)
+print(f"Similarity: {sim}")
+# Similarity: tensor([[38.6113]])
+
+decoded_query = model.decode(query_embed)
+decoded_document = model.decode(document_embed)
+
+for i in range(len(decoded_query)):
+    query_token, query_score = decoded_query[i]
+    doc_score = next((score for token, score in decoded_document if token == query_token), 0)
+    if doc_score != 0:
+        print(f"Token: {query_token}, Query score: {query_score:.4f}, Document score: {doc_score:.4f}")
+
+# Token: york, Query score: 2.7273, Document score: 2.9088
+# Token: now, Query score: 2.5734, Document score: 0.9208
+# Token: ny, Query score: 2.3895, Document score: 1.7237
+# Token: weather, Query score: 2.2184, Document score: 1.2368
+# Token: current, Query score: 1.8693, Document score: 1.4146
+# Token: today, Query score: 1.5888, Document score: 0.7450
+# Token: sunny, Query score: 1.4704, Document score: 0.9247
+# Token: nyc, Query score: 1.4374, Document score: 1.9737
+# Token: currently, Query score: 1.4347, Document score: 1.6019
+# Token: climate, Query score: 1.1605, Document score: 0.9794
+# Token: upstate, Query score: 1.0944, Document score: 0.7141
+# Token: forecast, Query score: 1.0471, Document score: 0.5519
+# Token: verve, Query score: 0.9268, Document score: 0.6692
+# Token: huh, Query score: 0.9126, Document score: 0.4486
+# Token: greene, Query score: 0.8960, Document score: 0.7706
+# Token: picturesque, Query score: 0.8779, Document score: 0.7120
+# Token: pleasantly, Query score: 0.8471, Document score: 0.4183
+# Token: windy, Query score: 0.8079, Document score: 0.2140
+# Token: favorable, Query score: 0.7537, Document score: 0.4925
+# Token: rain, Query score: 0.7519, Document score: 2.1456
+# Token: skies, Query score: 0.7277, Document score: 0.3818
+# Token: lena, Query score: 0.6995, Document score: 0.8593
+# Token: sunshine, Query score: 0.6895, Document score: 0.2410
+# Token: johnny, Query score: 0.6621, Document score: 0.3016
+# Token: skyline, Query score: 0.6604, Document score: 0.1933
+# Token: sasha, Query score: 0.6117, Document score: 0.2197
+# Token: vibe, Query score: 0.5962, Document score: 0.0414
+# Token: hardly, Query score: 0.5381, Document score: 0.7560
+# Token: prevailing, Query score: 0.4583, Document score: 0.4243
+# Token: unpredictable, Query score: 0.4539, Document score: 0.5073
+# Token: presently, Query score: 0.4350, Document score: 0.8463
+# Token: hail, Query score: 0.3674, Document score: 0.2496
+# Token: shivered, Query score: 0.3324, Document score: 0.5506
+# Token: wind, Query score: 0.3281, Document score: 0.1964
+# Token: rudy, Query score: 0.3052, Document score: 0.5785
+# Token: looming, Query score: 0.2797, Document score: 0.0357
+# Token: atmospheric, Query score: 0.2712, Document score: 0.0870
+# Token: vicky, Query score: 0.2471, Document score: 0.3490
+# Token: sandy, Query score: 0.2247, Document score: 0.2383
+# Token: crowded, Query score: 0.2154, Document score: 0.5737
+# Token: chilly, Query score: 0.1723, Document score: 0.1857
+# Token: blizzard, Query score: 0.1700, Document score: 0.4110
+# Token: ##cken, Query score: 0.1183, Document score: 0.0613
+# Token: unrest, Query score: 0.0923, Document score: 0.6363
+# Token: russ, Query score: 0.0624, Document score: 0.2127
+# Token: blackout, Query score: 0.0558, Document score: 0.5542
+# Token: kahn, Query score: 0.0549, Document score: 0.1589
+# Token: 2020, Query score: 0.0160, Document score: 0.0566
+# Token: nighttime, Query score: 0.0125, Document score: 0.3753
+```
 
 ## Usage (HuggingFace)
 This model is supposed to run inside OpenSearch cluster. But you can also use it outside the cluster, with HuggingFace models API. 

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "SparseEncoder",
+  "__version__": {
+    "sentence_transformers": "5.0.0",
+    "transformers": "4.50.3",
+    "pytorch": "2.6.0+cu124"
+  },
+  "prompts": {
+    "query": "",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "dot"
+}

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.sparse_encoder.models.MLMTransformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_SpladePooling",
+    "type": "sentence_transformers.sparse_encoder.models.SpladePooling"
+  }
+]

sentence_bert_config.json

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