cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 on Ascend NPU

1. 简介

本文档记录 cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 cross-encoder 模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

Cross-encoder 是一种用于信息检索精排（re-ranking）的模型。与 bi-encoder（如 sentence-transformers）不同，cross-encoder 同时接收 query 和 passage 作为输入，通过全注意力（full attention）交互计算相关性分数。该模型基于 XLM-RoBERTa（12 层，384 维隐藏层），在 MS MARCO 数据集上训练，输出单个 sigmoid 相关性分数（0-1），适用于搜索结果的精排阶段。

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1

2. 验证环境

• NPU：8 × Ascend 910B3
• 精度对比基准：CPU（x86, PyTorch 2.8.0）

3. 部署使用流程

3.1 环境准备

conda create -n cross-encoder--mmarco-mMiniLMv2-L12-H384-v1 python=3.11 -y
conda activate cross-encoder--mmarco-mMiniLMv2-L12-H384-v1

pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers sentencepiece numpy \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

XLM-RoBERTa 使用 sentencepiece tokenizer，需要安装 sentencepiece 包。

3.2 模型权重下载

HF_ENDPOINT=https://hf-mirror.com \
    huggingface-cli download cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 \
    --local-dir ./cross-encoder--mmarco-mMiniLMv2-L12-H384-v1

3.3 推理脚本使用

python inference.py --query "What is AI?" --passage "AI is artificial intelligence." --device npu
python inference.py --pairs_file pairs.tsv --device npu

编程接口：

from inference import CrossEncoder

ce = CrossEncoder(
    model_path="./cross-encoder--mmarco-mMiniLMv2-L12-H384-v1",
    device="npu"
)
scores = ce.predict([
    ("What is AI?", "Artificial intelligence is the simulation of human intelligence."),
    ("How to cook pasta?", "Boil water and add salt."),
])
# scores → array([0.92, 0.15])

4. Smoke 验证

python inference.py --query "What is machine learning?" \
    --passage "Machine learning is a subset of artificial intelligence." \
    --device npu

预期输出：相关性分数（0-1 之间的浮点数），高分表示 query 与 passage 相关；无运行时错误。

5. 性能参考

测试条件：20 对英文 query-passage 测试样本，batch_size=16，NPU 预热 1 轮。

Cross-encoder 的 query 和 passage 拼接后序列较长（512 tokens），NPU 加速比略低于短文本模型，但仍提供显著提升。

6. 精度评测

6.1 评测方法

分别在 CPU 和 NPU 上对 20 对 query-passage 样本推理，比较 sigmoid 相关性分数的精度：

• MAE：分数逐元素平均绝对误差
• 余弦相似度：分数向量方向一致性
• 最大误差：单样本最大分数差异

6.2 评测结果

结论：精度误差率 0.0001%，远低于 1% 要求，评测通过。

7. 迁移适配说明

7.1 模型结构

• Backbone：XLMRobertaModel（12 层 Transformer，384 维隐藏层，MiniLM 蒸馏版本）
• Classifier Head：线性层（384 → 1），二分类（相关/不相关）
• 激活函数：Sigmoid，输出 0-1 相关性分数
• Tokenizer：SentencePiece（spm.model），多语言支持
• 参数量：118M（MiniLM 蒸馏，比 XLM-RoBERTa-base 更轻量）

7.2 适配要点

1. 使用 AutoModelForSequenceClassification.from_pretrained() 加载
2. 输入格式为 (query, passage) 对，通过 tokenizer(query, passage) 拼接编码
3. model.to("npu:0") 迁移，sigmoid 激活在 NPU 上完成
4. sentencepiece tokenizer 从 spm.model 文件加载，CPU 端分词
5. ONNX 格式权重可用于跨平台加速部署

7.3 关键代码

import torch, torch_npu
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
).to("npu:0")
tokenizer = AutoTokenizer.from_pretrained(
    "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
)

query, passage = "What is AI?", "AI is artificial intelligence."
inputs = tokenizer(query, passage, return_tensors="pt", truncation=True)
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    score = torch.sigmoid(model(**inputs).logits).item()

8. 注意事项

1. sentencepiece 依赖：XLM-RoBERTa 使用 sentencepiece 分词，需安装 pip install sentencepiece，缺少会导致 tokenizer 初始化失败。
2. Cross-encoder 不可单独编码：该模型需同时传入 query 和 passage，不可先编码 passage 再存向量库（应使用 bi-encoder 实现）。
3. 序列长度限制：query + passage 拼接后总长度受 max_length=512 限制，长 passage 将被截断。
4. NPU 预热：首次推理触发算子编译（约 3-5 秒），建议生产环境预热后使用。
5. 与 ms-marco-MiniLM-L12-v2 对比：两者同属 MiniLM 蒸馏系列（12 层 384 维），mmarco 版本在多语言检索场景有优势，ms-marco 版本在纯英文场景更优。