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

1. 简介

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

BGE（BAAI General Embedding）是智源研究院（BAAI）推出的通用句嵌入模型系列。bge-base-en-v1.5 是基于 BERT-base（12层，768维）的对比学习嵌入模型，在 MTEB 基准的检索任务上表现优异（NDCG@10 达 58+）。相比 small 版本（384维），base 版本精度更高但推理速度约为其一半，适合对语义质量要求更高的 RAG 和语义搜索场景。

相关获取地址：

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

2. 验证环境

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

3. 部署使用流程

3.1 环境准备

conda create -n BAAI--bge-base-en-v1.5 python=3.11 -y
conda activate BAAI--bge-base-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 chocolate cake?" --device npu
python inference.py --batch_file queries.txt --device npu --batch_size 32

编程接口：

from inference import BGEEncoder
encoder = BGEEncoder(model_path="./BAAI--bge-base-en-v1.5", device="npu")
embeddings = encoder.encode(["What is the capital of France?"])
# embeddings.shape → (1, 768)

4. Smoke 验证

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

预期输出：768 维归一化嵌入向量，无运行时错误。

5. 性能参考

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

BGE-base（768维）在 NPU 上获得 23.6× 加速。768 维嵌入在 NPU 上的计算密度高于 384 维，加速比更大。但绝对吞吐（688 vs 792 sent/s）低于 BGE-small。

6. 精度评测

6.1 评测方法

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

• 余弦相似度：逐句 NPU 嵌入与 CPU 嵌入的方向一致性（>0.99 视为等价）
• 精度误差率：(1 - 平均余弦) × 100%（要求 < 1.0%）
• 元素级绝对误差：768 维向量的最大单维偏差
• MSE：逐元素均方误差
• 语义一致性（Pearson r）：NPU 与 CPU 各自计算的语义相似度矩阵的相关系数

6.2 评测结果

6.3 逐句余弦详情

全部 23 句余弦相似度 ≥ 0.999735。

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

7. 迁移适配说明

7.1 模型结构

• Backbone：BertModel（BERT-base，12 层，768 维）
• Pooling：Mean Pooling（对所有 token 嵌入取 mask 加权平均）
• Normalize：L2 归一化到单位球面
• Tokenizer：BERT WordPiece（vocab.txt）
• 参数量：约 110M（BERT-base 标准规模）

7.2 适配要点

1. AutoModel.from_pretrained() 加载，model.to("npu:0") 迁移
2. Mean Pooling 手动实现：(embeddings * mask).sum(1) / mask.sum(1)
3. 与 BGE-small 完全相同适配代码（仅模型维度不同 768 vs 384）
4. 支持 sentence-transformers 格式加载

7.3 关键代码

from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("bge-base-en-v1.5").to("npu:0")
tokenizer = AutoTokenizer.from_pretrained("bge-base-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-base vs small：base（768维）精度更高但速度约为 small（384维）的 87%。检索场景建议先 small 粗召回 + base 精排的组合。
2. 查询指令前缀：BGE 官方建议搜索查询添加 "Represent this sentence for searching relevant passages: " 前缀以发挥最佳检索效果。本适配脚本未强制添加，可自行在调用前拼接。
3. 首次 NPU 推理：BERT-base 12 层算子编译约 3-5 秒，建议预热后批量推理。
4. sentence-transformers 兼容：模型以 sentence-transformers 格式保存，可直接使用 SentenceTransformer 加载。NPU 适配使用手动 Mean Pooling 以简化依赖。
5. 精度对比 BGE-small：Base 版本由于模型更大（768维 vs 384维），浮点运算量增加导致误差略高（0.0022% vs 0.0001%），仍在安全范围内。