BGE-Reranker-Base 昇腾 NPU 适配

# 安装 torch_npu（如果已有可跳过）
pip install torch torch_npu -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 transformers
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple

python3 inference.py [选项]

# 基本推理
python3 inference.py

# 使用本地模型
python3 inference.py --model_path /path/to/bge-reranker-base

# 性能基准测试（20 次运行）
python3 inference.py --num_runs 20 --warmup_runs 5

import torch
import torch_npu
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 加载模型
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-base")
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-base")
model.eval()
model = model.npu()

# 准备数据
pairs = [
    ("什么是机器学习？", "机器学习是人工智能的一个分支。"),
    ("什么是机器学习？", "今天天气很好。"),
]
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=512)
inputs = {k: v.npu() for k, v in inputs.items()}

# 推理
with torch.no_grad():
    outputs = model(**inputs)
    scores = outputs.logits.squeeze(-1)
    print(scores)  # 正值表示相关，负值表示不相关

bge-reranker-base-npu/
├── README.md          # 本文件（部署文档）
├── inference.py       # NPU 推理脚本
└── eval_accuracy.py   # 精度验证脚本

modelscope download --model BAAI/bge-reranker-base --local_dir /path/to/model

npu-smi info  # 查看 NPU 显存使用情况

# NPU 推理
python3 inference.py --device npu

# CPU 推理
python3 inference.py --device cpu