Shitao michaelfeil commited on
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]>

Files changed (4) hide show
  1. .gitattributes +1 -0
  2. README.md +46 -0
  3. onnx/model.onnx +3 -0
  4. 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