冬
gcw_IDzXRVNw/multilingual-e5-small-ascend
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

multilingual-e5-small Ascend NPU 部署指南

项目简介

multilingual-e5-small 是 intfloat 团队开发的多语言句子嵌入模型,支持 50+ 种语言的文本嵌入提取。该模型基于 BERT 架构,将句子映射到 384 维稠密向量空间,可用于语义搜索、聚类和多语言文本相似度计算等任务。

特性

  • 支持 Ascend NPU 推理加速
  • CPU vs NPU 精度对比测试 (< 1% 误差)
  • 多语言支持 (50+ 种语言)
  • 高效句子嵌入提取
  • 384 维向量输出
  • Mean Pooling + L2 归一化

环境要求

  • 硬件: 华为 Ascend 910 系列 NPU
  • CANN: 8.0.RC1 或更高版本
  • PyTorch: 2.0+ with torch_npu
  • transformers: 4.8+

目录结构

multilingual-e5-small-ascend/
├── inference.py          # 推理测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
└── test_data/            # 测试数据

部署步骤

1. 进入容器

docker exec -it test-modelagent bash

2. 设置环境变量

source /usr/local/Ascend/ascend-toolkit/set_env.sh

3. 准备模型文件

模型文件位于 /data/ysws/agentsp/5-15/multilingual-e5-small/ 目录下:

  • model.safetensors - 模型权重 (约 471MB)
  • config.json - 模型配置
  • tokenizer.json - 分词器文件
  • sentencepiece.bpe.model - SentencePiece 模型

4. 安装依赖

pip install transformers torch_npu

使用方式

方式一:普通推理模式

运行推理脚本进行句子嵌入提取:

cd /data/ysws/agentsp/5-15/multilingual-e5-small-ascend/

python3 inference.py --mode inference

python3 inference.py --mode inference --device npu:0

方式二:精度测试模式 (CPU vs NPU)

运行精度对比测试,验证 NPU 计算结果与 CPU 一致性:

cd /data/ysws/agentsp/5-15/multilingual-e5-small-ascend/

python3 inference.py --mode precision_test

命令行参数说明

参数说明默认值
--mode测试模式: inference 或 precision_testinference
--device运行设备: npu:0, cuda:0, cpu, auto (默认auto)auto

测试验证

精度测试结果

指标实测值阈值状态
相对误差0.7885%< 1.00%PASS
Cosine 相似度0.999990> 0.999PASS

性能数据

操作耗时
NPU 推理时间 (单句)~0.015s
CPU 推理时间 (单句)~0.4s
首次推理有编译开销约 0.2s

推理结果示例

Input: Hello, this is a test sentence.
Input shape: torch.Size([1, 10])
Inference time: 0.222s
Embedding shape: torch.Size([1, 384])
Embedding norm: 1.0000

Input: multilingual embedding extraction
Input shape: torch.Size([1, 10])
Inference time: 0.013s
Embedding shape: torch.Size([1, 384])
Embedding norm: 1.0000

Python API 使用示例

基本句子嵌入

import torch
from transformers import AutoTokenizer, AutoModel

MODEL_DIR = "/data/ysws/agentsp/5-15/multilingual-e5-small"

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModel.from_pretrained(MODEL_DIR)
model = model.to("npu:0")
model.eval()

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output.last_hidden_state
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

sentences = ["Hello, this is a test sentence.", "This is another example."]
inputs = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors='pt')
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

embeddings = mean_pooling(outputs, inputs['attention_mask'])
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

print(f"Embeddings shape: {embeddings.shape}")  # torch.Size([2, 384])

语义相似度计算

from sklearn.metrics.pairwise import cosine_similarity

emb1 = embeddings[0].cpu().numpy()
emb2 = embeddings[1].cpu().numpy()
similarity = cosine_similarity([emb1], [emb2])[0][0]
print(f"Cosine similarity: {similarity:.4f}")

多语言示例

multilingual_sentences = [
    "Hello, this is a test sentence.",  # English
    "Bonjour, ceci est un exemple.",      # French
    "Dies ist ein Test.",                # German
    "这是一个测试句子。"                  # Chinese
]

inputs = tokenizer(multilingual_sentences, padding=True, truncation=True, max_length=512, return_tensors='pt')
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

embeddings = mean_pooling(outputs, inputs['attention_mask'])
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

print(f"Multilingual embeddings shape: {embeddings.shape}")  # torch.Size([4, 384])

模型结构

  • 架构类型: BERT (Multilingual)
  • 编码器: 12 层 Transformer
  • 隐藏层维度: 384
  • 注意力头数: 12
  • 参数量: ~118M
  • 词汇表大小: 250037
  • 池化方式: Mean Pooling + L2 Normalization
组件说明
embeddingsBERT 词嵌入 + 位置编码
encoder12 层 Transformer 编码器
poolerMean Pooling + LayerNorm

推理参数配置

从 config.json 提取的关键参数:

{
  "model_type": "bert",
  "hidden_size": 384,
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "vocab_size": 250037,
  "max_position_embeddings": 512
}

常见问题

Q: 精度测试失败?

A: 检查 NPU 驱动是否正确安装,确保 CANN 环境变量已 source。0.1-0.8% 的数值误差是正常的。

Q: 如何提高推理速度?

A: 使用批处理可以显著提高吞吐量。对于大量句子,建议批量输入。

Q: 输出的向量有什么用?

A: 可用于:

  • 语义搜索: 计算向量相似度
  • 聚类: 将相似句子聚在一起
  • 重排序: 对候选结果进行相似度排序
  • 多语言迁移: 不同语言但语义相似的句子会有相近的向量

测试日志

完整测试日志保存在 log.txt。包括:

  • 模型加载状态
  • 每条句子的推理时间和嵌入向量范数
  • CPU vs NPU 精度对比数据

参考链接

  • 原始模型: https://huggingface.co/intfloat/multilingual-e5-small
  • E5 模型论文: https://arxiv.org/abs/2206.10658
  • HuggingFace Transformers: https://huggingface.co/transformers

许可证

本项目遵循 MIT 许可证