rubert-tiny2 on Ascend NPU

1. 简介

本文档记录 cointegrated/rubert-tiny2 俄语轻量句嵌入模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

rubert-tiny2 是俄语 TinyBERT 的知识蒸馏版本（312 维隐藏层，极轻量），通过 Mean Pooling + L2 归一化输出句嵌入。模型在俄语语料上预训练，适用于俄语语义搜索、文本聚类和相似度计算。额外包含 LaBSE（Language-Agnostic BERT Sentence Embedding）适配器权重（tinybert-ru-labse-adapter-v2.pt），可将 312 维嵌入映射到多语言联合空间，实现跨语言检索。

TinyBERT 312 维的极轻量设计使其成为俄语 NLP 场景中推理速度最快的嵌入模型之一。

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/cointegrated/rubert-tiny2

2. 验证环境

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

3. 部署使用流程

3.1 环境准备

conda create -n rubert-tiny2 python=3.11 -y
conda activate rubert-tiny2

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

3.2 推理脚本使用

python inference.py --text "Привет мир" --device npu
python inference.py --batch_file russian_texts.txt --device npu

编程接口：

from inference import RuBERTEncoder
encoder = RuBERTEncoder(model_path="./rubert-tiny2", device="npu")
embeddings = encoder.encode(["Привет мир", "Как дела?"])
# embeddings.shape → (2, 312)

4. Smoke 验证

python inference.py --text "Привет мир" --device npu

预期输出：312 维归一化嵌入向量，无运行时错误。

5. 性能参考

测试条件：23 条多样化测试句子（含俄语/英语），batch_size=32，NPU 预热 1 轮。

TinyBERT 312 维极轻量设计使 NPU 吞吐高达 1703 sent/s，是目前适配模型中推理最快的嵌入模型。相比 768 维 BERT-base（~900 sent/s），速度提升近 2×。

6. 精度评测

6.1 评测方法

分别在 CPU 和 NPU 上对 23 条测试句子推理，比较 312 维嵌入向量的余弦相似度和语义相似度矩阵的 Pearson 相关系数。

6.2 评测结果

结论：精度误差率 0.0000%，NPU 与 CPU 输出完全一致，评测通过。

7. 迁移适配说明

7.1 模型结构

• Backbone：BertForPreTraining（TinyBERT 蒸馏版，312 维隐藏层，极轻量）
• Pooling：Mean Pooling（对所有 token 嵌入取 mask 加权平均）
• Normalize：L2 归一化到单位球面
• Tokenizer：BERT WordPiece（vocab.txt），俄语 Cyrillic 字符支持
• 额外组件：LaBSE adapter（tinybert-ru-labse-adapter-v2.pt），312 维 → 768 维多语言空间映射
• 参数量：约 29M（仅为 BERT-base 的 1/4）

7.2 适配要点

1. 使用 AutoModel.from_pretrained() 加载（BertForPreTraining 兼容 BertModel 接口）
2. model.to("npu:0") 迁移，312 维 TinyBERT 算子编译极快
3. 手动实现 Mean Pooling：(embeddings * mask).sum(1) / mask.sum(1)
4. 支持 sentence-transformers 格式（含 modules.json、1_Pooling、sentence_bert_config.json）
5. LaBSE adapter 为独立 .pt 文件，需通过 torch.load 单独加载并应用于嵌入输出

7.3 关键代码

import torch, torch_npu
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("rubert-tiny2").to("npu:0")
tokenizer = AutoTokenizer.from_pretrained("rubert-tiny2")

text = "Привет мир"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)
    embeddings = outputs.last_hidden_state
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    sentence_emb = (embeddings * mask).sum(1) / mask.sum(1)
    sentence_emb = torch.nn.functional.normalize(sentence_emb, p=2, dim=1)

8. 注意事项

1. 俄语特化模型：该模型在俄语语料上预训练，分词器支持 Cyrillic 字符。对英文和其他拉丁语系语言也兼容但非优化场景。中文等非 Cyrillic/拉丁语系文本将退化为字符级分词。
2. TinyBERT 极轻量：312 维 + 蒸馏架构使参数量仅 29M（BERT-base 的 1/4），NPU 吞吐高达 1700+ sent/s。适合高吞吐、低延迟的俄语 NLP 场景。
3. LaBSE 适配器：模型附带的 tinybert-ru-labse-adapter-v2.pt 是一个独立线性映射层，将 312 维嵌入投影到 LaBSE 768 维多语言空间，实现跨语言（俄-英等）相似度计算。适配器需单独加载：adapter = torch.load(adapter_path) → multilingual_emb = adapter(russian_emb)。
4. sentence-transformers 兼容：模型使用 sentence-transformers 格式保存（modules.json + sentence_bert_config.json），可直接用 SentenceTransformer 加载。NPU 适配使用手动 Mean Pooling 替代以简化依赖。
5. 首次 NPU 推理：TinyBERT 仅 312 维、层数少，算子编译约 1-2 秒，是所有模型中最快的预热时间。