自建数据集用BiLSTM+CRF实现命名实体识别(NER)

本文使用的代码来自：
GitHub - ZejunCao/NER_baseline: 使用多种方法做中文命名实体识别(NER)，代码包含详细注释
参考的代码解读来自：
逐行讲解BiLSTM+CRF实现命名实体识别(NER)_bilstm+crf命名实体识别-CSDN博客
本人根据自己项目的实际情况，需要使用深度学习模型实现命名实体识别，通过参考上述文章的代码解读成功运行，非常感谢上述博客作者的耐心解答！在跑代码时根据自己的需求进行了一些修改，在此记录一下，也供大家参考。

一、数据准备

该模型输入的数据需要为经过BIO标注格式的数据，我主要针对自己的数据格式修改了data_processor.py文件。本人使用的数据为逐条待抽取实体的.txt文本，对输入模型数据的处理步骤如下：

1. 获取.ann格式的标注文件

对逐条的.txt文件，使用精灵标注软件，标注每个标签得到每条.txt文件对应的.ann格式的标注文件

2. .ann格式转化为BIO标注后的txt文件

将ann标注形式的文件转换为BIO标注形式的

举例来说，本人未经处理的一条.txt数据为：

打齿。客户反映车辆发动机异响严重。检查发现车辆飞轮打齿损坏导致故障。更换飞轮，故障排除。。

该条数据经过精灵标注后得到的.ann文件为：

T1    故障件 0 2    打齿
T2    故障主要现象 9 16    发动机异响严重
T3    真实故障原因 23 29    飞轮打齿损坏

表示该条文本中标注了三种类型的实体，每种实体各有一个，将.ann格式转化为BIO标注后的txt文件，形式如下：

打,B-故障件
齿,I-故障件
。,O
客,O
户,O
反,O
映,O
车,O
辆,O
发,B-故障主要现象
动,I-故障主要现象
机,I-故障主要现象
异,I-故障主要现象
响,I-故障主要现象
严,I-故障主要现象
重,I-故障主要现象
。,O
...

在本人使用的代码中，还增加了将BIO标注后的txt文件划分为训练集和测试集，即划分后将单条BIO标注后的数据连接起来，最后得到两个txt文件：训练集和测试集，代码如下：

import glob
import random

def bratann2BIO_format(text, ann_str, fstream):
    # 将每一行的元素变为list，strip()删除的字符,按照split()中的符号进行每行元素分割为list的元素
    ann_list = ann_str.strip().split('\n')

    label = ['O' for _ in range(len(text))]  # 对所有的文字赋值为标签"O"
    for i, line in enumerate(ann_list):  # enumerate函数用于遍历序列中的元素以及它们的下标
        try:
            # line：T1    疾病 4 7  高血压
            T, typ, word = line.strip().split('\t')
            t, s, e = typ.split()  # 分别t=疾病 s=起始位置下标 e=结束位置
            s, e = int(s), int(e)
            label[s] = 'B-' + t
            while s < e - 1:
                s += 1
                label[s] = 'I-' + t
        except:
            continue

    for t, l in zip(list(text), label):  # list() 构造函数在 Python 中返回一个列表，将文本以字切分为列表
        # str.join(item1,itemm2)，join函数是一个字符串操作函数,使用str符号将item1和item2串联起来
        line = ','.join([t, l])
        print(line)
        fstream.write(line)
        fstream.write('\n')  # 每一个文本(一行)写完，然后进行换行
    fstream.write('\n')  # 使用换行符，将每一个文本用一个空行分开，在train.txt文档中可以很清晰的看到句与句的切分

def gen_NER_trainANDtest_data(root_dir,stream_train,stream_test,test_size=0.3):
    file_list = glob.glob(root_dir + '/*.ann')  # glob.glob() 函数的作用：在一个文件中，要遍历所有的文件内容
    random.seed(0)
    random.shuffle(file_list)
    n = int(len(file_list) * test_size)
    test_files = file_list[:n]
    train_files = file_list[n:]
    #生成训练集数据
    for ann_path in train_files:
        ann_path = ann_path.replace('\\', '/')
        # 获得txt路径
        txt_path = ann_path.replace('/outputs', '').replace('ann', 'txt')

        try:
            ft = open(txt_path, 'r', encoding='utf8')
            text = ft.read().strip()
            ft.close()
            fa = open(ann_path, 'r', encoding='utf8')
            ann = fa.read().strip()
            fa.close()
            if ann == '':
                continue
            bratann2BIO_format(text, ann, stream_train)
        except Exception as e:
            print(ann_path, e)
    stream_train.close()

    #生成测试集数据
    for ann_path in test_files:
        ann_path = ann_path.replace('\\', '/')
        # 获得txt路径
        txt_path = ann_path.replace('/outputs', '').replace('ann', 'txt')

        try:
            ft = open(txt_path, 'r', encoding='utf8')
            text = ft.read().strip()
            ft.close()
            fa = open(ann_path, 'r', encoding='utf8')
            ann = fa.read().strip()
            fa.close()
            if ann == '':
                continue
            bratann2BIO_format(text, ann, stream_test)
        except Exception as e:
            print(ann_path, e)
    stream_test.close()

