Jina Reranker v1 Turbo EN on Ascend NPU

1. 简介

本文档记录 jinaai/jina-reranker-v1-turbo-en 在 Ascend 910B3 NPU 环境下的适配与验证结果。

jina-reranker-v1-turbo-en 是 Jina AI 发布的轻量级交叉编码器重排序模型，基于 JinaBERT 架构。JinaBERT 是 BERT 的变体，支持对称双向 ALiBi（Attention with Linear Biases）位置编码，可处理最长 8192 token 的输入序列。该模型通过知识蒸馏技术，在保持 competitive 精度的同时实现极快的推理速度。

本仓库提供：

• inference.py：NPU 推理脚本，支持单条/批量查询-文档对评分
• eval.py：精度与性能评测脚本，对比 NPU 与 CPU 输出差异
• log.txt：评测运行日志

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/jinaai/jina-reranker-v1-turbo-en
• 镜像加速：https://hf-mirror.com

参考文档：

• https://jina.ai/reranker/
• https://arxiv.org/abs/2310.19923（JinaBERT论文）
• https://arxiv.org/abs/2108.12409（ALiBi论文）
• https://www.sbert.net/docs/pretrained_cross-encoders.html

2. 验证环境

• NPU：Ascend 910B3，1 逻辑卡
• 模型架构：JinaBertForSequenceClassification
• 隐藏层维度：384
• 注意力头数：12
• 隐藏层数：6
• 位置编码：ALiBi（支持外推到 8192 token）
• 词汇表大小：61,056
• 参数量：约 32M

3. 推理启动

启动前可先检查 NPU 可用性：

python3 -c "import torch; print(f'NPU available: {torch.npu.is_available()}')"

环境准备：

# 安装依赖
pip install torch torch_npu transformers

# 设置 NPU 可见设备（默认使用 0 号卡）
export ASCEND_RT_VISIBLE_DEVICES=0

已验证通过的推理命令：

单对查询-文档评分：

python inference.py \
  --query "What is machine learning?" \
  --passage "Machine learning is a subset of AI that enables systems to learn from data."

批量推理（JSONL 格式）：

# pairs.jsonl 每行：{"query": "...", "passage": "..."}
python inference.py --input-file pairs.jsonl

CPU 推理（参考基准）：

python inference.py --device cpu --query "test" --passage "example"

输出格式：

{
  "model": "jinaai--jina-reranker-v1-turbo-en",
  "device": "npu:0",
  "num_pairs": 1,
  "inference_time_seconds": 0.0021,
  "pairs_per_second": 476.19,
  "results": [
    {
      "query": "What is machine learning?",
      "passage": "Machine learning is a subset of AI...",
      "score": 4.527361
    }
  ]
}

4. Smoke 验证

python inference.py --query "test" --passage "example passage"

预期输出：

• 返回 JSON 格式结果
• score 字段为相关性得分（正数表示正相关）
• inference_time_seconds 应在毫秒级别（< 50ms）
• 无错误输出到 stderr

验证结果：

• 推理可在 NPU 上正常完成
• 输出格式符合预期 JSON 结构
• 评分结果具备语义区分度
• 首次推理包含模型加载时间，后续推理更快

5. 性能参考

测试条件：8 对查询-文档，batch_size=8，max_length=8192。

6. 精度评测

精度评测采用余弦距离（Cosine Distance）作为主要指标，对比 NPU 与 CPU 输出的评分向量。

结论：NPU 与 CPU 输出的评分向量余弦距离仅 0.000048%，远低于 1% 的阈值，精度通过验证。

7. 注意事项

1. 自定义代码：该模型依赖 JinaBERT 自定义实现（modeling_bert.py、configuration_bert.py），加载时需设置 trust_remote_code=True。
2. num_labels：模型为单输出回归任务（ranking score），加载时需指定 num_labels=1 及 ignore_mismatched_sizes=True 以匹配权重形状。
3. 长序列支持：得益于 ALiBi 位置编码，模型支持最长 8192 token 的输入序列，远超标准 BERT 的 512 token 限制。
4. 输入分隔符：查询和文档使用 </s><s> 作为分隔符拼接（JinaBERT 的特定格式）。
5. 评分含义：输出为原始 logits 值，正数表示正相关，负数表示负相关，绝对值越大置信度越高。
6. 重排序场景：该模型适用于检索后的精排阶段（Re-ranking），建议先通过稀疏检索或嵌入检索召回 Top-K 候选文档，再使用本模型重排序。