mxbai-edge-colbert-v0-32M Ascend NPU 部署指南

mxbai-edge-colbert-v0-32m-ascend/
├── inference.py          # 推理测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
├── test_sample.txt       # 测试样本
├── inference_result.json # 推理结果
└── precision_result.json # 精度测试结果

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-18-2/mxbai-edge-colbert-v0-32m-ascend/
python3 inference.py

cd /data/ysws/agentsp/5-18-2/mxbai-edge-colbert-v0-32m-ascend/
python3 inference.py precision_test

============================================================
mxbai-edge-colbert-v0-32M NPU Test
Output: /data/ysws/agentsp/5-18-2/mxbai-edge-colbert-v0-32m-ascend
============================================================

============================================================
mxbai-edge-colbert-v0-32M Inference Test (NPU)
============================================================
Device: npu:0
Model: /data/ysws/agentsp/5-18-2/mxbai-edge-colbert-v0-32m/mixedbread-ai/mxbai-edge-colbert-v0-32m

Loading tokenizer...
Loading model...
Loading weights: 100%|██████████| 62/62 [00:00<00:00, 7313.31it/s]

Input text: ['Hello, how are you today?']
Input shape: torch.Size([1, 9])
Embedding shape: torch.Size([1, 384])
Embedding sample (first 5): [0.004708211403340101, -0.016210922971367836, 0.09353338181972504, 0.004691631533205509, 0.00407747644931078]
Inference time: 0.300s

============================================================
Precision Test (CPU vs NPU)
============================================================

Loading model on CPU...
Loading weights: 100%|██████████| 62/62 [00:00<00:00, 4594.87it/s]
Running inference on CPU...

Loading model on NPU...
Loading weights: 100%|██████████| 62/62 [00:00<00:00, 5485.29it/s]
Running inference on NPU...

CPU inference time: 0.053s
NPU inference time: 0.014s
Speedup: 3.69x
Cosine similarity: 1.000000
Max absolute error: 9.805104e-05
Max relative error: 1.911812e-04 (0.0191%)
Status: PASS

============================================================
Creating Test Sample
============================================================
Saved test sample
  1. Hello, how are you today?
  2. The weather is nice today.
  3. I am very happy to see you.

============================================================
Test Complete!

import torch
from transformers import AutoTokenizer, AutoModel

MODEL_DIR = "/data/ysws/agentsp/5-18-2/mxbai-edge-colbert-v0-32m/mixedbread-ai/mxbai-edge-colbert-v0-32m"

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, how are you today?"]
inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True, max_length=128)
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([1, 384])

from sklearn.metrics.pairwise import cosine_similarity

# 计算两个句子的相似度
emb1 = embeddings[0].cpu().numpy()
# ... (compute second embedding)
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, max_length=128)
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"Batch embeddings shape: {embeddings.shape}")  # torch.Size([3, 384])

{
  "hidden_size": 384,
  "intermediate_size": 576,
  "num_attention_heads": 6,
  "num_hidden_layers": 10,
  "vocab_size": 50370,
  "model_type": "modernbert",
  "classifier_pooling": "mean"
}