all-MiniLM-L6-v2 Ascend NPU 部署指南

all-MiniLM-L6-v2-ascend/
├── inference.py          # 推理测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
└── test_sentences.txt   # 测试句子

docker exec -it test-modelagent bash

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

pip install transformers torch_npu

cd /data/ysws/agentsp/5-15/all-MiniLM-L6-v2-ascend/

# 使用默认测试句子
python3 inference.py

# 使用指定设备
python3 inference.py --device npu:0

cd /data/ysws/agentsp/5-15/all-MiniLM-L6-v2-ascend/

# 运行完整精度测试
python3 inference.py --mode precision_test

import torch
from transformers import AutoTokenizer, AutoModel

MODEL_DIR = "/data/ysws/agentsp/5-15/all-MiniLM-L6-v2"

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 = ["This is a test sentence", "Each sentence is converted"]
inputs = tokenizer(sentences, padding=True, truncation=True, 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'])

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}")

sentences = [
    "First sentence for embedding",
    "Second sentence for comparison",
    "Third sentence in the batch"
]

inputs = tokenizer(sentences, padding=True, truncation=True, 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'])

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

{
  "hidden_size": 384,
  "intermediate_size": 1536,
  "num_attention_heads": 12,
  "num_hidden_layers": 6,
  "vocab_size": 30522
}