基于 BiLSTM+CRF 实现 NER 任务 – 训练预测

孟宝亮2025-07-28我的文章1144

我们接下来编写训练函数、评估函数、预测函数。

1. 训练函数

由于我们希望批次输入训练数据，在使用 RNN、GRU、LSTM 时，可以使用 pad_sequence、packed_pad_sequence、pad_packed_sequence 等函数组合来实现批次数据的输入。在将输入进行编码时，使用 BertTokenizer，为了能够准确的按字分开，输入给 BertTokenizer 的句子最好用空格隔开，便于其能够正确的切分出字粒度的 Token。

def pad_batch_inputs(data, labels, tokenizer):

    # 函数需要返回一个按照内容长度从大到小排序过的，sentence 和 label, 还要返回 sentence 长度
    # 将批次数据的输入和标签值分开，并计算批次的输入长度
    data_inputs, data_length, data_labels = [], [], []
    for data_input, data_label in zip(data, labels):

        # 对输入句子进行编码
        data_input_encode = tokenizer.encode(data_input,
                                             return_tensors='pt',
                                             add_special_tokens=False)
        data_input_encode = data_input_encode.to(device)
        data_inputs.append(data_input_encode.squeeze())

        # 去除多余空格，计算句子长度
        data_input = ''.join(data_input.split())
        data_length.append(len(data_input))

        # 将标签转换为张量
        data_labels.append(torch.tensor(data_label, device=device))


    # 对一个批次的内容按照长度从大到小排序, 符号表示降序
    sorted_index = np.argsort(-np.asarray(data_length))
    # 根据长度的索引进行排序
    sorted_inputs, sorted_labels, sorted_length = [], [], []
    for index in sorted_index:
        sorted_inputs.append(data_inputs[index])
        sorted_labels.append(data_labels[index])
        sorted_length.append(data_length[index])

    # 对张量进行填充，使其变成长度一样的张量
    pad_inputs = pad_sequence(sorted_inputs)

    return pad_inputs, sorted_labels, sorted_length


def train():

    # 读取数据集
    train_data = load_from_disk('data/bilstm_crf_data')['train']
    # 构建分词器
    tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')
    # 构建模型
    model = NER(vocab_size=tokenizer.vocab_size, label_num=7).cuda(device)
    # 批次大小
    batch_size = 16
    # 优化器
    optimizer = optim.AdamW(model.parameters(), lr=3e-5)
    # 训练轮数
    num_epoch = 50

    # 开始训练
    def start_train(data_inputs, data_labels, tokenizer):

        # 对批量数据进行填充对齐
        pad_inputs, sorted_labels, sorted_length = \
            pad_batch_inputs(data_inputs, data_labels, tokenizer)

        # 计算损失
        loss = model(pad_inputs, sorted_labels, sorted_length)
        # 梯度清零
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 参数更新
        optimizer.step()
        # 统计损失
        nonlocal total_loss
        total_loss += loss.item()

    for epoch in range(num_epoch):
        # 统计损失
        total_loss = 0.0
        # 开始训练
        train_data.map(start_train,
                       input_columns=['data_inputs', 'data_labels'],
                       batched=True,
                       batch_size=batch_size,
                       fn_kwargs={'tokenizer': tokenizer},
                       desc='epoch: %d' % (epoch + 1))

        # 打印损失
        print('epoch: %d loss: %.3f' % (epoch + 1, total_loss))
        # 存储模型
        model.save_model('data/BiLSTM-CRF-%d.bin' % (epoch + 1))

if __name__ == '__main__':
    train()

训练过程输入损失如下：

