class Rerank_Model_Load(nn.Module):
    def __init__(self, model_name, use_gpu):
        super(Rerank_Model_Load, self).__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, model_name))
        if use_gpu:
            self.model = AutoModelForSequenceClassification.from_pretrained(os.path.join(model_path, model_name)).to("npu:0")
        else:
            self.model = AutoModelForSequenceClassification.from_pretrained(os.path.join(model_path, model_name)).to("cpu")
        self.use_gpu = use_gpu

    def forward(self, pairs):
        with torch.no_grad():
            if self.use_gpu:
                inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=1024).to("npu:0")
            else:
                inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=1024)
            scores = self.model(**inputs, return_dict=True).logits.view(-1, ).float().cpu().numpy().tolist()
        # 添加 normalize
        normalized_scores = [s / max(scores) for s in scores]
        return normalized_scores