add onnx files and readme for usage with onnx (#14)
Browse files- add onnx files (af0ab12eac918e67a5d9b4ce7c22641d2f2e907f)
- Upload model.onnx (1e1223f74591f1fcc92203ca252df268f3cd614c)
- Update README.md (3777b8d4e6a51aad2df2198421eebddf5a0222f7)
- Update README.md (4386d50bb4c9f2128d8a517e83d18c73050b379e)
Co-authored-by: Michael <[email protected]>
- README.md +46 -0
- onnx/model.onnx +3 -0
README.md
CHANGED
|
@@ -365,6 +365,52 @@ with torch.no_grad():
|
|
| 365 |
print(scores)
|
| 366 |
```
|
| 367 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 368 |
## Evaluation
|
| 369 |
|
| 370 |
`baai-general-embedding` models achieve **state-of-the-art performance on both MTEB and C-MTEB leaderboard!**
|
|
|
|
| 365 |
print(scores)
|
| 366 |
```
|
| 367 |
|
| 368 |
+
#### Usage reranker with the ONNX files
|
| 369 |
+
|
| 370 |
+
```python
|
| 371 |
+
from optimum.onnxruntime import ORTModelForSequenceClassification # type: ignore
|
| 372 |
+
|
| 373 |
+
import torch
|
| 374 |
+
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
| 375 |
+
|
| 376 |
+
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
|
| 377 |
+
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
|
| 378 |
+
model_ort = ORTModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base', file_name="onnx/model.onnx")
|
| 379 |
+
|
| 380 |
+
# Sentences we want sentence embeddings for
|
| 381 |
+
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.']]
|
| 382 |
+
|
| 383 |
+
# Tokenize sentences
|
| 384 |
+
encoded_input = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')
|
| 385 |
+
|
| 386 |
+
scores_ort = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()
|
| 387 |
+
# Compute token embeddings
|
| 388 |
+
with torch.inference_mode():
|
| 389 |
+
scores = model_ort(**encoded_input, return_dict=True).logits.view(-1, ).float()
|
| 390 |
+
|
| 391 |
+
# scores and scores_ort are identical
|
| 392 |
+
```
|
| 393 |
+
#### Usage reranker with infinity
|
| 394 |
+
|
| 395 |
+
Its also possible to deploy the onnx/torch files with the [infinity_emb](https://github.com/michaelfeil/infinity) pip package.
|
| 396 |
+
```python
|
| 397 |
+
import asyncio
|
| 398 |
+
from infinity_emb import AsyncEmbeddingEngine, EngineArgs
|
| 399 |
+
|
| 400 |
+
query='what is a panda?'
|
| 401 |
+
docs = ['The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear', "Paris is in France."]
|
| 402 |
+
|
| 403 |
+
engine = AsyncEmbeddingEngine.from_args(
|
| 404 |
+
EngineArgs(model_name_or_path = "BAAI/bge-reranker-base", device="cpu", engine="torch" # or engine="optimum" for onnx
|
| 405 |
+
))
|
| 406 |
+
|
| 407 |
+
async def main():
|
| 408 |
+
async with engine:
|
| 409 |
+
ranking, usage = await engine.rerank(query=query, docs=docs)
|
| 410 |
+
print(list(zip(ranking, docs)))
|
| 411 |
+
asyncio.run(main())
|
| 412 |
+
```
|
| 413 |
+
|
| 414 |
## Evaluation
|
| 415 |
|
| 416 |
`baai-general-embedding` models achieve **state-of-the-art performance on both MTEB and C-MTEB leaderboard!**
|
onnx/model.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:15b9a8c3da82eddf263df571281166e00e9308fe19d077084b642ebfcaf06d2b
|
| 3 |
+
size 1112459588
|