Y
gcw_hhrLdMXW/bge-base-en-v1.5
模型介绍
文件和版本
Pull Requests
讨论
分析

BAAI/bge-base-en-v1.5 昇腾 NPU 部署

模型简介

BGE (BAAI General Embedding) base-en-v1.5 是北京智源人工智能研究院（BAAI）发布的英文通用文本嵌入模型，基于 BERT-base 架构（12 层 Transformer，768 hidden，约 1.09 亿参数）。它把任意文本映射为 768 维语义向量，广泛用于语义检索、文本相似度计算、RAG 向量库构建等场景。

  • 架构：BertModel（BertModel）
  • 输入：英文文本
  • 输出：768 维归一化嵌入向量
  • 权重：model.safetensors（418MB）
  • 官方仓库：https://huggingface.co/BAAI/bge-base-en-v1.5

环境依赖清单

组件版本说明
昇腾 NPUAscend910_9362 ×2（64GB HBM）推理硬件
CANN8.5.1昇腾计算框架
torch2.9.0+cpuPyTorch
torch_npu2.9.0.post1昇腾 NPU 插件
transformers4.57.6模型加载
numpy-数值计算

安装依赖：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

分步推理操作流程

1. 环境检查

npu-smi info          # 确认 NPU 健康与显存
python3 -c "import torch, torch_npu; print(torch.npu.device_count())"

2. 运行推理脚本

cd bge-base-en-v1.5-npu-deployment
ASCEND_RT_VISIBLE_DEVICES=0 python3 inference.py

3. 自定义文本

python3 inference.py --texts "The quick brown fox" "jumps over the lazy dog"

4. 自定义设备与基准

python3 inference.py --device npu:0 --warmup 3 --runs 20 --max-length 128

测试用例与输出结果

测试输入（4 句英文）：

  1. "The weather is lovely today."
  2. "It's so sunny outside!"
  3. "He drove to the grocery store."
  4. "The cat is sleeping on the couch."

输出结果（昇腾 NPU，Ascend910_9362）：

嵌入向量（768 维，归一化 norm=1.0）：

[0] [0.0113, -0.0305, 0.0759, -0.0257, 0.0541, 0.0219, 0.0286, 0.0302, ...]
[1] [0.0303, -0.0274, 0.0601, -0.0445, 0.0683, 0.0082, 0.0112, 0.0149, ...]
[2] [-0.0555, 0.0101, 0.0062, 0.0426, 0.0599, 0.0465, 0.0107, 0.0574, ...]
[3] [-0.0378, -0.0528, 0.0474, -0.0233, 0.1046, -0.0031, 0.0161, 0.0240, ...]

余弦相似度矩阵（语义正确：天气相关句相似度高）：

        [0]     [1]     [2]     [3]
[0]    1.0000  0.8177  0.4997  0.3396
[1]    0.8177  1.0000  0.4803  0.3182
[2]    0.4997  0.4803  1.0000  0.3613
[3]    0.3396  0.3182  0.3613  1.0000

性能基准（warmup=3, runs=20）：

平均耗时: 10.50 ms / 4 句
吞吐量: 380.9 句/秒

验收结论：模型在昇腾 NPU 上完整跑通，嵌入向量归一化正确，相似度矩阵语义合理（天气句对相似度 0.82 显著高于无关句），推理性能 380 句/秒。