nlp_corom_sentence-embedding_chinese-tiny on Ascend NPU

1. 简介

本文档记录 CoROM 中文句嵌入模型在昇腾 NPU (Ascend 910B3) 上的迁移适配与精度验证结果。

该模型基于 BERT-tiny 架构（4 层 Transformer，256 维隐层），通过对比学习训练实现中文句子级别的语义向量表示，可用于语义搜索、文本聚类等任务。

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/damo/nlp_corom_sentence-embedding_chinese-tiny
• 权重下载地址（HuggingFace Mirror）：https://hf-mirror.com/damo/nlp_corom_sentence-embedding_chinese-tiny

2. 验证环境

• NPU：Ascend 910B3，单卡推理
• 模型大小：约 8.8M 参数
• 输入：中文文本，最大 512 tokens
• 输出：[1, seq_len, 256] 隐层表示

3. 推理启动

3.1 环境准备

conda create -n nlp_corom_sentence-embedding_chinese-tiny python=3.10 -y
conda activate nlp_corom_sentence-embedding_chinese-tiny

pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://mirrors.huaweicloud.com/repository/pypi/simple
pip install transformers modelscope numpy

3.2 下载模型权重

modelscope download --model damo/nlp_corom_sentence-embedding_chinese-tiny \
    --local_dir ./nlp_corom_sentence-embedding_chinese-tiny

3.3 命令行推理

python inference.py --text "这是一段测试文本" --device npu

3.4 Python API

from inference import run_inference
result = run_inference("测试文本", device="npu")
print(f"向量维度: {result['embedding'].shape}, 耗时: {result['time_ms']:.1f}ms")

4. Smoke 验证

python evaluate.py

验证结果：

• CPU 推理正常，输出 [1, seq_len, 256]
• NPU 推理正常，输出同形状
• 精度评测通过

5. 性能参考

测试条件：输入 4 个 token 的中文短句。

6. 精度评测

7. 注意事项

1. 加载方式：使用 transformers.AutoModel.from_pretrained 加载，为 HuggingFace 兼容格式
2. 输入格式：需配合 AutoTokenizer 分词，中文按字分词
3. 同系列模型：nlp_structbert_sentence-similarity_chinese-tiny 为同一架构的相似度版本，适配方式一致
4. 首次推理：NPU 首次推理有算子编译开销，后续稳定在 ~8ms