Bge M3 Onnx

🚀 [BGE - M3 ONNX 模型]

[BGE - M3](https://huggingface.co/BAAI/bge - m3) 藉助 HF Optimum 轉換為 ONNX 權重，以實現兼容性，例如與 ONNX Runtime 兼容。此 ONNX 模型可一次性輸出密集、稀疏和 ColBERT 嵌入表示，輸出為按上述表示順序排列的 numpy 數組列表。

🚀 快速開始

本部分將介紹如何使用轉換後的 ONNX 模型進行嵌入計算，以及如何導出 ONNX 權重。

✨ 主要特性

• 多類型嵌入輸出：該 ONNX 模型能夠同時輸出密集、稀疏和 ColBERT 嵌入表示。
• 規範化輸出：默認情況下，密集和 ColBERT 嵌入會進行規範化處理。若需要非規範化輸出，可修改 bgem3_model.py 中的代碼，並使用 export_onnx.py 腳本重新導出 ONNX 模型。
• 圖優化：模型應用了 “O2” 級別的圖優化，若需要不同級別的優化或不進行優化，可使用 export_onnx.py 腳本並傳入適當的優化參數。

📦 安裝指南

安裝 ONNX Runtime

若尚未安裝 ONNX Runtime Python 庫，可使用以下命令進行安裝：

pip install onnxruntime==1.17.0

安裝 HF Transformers

對於分詞操作，可使用 HF Transformers，通過以下命令進行安裝：

pip install transformers==4.37.2

克隆倉庫

使用 [Git LFS](https://git - lfs.com/) 克隆此倉庫以獲取 ONNX 模型文件。

💻 使用示例

基礎用法

以下代碼展示瞭如何使用該 ONNX 模型計算嵌入：

import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3")
ort_session = ort.InferenceSession("model.onnx")

inputs = tokenizer("BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", padding="longest", return_tensors="np")
inputs_onnx = {k: ort.OrtValue.ortvalue_from_numpy(v) for k, v in inputs.items()}

outputs = ort_session.run(None, inputs_onnx)

高級用法

若要從 ONNX 模型獲取稀疏表示的相同輸出，可使用以下代碼：

from collections import defaultdict


def process_token_weights(token_weights: np.ndarray, input_ids: list):
    # conver to dict
    result = defaultdict(int)
    unused_tokens = set(
        [
            tokenizer.cls_token_id,
            tokenizer.eos_token_id,
            tokenizer.pad_token_id,
            tokenizer.unk_token_id,
        ]
    )
    for w, idx in zip(token_weights, input_ids):
        if idx not in unused_tokens and w > 0:
            idx = str(idx)
            # w = int(w)
            if w > result[idx]:
                result[idx] = w
    return result


token_weights = outputs[1].squeeze(-1)
lexical_weights = list(
    map(process_token_weights, token_weights, inputs["input_ids"].tolist())
)

📚 詳細文檔

導出 ONNX 權重

可使用提供的自定義 BGE - M3 PyTorch 模型文件 bgem3_model.py 和 export_onnx.py ONNX 權重導出腳本（藉助 HF Optimum）來導出 ONNX 權重。

若有需要，可修改 bgem3_model.py 模型配置，例如移除嵌入規範化或不輸出所有三種嵌入表示。若修改了輸出表示的數量，則還需修改 export_onnx.py 中的 ONNX 輸出配置 BGEM3OnnxConfig。

安裝依賴

首先，使用以下命令安裝所需的 Python 依賴：

pip install -r requirements.txt

導出 ONNX 權重

然後，使用以下命令導出 ONNX 權重：

python export_onnx.py --output . --opset 17 --device cpu --optimize O2

可在 此處 瞭解更多關於可選優化級別的信息。

📄 許可證

本項目採用 MIT 許可證。