Ascend 3010如何让模型使用指定的npu进行推理

class RerankerV2:
    def __init__(self):
        from FlagEmbedding import FlagReranker
        import torch
        if hasattr(torch, "npu") and torch.npu.is_available():
            torch.npu.set_device(4)
        self.model = FlagReranker('BAAI/bge-reranker-v2-m3',
                                  use_fp16=True)

    def rank(self, query: str, documents: list[str]):
        return self.model.compute_score([(query, r) for r in documents], normalize=True)

# 本来想拆分出两个文件的，但是想了下这样写方便切换使用（毕竟这种rerank模型没有必要动态切换，之间也没有什么特别的优劣）
reranker = RerankerV2()


def test(query):
    attention = [
        "告诉我某月有多少个bug",
        "统计某年出现了多少个bug单"
        "查询某个字段的值",  # find 中 projection
        "统计某种数据的总数",  # countDocuments / $count
        "按照条件过滤数据",  # find + filter 条件
        "获取某个字段的唯一值",  # distinct
        "根据时间范围筛选记录",  # range filter
        "使用正则模糊匹配字段",  # $regex
        "根据多个条件联合查询",  # $and / $or
        "对查询结果排序",  # $sort
        "分页查询数据",  # skip + limit
        "判断某字段是否存在",  # $exists
        "聚合多个字段生成新字段",  # $project + $concat / $addFields
        "查询字段为 null 或缺失的记录",  # null / not exists
        "统计某字段的最大最小值",  # $group + $max / $min
        "筛选某字段包含关键字的记录",  # $regex or $text
        "获取最新一条记录",  # $sort + limit(1)
        "判断是否存在符合条件的记录"  # findOne / count > 0
    ]

    scores = reranker.rank(query, attention + [])
    print("query: ", query)
    print("document: ", attention)
    # print("scores: ", softmax(scores))
    print("scores: ", scores)

def test_rank():
    test("IOT平台2025年12月有多少个bug")
    test("IOT平台2025年12月有多少个需求？")
    test("系统管理员如何添加员工")
    test("今天天气真好啊")
    test("补一个设定 第一阶段的结尾（泽做全了通往另外分支的准备），泽穿戴好装备，便开始了计数，向着实验场的中心跑去。（这段也要改一下，改成冷冻仓+发射轨道，内部穿着宇航服，应对不时之需）（传送过程中冷冻仓被裁切掉，刚越过分支的时候，因气体逃逸导致泽在宇宙中疯狂翻滚，几次撞击后宇航服的氧气存量也不高了，大概40%左右，十来分钟） 突然，一只口袋空间在正面前打开，只是一秒，泽便消失在了口袋空间中。 通信全断。 他被吸入了真空，然后口袋从这个世界中消失，一群亮晶晶的尘雾从中间置换了出来，融入在了实验室的空气中。泽用的装备是为了在真空生存下去的，有点像是一个可以运动的冷冻仓，有一定的氧气存量，目的就是在那个地方可以生存一段时间，直到Niyredra能找泽")

reranker = RerankerV2()
class RerankerV2:
    def __init__(self):
        from FlagEmbedding import FlagReranker
        self.model = FlagReranker('BAAI/bge-reranker-v2-m3',
                                  use_fp16=True, devices=["npu:4"])

    def rank(self, query: str, documents: list[str]):
        return self.model.compute_score([(query, r) for r in documents], normalize=True)

# 本来想拆分出两个文件的，但是想了下这样写方便切换使用（毕竟这种rerank模型没有必要动态切换，之间也没有什么特别的优劣）
reranker = RerankerV2()