from atomgit_hub import snapshot_download
snapshot_download("Ascend-SACT/Bert",  local_dir = './download')

docker run -it -d --net=host --name bert --privileged --shm-size=300G \
--device=/dev/davinci0 \
--device=/dev/davinci1 \
--device=/dev/davinci2 \
--device=/dev/davinci3 \
--device=/dev/davinci4 \
--device=/dev/davinci5 \
--device=/dev/davinci6 \
--device=/dev/davinci7 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \
-v /mnt:/mnt \
-it f49277a2e0de bash

import torch
from transformers import BertTokenizer, BertModel
import numpy as np
import warnings

# 屏蔽权重未使用的非致命警告
warnings.filterwarnings("ignore")

# --------------------------
# 1. 配置参数（统一指定NPU设备）
# --------------------------
MODEL_PATH = "/workspace/bert-base-chinese"  # BERT模型路径
LABEL_LIST = ["O", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG", "I-ORG"]  # 实体标签列表
# 设备配置：强制指定昇腾NPU:0设备，确保模型和张量统一在此设备上
DEVICE = torch.device("npu:0")

# --------------------------
# 2. 自定义NER模型（BERT + 线性层 + 简化版CRF，修复索引越界）
# --------------------------
class BertCrfForNer(torch.nn.Module):
    def __init__(self, bert_model, num_labels, label_list):
        super().__init__()
        self.bert = bert_model
        self.hidden_size = bert_model.config.hidden_size
        self.num_labels = num_labels
        
        # 线性层：将BERT输出映射到标签空间
        self.classifier = torch.nn.Linear(self.hidden_size, num_labels)
        
        # CRF层参数（保留核心，简化命名避免混淆）
        self.transitions = torch.nn.Parameter(torch.randn(num_labels, num_labels))
        self.start_transitions = torch.nn.Parameter(torch.randn(num_labels))
        self.end_transitions = torch.nn.Parameter(torch.randn(num_labels))
        
        # 标签映射
        self.label2id = {label: i for i, label in enumerate(label_list)}
        self.id2label = {i: label for i, label in enumerate(label_list)}

    def forward(self, input_ids, attention_mask=None, labels=None):
        # BERT输出
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = outputs[0]   # (batch_size, seq_len, hidden_size)
        
        # 线性层映射到标签 logits
        logits = self.classifier(sequence_output)  # (batch_size, seq_len, num_labels)
        
        if labels is not None:
            # 训练时计算CRF损失（此处省略，推理不需要）
            return logits
        else:
            # 推理时使用【简化版Viterbi解码】，优先处理单样本，规避索引越界
            return self.simple_viterbi_decode(logits, attention_mask)

    def simple_viterbi_decode(self, logits, mask):
        """
        简化版Viterbi解码（修复索引越界，优先支持单样本，兼容小批量）
        摒弃复杂的嵌套路径列表，改用张量记录前驱索引，提升稳定性
        """
        batch_size, seq_len, num_labels = logits.shape
        # 转置为 (seq_len, batch_size, num_labels)，方便按时间步迭代
        logits = logits.transpose(0, 1)  # (T, B, C)
        mask = mask.transpose(0, 1).bool()  # (T, B)，转为布尔类型掩码

        # 初始化：第一步得分 = 起始转移得分 + 第一步logits
        prev_scores = self.start_transitions.unsqueeze(0) + logits[0]  # (B, C)
        # 记录每一步的前驱标签索引，用于回溯路径 (T, B, C)
        prev_label_indices = torch.zeros((seq_len, batch_size, num_labels), dtype=torch.long, device=logits.device)

        # 迭代计算每一步的最优得分和前驱索引
        for t in range(1, seq_len):
            # 计算当前步所有标签组合的得分：(B, C, C)
            # prev_scores (B, C) -> (B, C, 1)，加上转移矩阵 (C, C)，加上当前logits (B, 1, C)
            score_matrix = prev_scores.unsqueeze(2) + self.transitions.unsqueeze(0) + logits[t].unsqueeze(1)
            
            # 取每个标签的最优得分和对应的前驱标签索引 (B, C)
            current_scores, current_prev_indices = score_matrix.max(dim=1)
            
            # 掩码处理：padding位置继承上一步的得分和索引
            mask_t = mask[t].unsqueeze(1).expand(batch_size, num_labels)  # (B, C)
            prev_scores = torch.where(mask_t, current_scores, prev_scores)
            prev_label_indices[t] = torch.where(mask_t, current_prev_indices, prev_label_indices[t-1])

        # 加上结束转移得分，获取最终得分 (B, C)
        final_scores = prev_scores + self.end_transitions.unsqueeze(0)
        # 取每个样本的最优末尾标签 (B,)
        best_last_labels = final_scores.argmax(dim=1)

        # 回溯获取完整标签序列
        pred_labels = torch.zeros((batch_size, seq_len), dtype=torch.long, device=logits.device)
        pred_labels[:, -1] = best_last_labels  # 先填充最后一步的标签

        # 从后往前回溯
        for t in range(seq_len - 2, -1, -1):
            # 取出当前步的前驱标签索引 (B, C) -> 根据上一步的最优标签，取出对应前驱索引 (B,)
            pred_labels[:, t] = prev_label_indices[t+1, torch.arange(batch_size), pred_labels[:, t+1]]

        return pred_labels

# --------------------------
# 3. 实体提取工具（复用原逻辑，优化返回格式）
# --------------------------
def get_entity_bio(seq, id2label):
    """从BIO标签序列中提取实体（独立实现，不依赖原库）"""
    chunks = []
    chunk = [-1, -1, -1]  # [实体类型, 开始位置, 结束位置]
    for idx, tag_id in enumerate(seq):
        tag = id2label[tag_id]
        if tag.startswith("B-"):
            if chunk[2] != -1:
                chunks.append(chunk)
            chunk = [tag.split("-")[1], idx, idx]
            if idx == len(seq) - 1:
                chunks.append(chunk)
        elif tag.startswith("I-") and chunk[1] != -1:
            if tag.split("-")[1] == chunk[0]:
                chunk[2] = idx
            if idx == len(seq) - 1:
                chunks.append(chunk)
        else:
            if chunk[2] != -1:
                chunks.append(chunk)
            chunk = [-1, -1, -1]
    return chunks

# --------------------------
# 4. 推理函数（统一设备，修复张量同步问题）
# --------------------------
def ner_infer(text, model, tokenizer, id2label, device):
    # 1. 文本预处理（优化：默认不padding，减少无效索引，如需固定长度可调整）
    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding=False,  # 关闭默认max_length padding，避免多余索引
        truncation=True,
        max_length=128
    )
    # 核心：将输入张量同步到指定NPU设备，与模型设备保持一致
    input_ids = inputs["input_ids"].to(device)
    attention_mask = inputs["attention_mask"].to(device)
    
    # 2. 模型推理（使用外部传入的统一设备，不硬编码设备）
    model.to(device)
    model.eval()
    with torch.no_grad():  # 推理阶段关闭梯度计算，节省NPU资源
        # 简化版解码已返回 (batch_size, seq_len) 格式，无需额外索引取值
        pred_label_ids = model(input_ids=input_ids, attention_mask=attention_mask)
        # 推理结果转回CPU，方便后续numpy操作和分词解析
        pred_label_ids = pred_label_ids[0].cpu().numpy()
    
    # 3. 提取实体
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0].cpu().numpy())
    entities = get_entity_bio(pred_label_ids, id2label)
    
    # 4. 格式化结果（关联原始文本，过滤无效实体）
    result = []
    for entity in entities:
        entity_type, start, end = entity
        # 过滤[CLS]和[SEP]特殊标记及无效索引
        if start < 0 or end >= len(tokens) or tokens[start] in ["[CLS]", "[SEP]"] or tokens[end] in ["[CLS]", "[SEP]"]:
            continue
        # 拼接子词（处理BERT分词后的##前缀，还原完整实体）
        entity_text = "".join([t.replace("##", "") for t in tokens[start:end+1]])
        result.append({
            "type": entity_type,
            "text": entity_text,
            "start": start,
            "end": end
        })
    return result

