multilingual-e5-small 是 intfloat 团队开发的多语言句子嵌入模型,支持 50+ 种语言的文本嵌入提取。该模型基于 BERT 架构,将句子映射到 384 维稠密向量空间,可用于语义搜索、聚类和多语言文本相似度计算等任务。
multilingual-e5-small-ascend/
├── inference.py # 推理测试脚本
├── log.txt # 测试日志
├── README.md # 本文档
└── test_data/ # 测试数据docker exec -it test-modelagent bashsource /usr/local/Ascend/ascend-toolkit/set_env.sh模型文件位于 /data/ysws/agentsp/5-15/multilingual-e5-small/ 目录下:
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运行精度对比测试,验证 NPU 计算结果与 CPU 一致性:
cd /data/ysws/agentsp/5-15/multilingual-e5-small-ascend/
python3 inference.py --mode precision_test| 参数 | 说明 | 默认值 |
|---|---|---|
--mode | 测试模式: inference 或 precision_test | inference |
--device | 运行设备: npu:0, cuda:0, cpu, auto (默认auto) | auto |
| 指标 | 实测值 | 阈值 | 状态 |
|---|---|---|---|
| 相对误差 | 0.7885% | < 1.00% | PASS |
| Cosine 相似度 | 0.999990 | > 0.999 | PASS |
| 操作 | 耗时 |
|---|---|
| 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.0000import 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])| 组件 | 说明 |
|---|---|
| embeddings | BERT 词嵌入 + 位置编码 |
| encoder | 12 层 Transformer 编码器 |
| pooler | Mean 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
}A: 检查 NPU 驱动是否正确安装,确保 CANN 环境变量已 source。0.1-0.8% 的数值误差是正常的。
A: 使用批处理可以显著提高吞吐量。对于大量句子,建议批量输入。
A: 可用于:
完整测试日志保存在 log.txt。包括:
本项目遵循 MIT 许可证