bge-m3

This version of bge-m3 model has been converted to run on the Axera NPU using w8a16 quantization.

This model has been optimized with the following LoRA:

Compatible with Pulsar2 version: 6.0-dirty

Convert tools links:

For those who are interested in model conversion, you can try to export axmodel through

• bge-m3, original model repository

• Model Convert, which you can get the detail of guide

• Pulsar2 Link, How to Convert ONNX to axmodel

Support Platform

• AX650
  • AX650N DEMO Board
  • M4N-Dock(爱芯派Pro)
  • AI Pyramid
  • M.2 Accelerator card

How to use

Download all files from this repository to the device

(py312) root@ax650:~/bge-m3# tree
.
|-- README.md
|-- model
|   |-- bge-m3_full_b1_l512.onnx
|   |-- bge-m3_full_b1_l512.onnx.data
|   `-- bge-m3_u16_npu3.axmodel
|-- python
|   |-- axmodel_infer.py
|   |-- compare_bge_m3_onnx_axmodel.py
|   `-- onnx_infer.py
|-- quant
|   |-- bge-m3.json
|   `-- calib_tokens_m3.tar.gz
`-- requirements.txt

Inference

Inference with AX650 Host, such as M4N-Dock(爱芯派Pro)

run with python3 axmodel_infer.py

root@ax650:~/bge# python3 axmodel_infer_bgem3.py
[INFO] Available providers:  ['AxEngineExecutionProvider', 'AXCLRTExecutionProvider']
[INFO] Using provider: AxEngineExecutionProvider
[INFO] Chip type: ChipType.MC50
[INFO] VNPU type: VNPUType.DISABLED
[INFO] Engine version: 2.12.0s
[INFO] Model type: 2 (triple core)
[INFO] Compiler version: 6.0-dirty 71f24c74-dirty
[[0.60657114 0.3339174 ]
 [0.34422207 0.66171443]]
0.76533616
0.4494059
{'colbert': [0.7653361558914185, 0.4494059085845947, 0.4402206242084503, 0.7766788601875305], 'sparse': [0.18121449201226758, 0.007629240866828535, 0.0, 0.17698916647350543], 'dense': [0.6065711975097656, 0.33391737937927246, 0.3442220687866211, 0.661714494228363], 'sparse+dense': [0.4647856290105996, 0.22515466654179112, 0.22948137919108072, 0.5001393849767437], 'colbert+sparse+dense': [0.5850058397629272, 0.3148551633589126, 0.3137770771980286, 0.6107551750610585]}

Compare results for fp32/quantized model

run with python3 compare_bge_m3_onnx_axmodel.py to compare the results of the onnx model and the quantized axmodel

root@ax650:~/bge# python3 compare_bge_m3_onnx_axmodel.py
[INFO] Available providers:  ['AxEngineExecutionProvider', 'AXCLRTExecutionProvider']
query count: 50
passage count: 50
embedding text count: 100
[INFO] Using provider: AxEngineExecutionProvider
[INFO] Chip type: ChipType.MC50
[INFO] VNPU type: VNPUType.DISABLED
[INFO] Engine version: 2.12.0s
[INFO] Model type: 2 (triple core)
[INFO] Compiler version: 6.0-dirty 71f24c74-dirty

...
[98] passage[48] The model output for text embedding may include dense vectors, sparse token weights, and token-level vectors.
dense_vecs: shape=(1, 1024), max_abs=0.009957, mean_abs=0.002286
dense cosine=0.995690
lexical_weights: onnx_keys=19, ax_keys=19, common=19, max_abs=0.010975, mean_abs=0.004431
colbert_vecs: shape=(31, 1024), max_abs=0.102679, mean_abs=0.003649

[99] passage[49] This passage describes how vector search is used in a BGE-M3 text embedding workflow.
dense_vecs: shape=(1, 1024), max_abs=0.010877, mean_abs=0.002496
dense cosine=0.994958
lexical_weights: onnx_keys=18, ax_keys=18, common=18, max_abs=0.014455, mean_abs=0.005134
colbert_vecs: shape=(25, 1024), max_abs=0.024213, mean_abs=0.003087

========== Embedding Summary ==========
dense_max_abs: mean=0.007735, max=0.015179, p95=0.010877
dense_mean_abs: mean=0.001762, max=0.002566, p95=0.002496
dense_cosine: mean=0.997326, min=0.994669, p95=0.999105
lexical_max_abs: mean=0.010224, max=0.048042, p95=0.018426
lexical_mean_abs: mean=0.004123, max=0.007477, p95=0.006482
colbert_max_abs: mean=0.022130, max=0.102679, p95=0.095120
colbert_mean_abs: mean=0.001964, max=0.003809, p95=0.003458
...

Pair 48: q='What is BGE M3?', p='The model output for text embedding may include de'
  dense: onnx=0.239425, ax=0.235858, abs_diff=0.003567
  sparse: onnx=0.000000, ax=0.000000, abs_diff=0.000000
  colbert: onnx=0.323061, ax=0.309670, abs_diff=0.013391
  sparse+dense: onnx=0.159617, ax=0.157239, abs_diff=0.002378
  colbert+sparse+dense: onnx=0.224994, ax=0.218211, abs_diff=0.006783

Pair 49: q='What is BGE M3?', p='This passage describes how vector search is used i'
  dense: onnx=0.512933, ax=0.491597, abs_diff=0.021336
  sparse: onnx=0.118533, ax=0.108602, abs_diff=0.009931
  colbert: onnx=0.701150, ax=0.690015, abs_diff=0.011135
  sparse+dense: onnx=0.381466, ax=0.363932, abs_diff=0.017534
  colbert+sparse+dense: onnx=0.509340, ax=0.494365, abs_diff=0.014975

========== Pair Score Summary ==========
dense: mean_abs=0.005527, max_abs=0.021336, p95_abs=0.019722
sparse: mean_abs=0.001502, max_abs=0.009931, p95_abs=0.009504
colbert: mean_abs=0.008519, max_abs=0.016579, p95_abs=0.014835
sparse+dense: mean_abs=0.004185, max_abs=0.017534, p95_abs=0.016316
colbert+sparse+dense: mean_abs=0.005623, max_abs=0.014975, p95_abs=0.014281

========== Final Conclusion ==========
accuracy level: GOOD
dense cosine min: 0.994669
dense mean abs diff: 0.001762
fusion mean abs diff: 0.005623
fusion max abs diff: 0.014975
量化后整体精度掉点较小，dense 相似度和最终融合分数都比较稳定。