# --------------------------
# 5. 主函数（加载模型并测试，全流程NPU兼容）
# --------------------------
if __name__ == "__main__":
    # 加载分词器和BERT预训练模型
    tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
    bert_model = BertModel.from_pretrained(MODEL_PATH)
    
    # 初始化NER模型，并默认加载到NPU设备
    num_labels = len(LABEL_LIST)
    model = BertCrfForNer(bert_model, num_labels, LABEL_LIST).to(DEVICE)
    
    # （可选）加载训练好的NER权重（如需加载，确保权重映射到NPU设备）
    # model.load_state_dict(torch.load("ner_model_weights.pth", map_location=DEVICE))
    
    # 测试文本推理
    test_text = "张三在北京市海淀区的清华大学工作"
    # 传入统一的NPU设备，确保模型和输入张量设备一致
    entities = ner_infer(test_text, model, tokenizer, model.id2label, DEVICE)
    print("识别结果：", entities)

python test.py 
识别结果： [{'type': 'ORG', 'text': '张', 'start': 1, 'end': 1}, {'type': 'ORG', 'text': '在', 'start': 3, 'end': 3}, {'type': 'ORG', 'text': '京', 'start': 5, 'end': 5}, {'type': 'ORG', 'text': '海', 'start': 7, 'end': 7}, {'type': 'ORG', 'text': '区', 'start': 9, 'end': 9}, {'type': 'ORG', 'text': '清', 'start': 11, 'end': 11}, {'type': 'ORG', 'text': '大', 'start': 13, 'end': 13}]