Bge M3 Onnx

🚀 BGE - M3 ONNXモデル

[BGE - M3](https://huggingface.co/BAAI/bge - m3)をHF Optimumを使用してONNXウェイトに変換し、ONNX Runtimeなどと互換性を持たせます。

このONNXモデルは、密ベクトル、疎ベクトル、およびColBERT埋め込み表現を一度に出力します。出力は、前述の表現の順序で並んだnumpy配列のリストです。

⚠️ 重要提示

密ベクトルとColBERT埋め込みは、元のFlagEmbeddingライブラリのデフォルトの動作と同様に正規化されています。正規化されていない出力が必要な場合は、bgem3_model.pyのコードを変更し、export_onnx.pyスクリプトを使用してONNXエクスポートを再実行できます。

このONNXモデルには、"O2"レベルのグラフ最適化が適用されています。最適化レベルの詳細については、こちらを参照してください。異なる最適化レベルまたは最適化なしのONNXモデルが必要な場合は、適切な最適化引数を指定してexport_onnx.pyスクリプトを再実行できます。

🚀 クイックスタート

📦 インストール

必要なライブラリをインストールします。

pip install onnxruntime==1.17.0
pip install transformers==4.37.2

[Git LFS](https://git - lfs.com/)を使用してこのリポジトリをクローンし、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)

高度な使用法

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ファイルと、HF Optimumを利用したexport_onnx.py ONNXウェイトエクスポートスクリプトを使用して、ONNXウェイトをエクスポートできます。必要に応じて、bgem3_model.pyのモデル構成を変更して、たとえば埋め込みの正規化を削除したり、3つの埋め込み表現をすべて出力しないようにできます。出力表現の数を変更した場合は、export_onnx.pyのONNX出力構成BGEM3OnnxConfigも変更する必要があります。

まず、必要なPythonパッケージをインストールします。

pip install -r requirements.txt

次に、以下のコマンドでONNXウェイトをエクスポートします。

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

オプションの最適化レベルの詳細については、こちらを参照してください。

📄 ライセンス

このプロジェクトはMITライセンスの下でライセンスされています。