epoch: 1 loss: 646530.528
epoch: 2 loss: 322091.170
epoch: 3 loss: 231602.128
epoch: 4 loss: 182503.388
epoch: 5 loss: 149661.446
epoch: 6 loss: 125265.489
epoch: 7 loss: 106000.216
epoch: 8 loss: 90156.823
epoch: 9 loss: 76753.042
epoch: 10 loss: 65195.731
epoch: 11 loss: 55057.508
epoch: 12 loss: 46142.997
epoch: 13 loss: 38262.882
epoch: 14 loss: 31367.848
epoch: 15 loss: 25804.164
epoch: 16 loss: 23088.218
epoch: 17 loss: 18572.793
epoch: 18 loss: 15087.369
epoch: 19 loss: 13349.071
epoch: 20 loss: 11036.982
epoch: 21 loss: 9877.440
epoch: 22 loss: 8125.003
epoch: 23 loss: 6185.219
epoch: 24 loss: 5038.318
epoch: 25 loss: 4788.188
epoch: 26 loss: 3986.904
epoch: 27 loss: 3301.328
epoch: 28 loss: 2872.219
epoch: 29 loss: 2691.445
epoch: 30 loss: 3137.029
epoch: 31 loss: 1619.925
epoch: 32 loss: 2354.775
epoch: 33 loss: 2398.700
epoch: 34 loss: 1662.177
epoch: 35 loss: 1534.003
epoch: 36 loss: 1720.913
epoch: 37 loss: 1415.495
epoch: 38 loss: 1402.394
epoch: 39 loss: 1028.855
epoch: 40 loss: 1089.962
epoch: 41 loss: 1016.576
epoch: 42 loss: 1140.854
epoch: 43 loss: 1103.749
epoch: 44 loss: 735.662
epoch: 45 loss: 926.108
epoch: 46 loss: 1021.045
epoch: 47 loss: 1039.586
epoch: 48 loss: 608.818
epoch: 49 loss: 846.596
epoch: 50 loss: 685.285

最终得到的所有模型文件链接为：https://www.aliyundrive.com/s/HV1EoeYvKQ3 提取码: su42

2. 评估函数

我们这里就对 data/BiLSTM-CRF-50.bin 模型进行评估，当然你可以使用其他的模型的来评估。评估主要的步骤：

  1. 计算下测试集中各个不同类别的实体名称，并将这些名称分别存储到不同列表中；
  2. 将测试集送入模型得到预测，并存储不同类别的实体名称；
  3. 计算不同类别的实体的精度、召回率，以及整个测试集的准确率。
def evaluate():

    # 读取测试数据
    valid_data = load_from_disk('data/bilstm_crf_data')['valid']

    # 1. 计算各个不同类别总实体数量

    # 计算测试集实体数量
    total_entities = {'ORG': [], 'PER': [], 'LOC': []}
    def calculate_handler(data_inputs, data_labels):
        # 将 data_inputs 转换为没有空格隔开的句子
        data_inputs = ''.join(data_inputs.split())

        # 提取句子中的实体
        extract_entities = extract_decode(data_labels, data_inputs)
        # 统计每种实体的数量
        nonlocal total_entities
        for key, value in extract_entities.items():
            total_entities[key].extend(value)

    # 统计不同实体的数量
    valid_data.map(calculate_handler, input_columns=['data_inputs', 'data_labels'])
    print(total_entities)


    # 2. 计算模型预测的各个类别实体数量
    model_param = torch.load('data/BiLSTM-CRF-50.bin')
    model = NER(**model_param['init']).cuda(device)
    model.load_state_dict(model_param['state'])

    # 构建分词器
    tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')

    model_entities = {'ORG': [], 'PER': [], 'LOC': []}
    def start_evaluate(data_inputs):

        # 对输入文本进行分词
        model_inputs = tokenizer.encode(data_inputs, add_special_tokens=False, return_tensors='pt')[0]
        model_inputs = model_inputs.to(device)
        # 文本送入模型进行计算
        with torch.no_grad():
            label_list = model.predict(model_inputs)

        # 统计预测的实体数量
        text = ''.join(data_inputs.split())

        # 从预测结果提取实体名字
        extract_entities = extract_decode(label_list, text)
        nonlocal model_entities
        for key, value in extract_entities.items():
            model_entities[key].extend(value)

    # 统计预测不同实体的数量
    valid_data.map(start_evaluate, input_columns=['data_inputs'], batched=False)
    print(model_entities)

    # 3. 统计每个类别的召回率
    total_pred_correct = 0
    total_true_correct = 0
    for key in total_entities.keys():

        # 获得当前 key 类别真实和模型预测实体列表
        true_entities = total_entities[key]
        true_entities_num = len(true_entities)
        pred_entities = model_entities[key]

        # 分解预测实体中，pred_correct 表示预测正确，pred_incorrect 表示预测错误
        pred_correct, pred_incorrect = 0, 0
        for pred_entity in pred_entities:
            if pred_entity in true_entities:
                pred_correct += 1
                continue
            pred_incorrect += 1

        # 模型预测的 key 类别的实体数量
        model_pred_key_num = true_entities_num + pred_incorrect

        # 计算共预测正确多少个实体
        total_pred_correct += pred_correct
        # 计算共有多少个真实的实体
        total_true_correct += true_entities_num

        # 计算精度
        print(key, '查全率: %.3f' % (pred_correct / true_entities_num))
        print(key, '查准率: %.3f' % (pred_correct / model_pred_key_num))
        print('-' * 50)

    print('准确率: %.3f' % (total_pred_correct / total_true_correct))


