Sentence Transformers integration

1_Pooling/config.json

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

README.md

@@ -2934,6 +2934,38 @@ Based on the [intfloat/e5-large-unsupervised](https://huggingface.co/intfloat/e5
 
 ## Usage
 
+### Using Sentence Transformers
+
+You can use the sentence-transformers package to use an snowflake-arctic-embed model, as shown below. 
+
+```python
+from sentence_transformers import SentenceTransformer
+
+model = SentenceTransformer("Snowflake/snowflake-arctic-embed-xs")
+
+queries = ['what is snowflake?', 'Where can I get the best tacos?']
+documents = ['The Data Cloud!', 'Mexico City of Course!']
+
+query_embeddings = model.encode(queries, prompt_name="query")
+document_embeddings = model.encode(documents)
+
+scores = query_embeddings @ document_embeddings.T
+for query, query_scores in zip(queries, scores):
+    doc_score_pairs = list(zip(documents, query_scores))
+    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
+    # Output passages & scores
+    print("Query:", query)
+    for document, score in doc_score_pairs:
+        print(score, document)
+```
+```
+Query: what is snowflake?
+0.57515126 The Data Cloud!
+0.45798576 Mexico City of Course!
+Query: Where can I get the best tacos?
+0.5636022 Mexico City of Course!
+0.5044898 The Data Cloud!
+```
 
 ### Using Huggingface transformers
 

config_sentence_transformers.json

@@ -0,0 +1,11 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.7.0.dev0",
+    "transformers": "4.39.3",
+    "pytorch": "2.1.0+cu121"
+  },
+  "prompts": {
+    "query": "Represent this sentence for searching relevant passages: "
+  },
+  "default_prompt_name": null
+}

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": 512,
+  "do_lower_case": false
+}