Build in query prompt to Sentence Transformers Config (#19)

- Build in query prompt to Sentence Transformers Config (b192313ec7397fa525650967f22d04201dfa0788)


Co-authored-by: Jonathan Wang <[email protected]>

README.md

@@ -2621,7 +2621,7 @@ Here, we provide several ways to produce sentence embeddings. Please note that y
 
 ## Quickstart
 
-Here, we provide several ways to produce sentence embeddings. Please note that you have to provide the prompt `Represent this sentence for searching relevant passages:` for query if you want to use it for retrieval. Besides that you don't need any prompt.
+Here, we provide several ways to produce sentence embeddings. Please note that you have to provide the prompt `Represent this sentence for searching relevant passages: ` for query if you want to use it for retrieval. Besides that you don't need any prompt.
 
 ### sentence-transformers
 
@@ -2640,11 +2640,11 @@ dimensions = 512
 # 2. load model
 model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", truncate_dim=dimensions)
 
-# For retrieval you need to pass this prompt.
-query = 'Represent this sentence for searching relevant passages: A man is eating a piece of bread'
+# The prompt used for query retrieval tasks:
+# query_prompt = 'Represent this sentence for searching relevant passages: '
 
+query = "A man is eating a piece of bread"
 docs = [
-    query,
     "A man is eating food.",
     "A man is eating pasta.",
     "The girl is carrying a baby.",
@@ -2652,19 +2652,24 @@ docs = [
 ]
 
 # 2. Encode
-embeddings = model.encode(docs)
+query_embedding = model.encode(query, prompt_name="query")
+# Equivalent Alternatives:
+# query_embedding = model.encode(query_prompt + query)
+# query_embedding = model.encode(query, prompt=query_prompt)
+
+docs_embeddings = model.encode(docs)
 
 # Optional: Quantize the embeddings
-binary_embeddings = quantize_embeddings(embeddings, precision="ubinary")
+binary_query_embedding = quantize_embeddings(query_embedding, precision="ubinary")
+binary_docs_embeddings = quantize_embeddings(docs_embeddings, precision="ubinary")
 
-similarities = cos_sim(embeddings[0], embeddings[1:])
+similarities = cos_sim(query_embedding, docs_embeddings)
 print('similarities:', similarities)
 
 
-```
 ### Transformers
 
-```python
+
 from typing import Dict
 
 import torch
@@ -2712,18 +2717,18 @@ embeddings = pooling(outputs, inputs, 'cls')
 
 similarities = cos_sim(embeddings[0], embeddings[1:])
 print('similarities:', similarities)
-```
+
 
 ### Transformers.js
 
 If you haven't already, you can install the [Transformers.js](https://huggingface.co/docs/transformers.js) JavaScript library from [NPM](https://www.npmjs.com/package/@xenova/transformers) using:
-```bash
+
 npm i @xenova/transformers
-```
+
 
 You can then use the model to compute embeddings like this:
 
-```js
+
 import { pipeline, cos_sim } from '@xenova/transformers';
 
 // Create a feature extraction pipeline
@@ -2745,13 +2750,13 @@ const output = await extractor(docs, { pooling: 'cls' });
 const [source_embeddings, ...document_embeddings ] = output.tolist();
 const similarities = document_embeddings.map(x => cos_sim(source_embeddings, x));
 console.log(similarities); // [0.7919578577247139, 0.6369278664248345, 0.16512018371357193, 0.3620778366720027]
-```
+
 
 ### Using API
 
 You can use the model via our API as follows:
 
-```python
+
 from mixedbread_ai.client import MixedbreadAI, EncodingFormat
 from sklearn.metrics.pairwise import cosine_similarity
 import os
@@ -2773,10 +2778,9 @@ res = mxbai.embeddings(
 
 encoded_embeddings = res.data[0].embedding
 print(res.dimensions, encoded_embeddings.ubinary, encoded_embeddings.float_, encoded_embeddings.int_8)
-```
 
-The API comes with native int8 and binary quantization support! Check out the [docs](https://mixedbread.ai/docs) for more information.
 
+The API comes with native int8 and binary quantization support! Check out the [docs](https://mixedbread.ai/docs) for more information.
 ## Evaluation
 As of March 2024, our model archives SOTA performance for Bert-large sized models on the [MTEB](https://huggingface.co/spaces/mteb/leaderboard). It ourperforms commercial models like OpenAIs text-embedding-3-large and matches the performance of model 20x it's size like the [echo-mistral-7b](https://huggingface.co/jspringer/echo-mistral-7b-instruct-lasttoken). Our model was trained with no overlap of the MTEB data, which indicates that our model generalizes well across several domains, tasks and text length. We know there are some limitations with this model, which will be fixed in v2. 
 

config_sentence_transformers.json

@@ -4,6 +4,9 @@
     "transformers": "4.37.0",
     "pytorch": "2.1.0+cu121"
   },
-  "prompts": {},
-  "default_prompt_name": null
+  "prompts": {
+    "query": "Represent this sentence for searching relevant passages: ",
+    "text": ""
+  },
+  "default_prompt_name": "text"
 }