Adding Onnx Support

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+onnx/model.onnx_data filter=lfs diff=lfs merge=lfs -text

README.md

@@ -364,6 +364,52 @@ with torch.no_grad():
     print(scores)
 ```
 
+#### Usage reranker with the ONNX files
+
+```python
+from optimum.onnxruntime import ORTModelForSequenceClassification  # type: ignore
+
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
+model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
+model_ort = ORTModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base', file_name="onnx/model.onnx")
+
+# Sentences we want sentence embeddings for
+pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
+
+# Tokenize sentences
+encoded_input = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')
+
+scores_ort = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()
+# Compute token embeddings
+with torch.inference_mode():
+    scores = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()
+
+# scores and scores_ort are identical
+```
+#### Usage reranker with infinity
+
+Its also possible to deploy the onnx/torch files with the [infinity_emb](https://github.com/michaelfeil/infinity) pip package.
+```python
+import asyncio
+from infinity_emb import AsyncEmbeddingEngine, EngineArgs
+
+query='what is a panda?'
+docs = ['The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear', "Paris is in France."]
+
+engine = AsyncEmbeddingEngine.from_args(
+    EngineArgs(model_name_or_path = "BAAI/bge-reranker-base", device="cpu", engine="torch" # or engine="optimum" for onnx
+))
+
+async def main(): 
+    async with engine:
+        ranking, usage = await engine.rerank(query=query, docs=docs)
+        print(list(zip(ranking, docs)))
+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:0528834a83cbe4b37fd20a887dd8f3dbdc6d924ffaaab51278d0a05364410117
+size 618476

onnx/model.onnx_data

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