Tom Aarsen commited on
Commit
88d5ea6
·
1 Parent(s): d1cf1c1

Add default query for 'model.encode_query' usage

Browse files
Files changed (2) hide show
  1. README.md +36 -0
  2. config_sentence_transformers.json +1 -1
README.md CHANGED
@@ -157,6 +157,42 @@ print(embeddings)
157
  '''
158
  ```
159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
160
  ### vllm support
161
  ```
162
  pip install -U vllm==0.8.5
 
157
  '''
158
  ```
159
 
160
+ Or you can use `encode_query` and `encode_document` to automatically add the default prompt for queries (`"Instruct: Given a query, retrieve documents that answer the query \n Query: "`) and documents (`""`), respectively.
161
+
162
+ ```python
163
+ from sentence_transformers import SentenceTransformer
164
+ import torch
165
+
166
+ model = SentenceTransformer(
167
+ "KaLM-Embedding/KaLM-embedding-multilingual-mini-instruct-v2.5",
168
+ trust_remote_code=True,
169
+ model_kwargs={
170
+ "torch_dtype": torch.bfloat16,
171
+ "attn_implementation": "flash_attention_2", # Optional
172
+ },
173
+ )
174
+ model.max_seq_length = 512
175
+
176
+ queries = [
177
+ "What is the capital of China?",
178
+ "Explain gravity",
179
+ ]
180
+ documents = [
181
+ "The capital of China is Beijing.",
182
+ "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
183
+ ]
184
+
185
+ query_embeddings = model.encode_query(queries)
186
+ document_embeddings = model.encode_document(documents)
187
+
188
+ similarities = model.similarity(query_embeddings, document_embeddings)
189
+ print(similarities)
190
+ '''
191
+ tensor([[0.9034, 0.2563],
192
+ [0.3153, 0.7396]])
193
+ '''
194
+ ```
195
+
196
  ### vllm support
197
  ```
198
  pip install -U vllm==0.8.5
config_sentence_transformers.json CHANGED
@@ -5,7 +5,7 @@
5
  "pytorch": "2.1.2+cu121"
6
  },
7
  "prompts": {
8
- "query": "",
9
  "document": ""
10
  },
11
  "default_prompt_name": null,
 
5
  "pytorch": "2.1.2+cu121"
6
  },
7
  "prompts": {
8
+ "query": "Instruct: Given a query, retrieve documents that answer the query \n Query: ",
9
  "document": ""
10
  },
11
  "default_prompt_name": null,