onnx support

README.md

@@ -2907,6 +2907,55 @@ with torch.no_grad():
     print(scores)
 ```
 
+#### Usage of the ONNX files
+
+```python
+from optimum.onnxruntime import ORTModelForFeatureExtraction  # type: ignore
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5')
+model = AutoModel.from_pretrained('BAAI/bge-small-en-v1.5')
+model_ort = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx")
+
+# Sentences we want sentence embeddings for
+sentences = ["样例数据-1", "样例数据-2"]
+
+# Tokenize sentences
+encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+# for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
+# encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
+
+model_output_ort = model_ort(**encoded_input)
+# Compute token embeddings
+with torch.no_grad():
+    model_output = model(**encoded_input)
+
+# model_output and model_output_ort are identical
+
+```
+
+#### Usage via infinity
+Its also possible to deploy the onnx files with the [infinity_emb](https://github.com/michaelfeil/infinity) pip package.
+Recommended is `device="cuda", engine="torch"` with flash attention on gpu, and `device="cpu", engine="optimum"` for onnx inference.
+
+```python
+import asyncio
+from infinity_emb import AsyncEmbeddingEngine, EngineArgs
+
+sentences = ["Embed this is sentence via Infinity.", "Paris is in France."]
+engine = AsyncEmbeddingEngine.from_args(
+    EngineArgs(model_name_or_path = "BAAI/bge-small-en-v1.5", device="cpu", engine="optimum" # or engine="torch"
+))
+
+async def main(): 
+    async with engine:
+        embeddings, usage = await engine.embed(sentences=sentences)
+asyncio.run(main())
+```
+
+
 ## Evaluation  
 
 `baai-general-embedding` models achieve **state-of-the-art performance on both MTEB and C-MTEB leaderboard!**

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:828e1496d7fabb79cfa4dcd84fa38625c0d3d21da474a00f08db0f559940cf35
+size 133093490