我们使用的是 MSRA 中文 NER 数据集，该数据集共包含三个目录：test、valid、train，分别对应了测试集、验证集、训练集。训练集有：42000 条数据，验证集有 3000 条数据，测试集有 3442 条数据。 另外也包括了一个定义了预测标签的 tags.txt 文件，数据集使用的是 BIO 的标签体系，即：O、B-PER、I-PER、B-ORG、I-ORG、B-LOC、I-LOC，目录结构为：

msra
├── tags.txt
├── test
│   ├── sentences.txt
│   └── tags.txt
├── train
│   ├── sentences.txt
│   └── tags.txt
└── valid
    ├── sentences.txt
    └── tags.txt

数据集链接：https://www.aliyundrive.com/s/HkNk51zog6gi 提取码: 60oq

import pandas as pd
import torch
import pickle
import os
from datasets import Dataset, DatasetDict
from transformers import BertTokenizer
import datasets
import transformers

# 显示进度条
datasets.disable_progress_bar()
# 设置日志级别
datasets.logging.set_verbosity(datasets.logging.CRITICAL)
transformers.logging.set_verbosity(transformers.logging.CRITICAL)


# 修改工作目录
os.chdir(os.path.dirname(os.path.abspath(__file__)))

1. 读取数据集

我们将 valid、train 数据集合并为 train 数据集。由于我们使用的是 Bert 模型来完成 NER 任务，而 Bert 模型对输入有 max_len 限制。所以，我们得把长度大于 505 的数据过滤掉。当然这里的 505 可以修改，只要总长度不超过 512 就可以。最终，我们会生成两个 csv 文件，01-训练集.csv 和 02-测试集.csv 分别存储训练集和测试集数据。

# 加载训练集
def load_train_corpus():

    train_path = ['msra/train/sentences.txt', 'msra/train/tags.txt']
    valid_path = ['msra/valid/sentences.txt', 'msra/valid/tags.txt']
    data_path = [train_path, valid_path]

    data_inputs, data_labels = [], []
    max_len = 0
    for x_path, y_path in data_path:
        for data_input, data_label in zip(open(x_path), open(y_path)):

            data_input = data_input.split()
            data_label = data_label.split()

            if len(data_input) > max_len:
                max_len = len(data_input)

            if len(data_input) > 505:
                continue

            if len(data_input) != len(data_label):
                continue

            data_labels.append(' '.join(data_label))
            data_inputs.append(' '.join(data_input))

    # 数据存储
    train_data = pd.DataFrame()
    train_data['data_inputs'] = data_inputs
    train_data['data_labels'] = data_labels
    train_data.to_csv('data/01-训练集.csv')

    print('训练集数据量:', len(train_data), '最大句子长度:', max_len)


# 加载测试集
def load_valid_corpus():

    x_path = 'msra/test/sentences.txt'
    y_path = 'msra/test/tags.txt'

    data_inputs, data_labels = [], []
    max_len = 0
    for data_input, data_label in zip(open(x_path), open(y_path)):

        data_input = data_input.split()
        data_label = data_label.split()

        if len(data_input) > max_len:
            max_len = len(data_input)

        if len(data_input) > 505:
            continue

        if len(data_input) != len(data_label):
            continue

        data_labels.append(' '.join(data_label))
        data_inputs.append(' '.join(data_input))

    # 数据存储
    valid_data = pd.DataFrame()
    valid_data['data_inputs'] = data_inputs
    valid_data['data_labels'] = data_labels
    valid_data.to_csv('data/02-测试集.csv')

    print('测试集数据量:', len(valid_data), '最大句子长度:', max_len)

程序的输出结果：

训练集数据量: 44968 最大句子长度: 746
测试集数据量: 3438 最大句子长度: 2427

2. 标签数值化

这里要做的事情就是将 O、B-LOC、I-PER…等标签映射为数字索引表示。转换完成之后，存储到 data/03-train，具体实现代码如下：

def data_handler(data_labels_batch, data_inputs):

    labels = []
    for data_labels in data_labels_batch:
        label = []
        for data_label in data_labels.split():
            label.append(label_to_index[data_label])
        labels.append(label)

    return {'data_label_ids': labels, 'data_inputs': data_inputs}


if __name__ == '__main__':

    load_train_corpus()
    load_valid_corpus()
    
    train_data = pd.read_csv('data/01-训练集.csv')
    valid_data = pd.read_csv('data/02-测试集.csv')

    train_data = Dataset.from_pandas(train_data)
    valid_data = Dataset.from_pandas(valid_data)

    # 由于在 centos 里 load_dataset 总是报错，所以按照以下步骤加载
    # 1. 先使用 pd.read_csv 从 csv 文件加载数据
    # 2. 再使用 Dataset.from_pandas 从 DataFram 加载数据
    # 3. 最后使用 DatasetDict 构建包含训练集和测试集的数据对象
    data = DatasetDict({'train': train_data, 'valid': valid_data})

    # input_columns: 指定传递到 data_handler 中的列，以位置参数的形式传递
    # fn_kwargs: 指定传递到 data_handler 中的自定义关键字参数
    train_data = data.map(data_handler,
                          input_columns=['data_labels', 'data_inputs'],
                          batched=True,
                          batch_size=1000)

    train_data.save_to_disk('data/03-train')