KBLab_megatron.bert-base.unigram-32k-no_pretok.25k-steps

1. 简介

本文档记录 KBLab_megatron.bert-base.unigram-32k-no_pretok.25k-steps 在昇腾 NPU（Ascend910）环境的快速部署与验证结果。

AutoModel 文本嵌入模型，基于 HuggingFace transformers + sentence-transformers 框架，支持将文本转换为固定维度的语义向量，适用于句子相似度计算、语义检索等场景。

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/KBLab_megatron.bert-base.unigram-32k-no_pretok.25k-steps

参考文档：

• https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoModel
• https://www.sbert.net/

2. 验证环境

• NPU：Ascend910（单卡）
• 嵌入维度：768
• 最大序列长度：128
• 推理框架：PyTorch + transformers

3. 快速部署

3.1 环境准备

pip install transformers torch sentence-transformers

3.2 推理代码

import torch
from transformers import AutoTokenizer, AutoModel

device = torch.device("npu:0" if torch.npu.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained("KBLab_megatron.bert-base.unigram-32k-no_pretok.25k-steps")
model = AutoModel.from_pretrained("KBLab_megatron.bert-base.unigram-32k-no_pretok.25k-steps")
model = model.to(device).eval()

texts = ["What is machine learning?", "Hello world"]
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=128)
inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)
embeddings = outputs.pooler_output
print(f"嵌入维度: {embeddings.shape}")

4. Smoke 验证

python3 inference.py

验证结果：

• 模型成功加载到 npu:0
• 输出嵌入向量（batch_size × embedding_dim）
• 推理过程无报错

5. 性能参考

测试条件：FP32 / batch=2 / max_length=128 / warmup=5 / timed=50 runs，Ascend910 单卡。

6. 精度评测

NPU 与 CPU 输出对比，使用 5 组英文测试文本，比较嵌入向量的 cosine 相似度与逐元素误差。

精度要求：cosine 相似度 > 0.999，逐元素误差 < 0.01。

7. 注意事项

• 该模型使用 AutoModel 加载，输出 pooler_output 作为句子嵌入
• sentence-transformers 库可提供更丰富的 pooling 策略（如 CLS、MEAN、MAX）
• 嵌入模型精度验证以 cosine 相似度为主指标，logits 相对误差为辅助指标