BAAI/bge-reranker-v2-minicpm-layerwise on Ascend NPU

import torch
import torch_npu
from transformers import AutoTokenizer

# 手动加载自定义模块
exec(open("configuration_minicpm_reranker.py").read())
exec(open("modeling_minicpm_reranker.py").read().replace(
    "from .configuration_minicpm_reranker import LayerWiseMiniCPMConfig", "# handled"
))

config = LayerWiseMiniCPMConfig.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = LayerWiseMiniCPMForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.bfloat16, config=config
)
model = model.to("npu:1").eval()

# 推理：给定 query + passage，获取相关性评分
text = query + "\n" + passage
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
inputs = {k: v.to("npu:1") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs, cutoff_layers=[12], output_hidden_states=True)

logits = outputs.logits
score = logits[0][0, -1].item() if isinstance(logits, tuple) else logits[0, -1].item()

# 测试 query-passage 相关性打分
pairs = [
    ("what is panda?", "The giant panda is a bear species endemic to China."),  # 相关
    ("what is panda?", "hi"),  # 不相关
]

for query, passage in pairs:
    text = query + "\n" + passage
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    inputs = {k: v.to("npu:1") for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs, cutoff_layers=[12], output_hidden_states=True)
    logits = outputs.logits
    score = logits[0][0, -1].item() if isinstance(logits, tuple) else logits[0, -1].item()
    print(f"Score: {score:.4f}")