def extract_decode(label_list, text):
    """
    :param label_list: 模型输出的包含标签序列的一维列表
    :param text: 模型输入的句子
    :return: 提取到的实体名字
    """

    labels = ['O', 'B-ORG', 'I-ORG', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC']
    label_to_index = {label: index for index, label in enumerate(labels)}
    B_ORG, I_ORG = label_to_index['B-ORG'], label_to_index['I-ORG']
    B_PER, I_PER = label_to_index['B-PER'], label_to_index['I-PER']
    B_LOC, I_LOC = label_to_index['B-LOC'], label_to_index['I-LOC']

    # 提取连续的标签代表的实体
    def extract_word(start_index, next_label):

        # index 表示最后索引的位置
        index, entity = start_index + 1, [text[start_index]]
        for index in range(start_index + 1, len(label_list)):
            if label_list[index] != next_label:
                break
            entity.append(text[index])

        return index, ''.join(entity)

    # 存储提取的命名实体
    extract_entites, index = {'ORG': [], 'PER': [],  'LOC': []}, 0
    # 映射下一个持续的标签
    next_label = {B_ORG: I_ORG, B_PER: I_PER, B_LOC: I_LOC}
    # 映射词的所属类别
    word_class = {B_ORG: 'ORG', B_PER: 'PER', B_LOC: 'LOC'}

    while index < len(label_list):
        # 获得当前位置的标签
        label = label_list[index]
        if label in next_label.keys():
            # 将当前位置和对应的下一个持续标签传递到 extract_word 函数
            index, word = extract_word(index, next_label[label])
            extract_entites[word_class[label]].append(word)
            continue
        index += 1

    return extract_entites


if __name__ == '__main__':
    evaluate()

程序输出结果：

ORG 查全率: 0.762
ORG 查准率: 0.655
--------------------------------------------------
PER 查全率: 0.824
PER 查准率: 0.716
--------------------------------------------------
LOC 查全率: 0.846
LOC 查准率: 0.778
--------------------------------------------------
准确率: 0.821

3. 预测函数

预测函数就是输入一个句子，提取句子中的实体。步骤如下：

  1. 输入的句子先添加上空格隔开；
  2. BertTokenizer 进行编码；
  3. 使用 model 进行预测；
  4. 使用维特比解码解析出实体名称。
def entity_extract(text):

    # 构建分词器
    tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')
    # 初始化模型
    model_param = torch.load('data/BiLSTM-CRF-48.bin')
    model = NER(**model_param['init']).cuda(device)
    model.load_state_dict(model_param['state'])

    # 我们先按字将其分开，并在字之间添加空格，便于 Bert 分词器能够准确按字分割
    input_text = ' '.join(list(text))
    model_inputs = tokenizer.encode(input_text, add_special_tokens=False, return_tensors='pt')[0]
    model_inputs = model_inputs.to(device)

    with torch.no_grad():
        outputs = model.predict(model_inputs)

    return extract_decode(outputs, ''.join(input_text.split()))


if __name__ == '__main__':
    text = '我要感谢洛杉矶市民议政论坛、亚洲协会南加中心、美中关系全国委员会、美中友协美西分会等友好团体的盛情款待。'
    result = entity_extract(text)
    print(result)

程序输出结果：

{'ORG': ['亚洲协会南加中心', '美中关系全国委员会', '美中友协美西分会'], 'PER': [], 'LOC': ['洛杉矶']}
未经允许不得转载：一亩三分地 » 基于 BiLSTM+CRF 实现 NER 任务 – 训练预测
2 赞

文章目录

©2024 老孟笔记 京ICP备2021009384号-1anbei 冀公网安备13050302001966号