michaelfeil commited on
Commit
2176106
1 Parent(s): c5ac6c3

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +49 -0
README.md CHANGED
@@ -2907,6 +2907,55 @@ with torch.no_grad():
2907
  print(scores)
2908
  ```
2909
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2910
  ## Evaluation
2911
 
2912
  `baai-general-embedding` models achieve **state-of-the-art performance on both MTEB and C-MTEB leaderboard!**
 
2907
  print(scores)
2908
  ```
2909
 
2910
+ #### Usage of the ONNX files
2911
+
2912
+ ```python
2913
+ from optimum.onnxruntime import ORTModelForFeatureExtraction # type: ignore
2914
+
2915
+ import torch
2916
+ from transformers import AutoModel, AutoTokenizer
2917
+
2918
+ tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5')
2919
+ model = AutoModel.from_pretrained('BAAI/bge-small-en-v1.5')
2920
+ model_ort = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx")
2921
+
2922
+ # Sentences we want sentence embeddings for
2923
+ sentences = ["样例数据-1", "样例数据-2"]
2924
+
2925
+ # Tokenize sentences
2926
+ encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
2927
+ # for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
2928
+ # encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
2929
+
2930
+ model_output_ort = model_ort(**encoded_input)
2931
+ # Compute token embeddings
2932
+ with torch.no_grad():
2933
+ model_output = model(**encoded_input)
2934
+
2935
+ # model_output and model_output_ort are identical
2936
+
2937
+ ```
2938
+
2939
+ #### Usage via infinity
2940
+ Its also possible to deploy the onnx files with the [infinity_emb](https://github.com/michaelfeil/infinity) pip package.
2941
+ Recommended is `device="cuda", engine="torch"` with flash attention on gpu, and `device="cpu", engine="optimum"` for onnx inference.
2942
+
2943
+ ```python
2944
+ import asyncio
2945
+ from infinity_emb import AsyncEmbeddingEngine, EngineArgs
2946
+
2947
+ sentences = ["Embed this is sentence via Infinity.", "Paris is in France."]
2948
+ engine = AsyncEmbeddingEngine.from_args(
2949
+ EngineArgs(model_name_or_path = "BAAI/bge-small-en-v1.5", device="cpu", engine="optimum" # or engine="torch"
2950
+ ))
2951
+
2952
+ async def main():
2953
+ async with engine:
2954
+ embeddings, usage = await engine.embed(sentences=sentences)
2955
+ asyncio.run(main())
2956
+ ```
2957
+
2958
+
2959
  ## Evaluation
2960
 
2961
  `baai-general-embedding` models achieve **state-of-the-art performance on both MTEB and C-MTEB leaderboard!**