BAAI--llm-embedder on Ascend NPU

1. 简介

本文档记录 BAAI/llm-embedder 在 Ascend 910B3 NPU 环境下的适配与验证结果。llm-embedder 是一个统一的文本嵌入模型，支持检索、分类、聚类和语义搜索等任务。

本仓库提供：

• inference.py：NPU 推理脚本，支持单条/批量文本嵌入生成
• eval.py：精度与性能评测脚本，对比 NPU 与 CPU 输出差异
• log.txt：评测运行日志

相关获取地址：

• 模型权重（HuggingFace）：https://huggingface.co/BAAI/llm-embedder
• 镜像加速：https://hf-mirror.com

参考文档：

• https://github.com/FlagOpen/FlagEmbedding
• https://python.langchain.com/docs/integrations/text_embedding/llm_embedder/

2. 验证环境

• NPU：Ascend 910B3，1 逻辑卡
• 模型架构：BertModel（12 层，768 隐藏维，12 注意力头）
• 最大序列长度：512
• 嵌入维度：768

3. 推理启动

环境准备

# 安装依赖
pip install torch torch_npu transformers sentence-transformers

# 设置 NPU 可见设备
export ASCEND_RT_VISIBLE_DEVICES=0

单条文本推理

python inference.py --input "Your text here"

批量文本推理

# 从文件读取（每行一条）
python inference.py --input-file texts.txt

# 指定模型路径
python inference.py --model-path /path/to/BAAI--llm-embedder --input "Your text"

CPU 推理（参考基准）

python inference.py --device cpu --input "Your text"

输出格式

{
  "model": "BAAI--llm-embedder",
  "device": "npu:0",
  "num_texts": 1,
  "dimension": 768,
  "pooling": "cls",
  "normalize": true,
  "inference_time_seconds": 0.0021,
  "texts_per_second": 476.19,
  "embeddings": [[0.0012, -0.0034, ...]]
}

4. Smoke 验证

python inference.py --input "Hello, world!"

预期输出：

• 返回 JSON 格式嵌入向量
• dimension 为 768
• 推理时间在毫秒级别

5. 性能参考

测试条件：10 条文本，batch_size=32，max_length=512。

6. 精度评测

精度评测采用余弦距离（Cosine Distance）作为主要指标，对比 NPU 与 CPU 的输出嵌入向量。

结论：NPU 与 CPU 输出的嵌入向量余弦距离仅 0.000751%，远低于 1% 的阈值，精度通过验证。

7. 注意事项

1. 嵌入归一化：默认对输出向量进行 L2 归一化处理，可通过 --no-normalize 关闭。
2. 池化策略：默认使用 CLS token 池化，可通过 --pooling mean 切换为均值池化。
3. 序列长度：默认最大序列长度为 512 token，超出部分会被截断。
4. 数值精度：NPU 与 CPU 在极低幅值维度（< 1e-5）存在微小差异，对整体嵌入质量无影响（余弦相似度 > 0.99999）。