BAAI/bge-small-en-v1.5 on Ascend NPU

1. 简介

本文档记录 BAAI/bge-small-en-v1.5 BGE 句嵌入模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

BGE（BAAI General Embedding）是智源研究院（BAAI）推出的通用句嵌入模型系列。bge-small-en-v1.5 是轻量版本，基于 BERT（384维），通过对比学习在大量检索数据上训练，专为语义搜索和 RAG（检索增强生成）场景优化。相比 all-MiniLM-L6-v2，BGE 在 MTEB 基准上表现更优（检索任务提升 3-5%）。

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/BAAI/bge-small-en-v1.5

2. 验证环境

• NPU：8 × Ascend 910B3
• 精度对比基准：CPU（x86, PyTorch 2.8.0）

3. 部署使用流程

3.1 环境准备

conda create -n BAAI--bge-small-en-v1.5 python=3.11 -y
conda activate BAAI--bge-small-en-v1.5

pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers sentence-transformers numpy \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

3.2 推理

python inference.py --text "How to bake a cake?" --device npu

from inference import BGEEncoder
encoder = BGEEncoder(model_path="./BAAI--bge-small-en-v1.5", device="npu")
embeddings = encoder.encode(["What is machine learning?"])
# embeddings.shape → (1, 384)

4. Smoke 验证

python inference.py --text "How to bake a cake?" --device npu

预期输出：384 维嵌入向量，无运行时错误。

5. 性能参考

测试条件：23 条多样化测试句子（含中英文、短句、长句、特殊字符），batch_size=32，NPU 预热 1 轮。

BGE-small（384维）在 NPU 上获得 8.2× 加速。相比 BERT-base（768维），384维的轻量设计使推理吞吐提升近 2×。

6. 精度评测

6.1 评测方法

分别在 CPU 和 NPU 上对 23 条多样化测试句子（含英文、中文、法文、标点、数字、短句、长句）进行推理，从以下维度评估 NPU 与 CPU 输出精度：

• 余弦相似度：每条句子的 NPU 嵌入与 CPU 嵌入的余弦相似度，衡量逐句向量方向一致性。余弦 > 0.99 视为等价
• 相对误差（精度误差率）：(1 - 平均余弦相似度) × 100%，要求 < 1.0%
• 元素级绝对误差：NPU 嵌入与 CPU 嵌入间 384 维向量的最大元素偏差，衡量数值精度
• MSE：逐元素均方误差，衡量整体偏差水平
• 语义一致性（Pearson r）：NPU 与 CPU 各自计算的 23×23 语义相似度矩阵之间的相关系数，衡量相对语义结构的保持能力

6.2 评测结果

6.3 逐句余弦相似度

全部 23 句余弦相似度 ≥ 0.999998，语义等价性验证通过。

结论：精度误差率 0.0001%，远低于 1% 阈值。NPU 与 CPU 输出在向量方向、数值精度和语义结构上均完全等价，评测通过。

7. 迁移适配说明

7.1 模型结构

• BertModel（BERT-base 384维轻量版）
• Mean Pooling + L2 Normalize（与 MiniLM 相同）
• sentence-transformers 格式（含 Pooling+Normalize 模块配置）

7.2 适配要点

1. AutoModel.from_pretrained() 加载，model.to("npu:0") 迁移
2. Mean Pooling 手动实现：(embeddings*mask).sum(1)/mask.sum(1)
3. 与 all-MiniLM-L6-v2 完全相同适配代码

7.3 关键代码

from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("bge-small-en-v1.5").to("npu:0")
tokenizer = AutoTokenizer.from_pretrained("bge-small-en-v1.5")
inputs = tokenizer(text, return_tensors="pt"); inputs = {k:v.to("npu:0") for k,v in inputs.items()}
with torch.no_grad():
    emb = model(**inputs).last_hidden_state
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    embedding = torch.nn.functional.normalize((emb*mask).sum(1)/mask.sum(1), p=2, dim=1)

8. 注意事项

1. BGE 检索优化：相比 MiniLM，BGE 在检索/排序任务上更优，特别适合 RAG pipeline 中的 embedding 环节
2. 384维轻量：速度与 MiniLM-L6 相近，适合大规模向量检索
3. 输入需加前缀：BGE 官方建议检索查询加 "Represent this sentence for searching relevant passages: " 前缀（本适配脚本未强制，可自行添加）
4. 首次 NPU 推理：384维轻量模型，算子编译约 2-3 秒