🚀 [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):
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)
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 许可证。