使用上述代码，指定root_dir（标注文件ann所在文件夹目录）和stream_train（训练集文本输出路径）、stream_test（测试集文本输出路径），最后调用gen_NER_trainANDtest_data函数即可

到此处即完成数据准备工作，下面进入预测

---

二、预测

使用BiLSTM-CRF算法训练模型、测试效果部份的代码，全部使用参考博客中的代码，未改动。

在使用predict.py文件对未标注数据进行预测时，改动了代码，在预测代码中调用模型之后，做了批量预测和结果文字版呈现，代码如下：

#批量预测
model.eval()
model.state = 'pred'

folder_path = '...' # 待预测的.txt文件所在的文件夹路径
# 用于存储预测结果的列表
predictions = []
fenkuaiTQ=[]
Text=[]
# 遍历文件夹中的所有.txt文件
for filename in os.listdir(folder_path):
    if filename.endswith('.txt'):
        # 构造文件的完整路径
        file_path = os.path.join(folder_path, filename)

        # 打开并读取文件中的文本
        with open(file_path, 'r', encoding='utf-8') as file:
            text = file.read().strip()  # 读取并去除可能的空白字符

        # 对文本进行预测
        with torch.no_grad():
            # 将文本转换为索引序列
            text_indices = [vocab.get(t, vocab['UNK']) for t in text]
            # 创建序列长度张量并移动到设备
            seq_len = torch.tensor(len(text_indices), dtype=torch.long).unsqueeze(0)
            seq_len = seq_len.to(device)
            # 创建文本索引张量并移动到设备
            text_tensor = torch.tensor(text_indices, dtype=torch.long).unsqueeze(0)
            text_tensor = text_tensor.to(device)
            # 进行预测
            batch_tag = model(text_tensor, seq_len)
            # 将预测的索引转换回标签
            pred = [label_map_inv[t] for t in batch_tag]
            # 存储预测结果
            Text.append(text)
            predictions.append(pred)
            fenkuaiTQ.append(chunks_extract(pred))

        # 打印当前文本的预测结果和分块提取结果
        #print(f"预测的文本: {text}")
        #print(f"预测结果: {pred}")
        #print(f"分块提取结果: {chunks_extract(pred)}\n")

# 循环结束后，predictions列表将包含所有文件的预测结果
predictions_df = pd.DataFrame({
    'Predictions': predictions,
    '抽取结果': fenkuaiTQ,
    'Text': Text
})
#predictions_df['Predictions'] = predictions_df['Predictions'].apply(ast.literal_eval)#将'Predictions'列由str转换为list

#将预测结果转换为对应文字
def extract_text_based_on_predictions(row):
    predictions = row['Predictions']
    text = row['Text']
    result = []

    # 用于存储当前标签的开始索引
    current_label_start = None

    for i, pred in enumerate(predictions):
        if pred.startswith('B-'):
            # 找到新的标签开始处，保存索引
            current_label_start = i
            current_label = pred[2:]
            result.append(f"{current_label}：{text[i]}")
        elif pred.startswith('I-') and current_label_start is not None:
            # 如果是内部标签，继续添加到当前标签的结果中
            result[-1] += text[i] + ""
        elif pred == 'O':
            if current_label_start is not None:
                # 遇到'O'标签，结束当前标签的文本收集
                current_label_start = None
    # 将结果列表转换为字符串
    return '。'.join(result)

# 应用函数到每一行，并创建新列
predictions_df['Reconstructed_Text'] = predictions_df.apply(extract_text_based_on_predictions, axis=1)

# 将 DataFrame 输出为 Excel 文件
predictions_df.to_excel('路径.xlsx', index=False,header=True)