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