阿里天池 NLP 入门赛 Bert 方案 -1 数据预处理

前言

这篇文章用于记录阿里天池 NLP 入门赛，详细讲解了整个数据处理流程，以及如何从零构建一个模型，适合新手入门。

赛题以新闻数据为赛题数据，数据集报名后可见并可下载。赛题数据为新闻文本，并按照字符级别进行匿名处理。整合划分出 14 个候选分类类别：财经、彩票、房产、股票、家居、教育、科技、社会、时尚、时政、体育、星座、游戏、娱乐的文本数据。实质上是一个 14 分类问题。

赛题数据由以下几个部分构成：训练集 20w 条样本，测试集 A 包括 5w 条样本，测试集 B 包括 5w 条样本。

比赛地址：https://tianchi.aliyun.com/competition/entrance/531810/introduction

数据可以通过上面的链接下载。

为什么写篇文章

首先，这篇文章的代码全部都来源于 Datawhale 提供的开源代码，我添加了自己的笔记，帮助新手更好地理解这个代码。

1. Datawhale 提供的代码有哪些需要改进？

Datawhale 提供的代码里包含了数据处理，以及从 0 到 1 模型建立的完整流程。但是和前面提供的 basesline 的都不太一样，它包含了非常多数据处理的细节，模型也是由 3 个部分构成，所以看起来难度陡然上升。

其次，代码里的注释非常少，也没有讲解整个数据处理和网络的整体流程。这些对于新手来说，增加了理解的门槛。
在数据竞赛方面，我也是一个新人，花了一天的时间，仔细研究数据在一种每一个步骤的转化，对于一些难以理解的代码，在群里询问之后，也得到了 Datawhale 成员的热心解答。最终才明白了全部的代码。

2. 我做了什么改进？

所以，为了减少对于新手的阅读难度，我添加了一些内容。

1. 首先，梳理了整个流程，包括两大部分：数据处理和模型。

   因为代码不是从上到下顺序阅读的。因此，更容易让人理解的做法是：先从整体上给出宏观的数据转换流程图，其中要包括数据在每一步的 shape，以及包含的转换步骤，让读者心中有一个框架图，再带着这个框架图去看细节，会更加了然于胸。

2. 其次，除了了解了整体流程，在真正的代码细节里，读者可能还是会看不懂某一段小逻辑。因此，我在原有代码的基础之上增添了许多注释，以降低代码的理解门槛。

在 上一篇文章 中，我们讲解了使用 TextCNN + Attention + LSTM 来构建模型。这篇文章，我们来讲解使用 Bert 方案构建模型。

Bert 来自于 Transformer 的编码器（Encoder）。

如果你对 Bert 或者 Transformer 不了解，请查看这篇 图解 Transformer。

由于数据是经过脱敏的，因此不能直接使用网络上开源的训练好的模型和词向量，我们在这个数据集上，从 0 开始预训练 Bert。

这里使用 Google 提供的 Bert 代码 来进行训练。然后使用 HuggingFace 提供的转换代码来把训练好的 Tensorflow 模型转换为 PyTorch 模型。后续在这个模型的基础上，使用 PyTorch 进行微调。

代码地址：https://github.com/zhangxiann/Tianchi-NLP-Beginner

分为 3 篇文章介绍：

• 数据预处理
• Bert 源码讲解
• Bert 预训练与分类

这篇文章，我们来看下数据预处理部分。

数据预处理

我们知道，在 Bert 训练的目标有两个：

• Masked LM：输入给 BERT 的句子中，会随机选择一定数量的字或者单词置换为一个特殊的 token，称为 MASK。把数据输入 Bert，预测这些 MASK 本来的字或词。
• Next Sentence Prediction：把两个句子拼接起来，中间添加一个 SEP 的分隔符，并在最前面添加一个 CLS 的 token，把数据输入 Bert，判断这两个句子是不是属于前后关系。

在我们的任务中，只有文本分类，而没有上下文句子的推理，因此只进行 Masked LM，而不需要 Next Sentence Prediction。

所以数据预处理步骤，主要是对文本进行 mask，获得 mask 后的数据。

1. 数据准备

原始数据是 csv 格式，每行表示一篇文章和对应的标签。我们需要把所有文章存放到普通的文件中（不保存标签），每篇文章之间使用空行分隔。这里把 14 种类别的文章，平均分为 10 份。对应的代码文件为 prepare_data.py，代码如下：

# split data to 10 fold
fold_num = 10
import os
import pandas as pd
import numpy as np
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(levelname)s: %(message)s')
dir = 'data'
data_file = os.path.join(dir,'train_set.csv')


def all_data2fold(fold_num):
    fold_data = []
    f = pd.read_csv(data_file, sep='\t', encoding='UTF-8')
    texts = f['text'].tolist()
    labels = f['label'].tolist()

    total = len(labels)

    index = list(range(total))
    # 打乱数据
    np.random.shuffle(index)

    # all_texts 和 all_labels 都是 shuffle 之后的数据
    all_texts = []
    all_labels = []
    for i in index:
        all_texts.append(texts[i])
        all_labels.append(labels[i])

    # 构造一个 dict，key 为 label，value 是一个 list，存储的是该类对应的 index
    label2id = {}
    for i in range(total):
        label = str(all_labels[i])
        if label not in label2id:
            label2id[label] = [i]
        else:
            label2id[label].append(i)

    # all_index 是一个 list，里面包括 10 个 list，称为 10 个 fold，存储 10 个 fold 对应的 index
    all_index = [[] for _ in range(fold_num)]
    for label, data in label2id.items():
        # print(label, len(data))
        batch_size = int(len(data) / fold_num)
        # other 表示多出来的数据，other 的数据量是小于 fold_num 的
        other = len(data) - batch_size * fold_num
        # 把每一类对应的 index，添加到每个 fold 里面去
        for i in range(fold_num):
            # 如果 i < other，那么将一个数据添加到这一轮 batch 的数据中
            cur_batch_size = batch_size + 1 if i < other else batch_size
            # print(cur_batch_size)
            # batch_data 是该轮 batch 对应的索引
            batch_data = [data[i * batch_size + b] for b in range(cur_batch_size)]
            all_index[i].extend(batch_data)

    batch_size = int(total / fold_num)
    other_texts = []
    other_labels = []
    other_num = 0
    start = 0

    # 由于上面在分 batch 的过程中，每个 batch 的数据量不一样，这里是把数据平均到每个 batch
    for fold in range(fold_num):
        num = len(all_index[fold])
        texts = [all_texts[i] for i in all_index[fold]]
        labels = [all_labels[i] for i in all_index[fold]]

        if num > batch_size:  # 如果大于 batch_size 那么就取 batch_size 大小的数据
            fold_texts = texts[:batch_size]
            other_texts.extend(texts[batch_size:])
            fold_labels = labels[:batch_size]
            other_labels.extend(labels[batch_size:])
            other_num += num - batch_size
        elif num < batch_size:  # 如果小于 batch_size，那么就补全到 batch_size 的大小
            end = start + batch_size - num
            fold_texts = texts + other_texts[start: end]
            fold_labels = labels + other_labels[start: end]
            start = end
        else:
            fold_texts = texts
            fold_labels = labels

        assert batch_size == len(fold_labels)

        # shuffle
        index = list(range(batch_size))
        np.random.shuffle(index)
        # 这里是为了打乱数据
        shuffle_fold_texts = []
        shuffle_fold_labels = []
        for i in index:
            shuffle_fold_texts.append(fold_texts[i])
            shuffle_fold_labels.append(fold_labels[i])

        data = {'label': shuffle_fold_labels, 'text': shuffle_fold_texts}
        fold_data.append(data)

    logging.info("Fold lens %s", str([len(data['label']) for data in fold_data]))

    return fold_data


# fold_data 是一个 list，有 10 个元素，每个元素是 dict，包括 label 和 text
fold_data = all_data2fold(10)
for i in range(0, 10):
    data = fold_data[i]

    path = os.path.join(dir, "train_" + str(i))
    my_open = open(path, 'w')
    # 打开文件，采用写入模式
    # 若文件不存在,创建，若存在，清空并写入
    for text in data['text']:
        my_open.write(text)
        my_open.write('\n') # 换行
        my_open.write('\n') # 添加一个空行，作为文章之间的分隔符
    logging.info("complete train_" + str(i))
    my_open.close()

其中 all_data2fold 的作用把每个类别的数据平均分为 10 份，返回的数据是 fold_data，是一个 list，有 10 个元素，每个元素是 dict，包括 label 和 text 的 list。

然后分别写入到 10 个文件中：train_0、train_1…train_9。每篇文章之间添加一个空行作为文章之间的分隔符。

生成的数据类似下面：

2564 1846 3605 6357 3530 . . .

4399 1121 648 3750 1679 . . .

3659 6242 2106 2662 5560 . . .
.
.
.

2. 创建字典

接着，我们创建字典，保存训练集中所有的单词，并添加 [PAD]，[UNK]， [CLS]，[SEP]， [MASK] 等 token。对应的代码文件为 create_vocab.py，代码如下：

from collections import Counter
import pandas as pd
import os.path as osp
import os

word_counter = Counter()
# 计算每个词出现的次数
data_file = osp.join('data','train_set.csv')
f = pd.read_csv(data_file, sep='\t', encoding='UTF-8')
data = f['text'].tolist()
for text in data:
    words = text.split()

    for word in words:
        word_counter[word] += 1

words = word_counter.keys()
path = os.path.join('bert-mini', "vocab.txt")
my_open = open(path, 'w')

# 打开文件，采用写入模式
# 若文件不存在,创建，若存在，清空并写入
extra_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
my_open.writelines("\n".join(extra_tokens))

my_open.writelines("\n".join(words))

my_open.close()

3. 对数据进行 mask

由于上面把数据分成了 10 份，因此使用脚本 create_pretraining_data.sh，分别对 10 份数据进行 mask 处理。

脚本

10 份数据的完整脚本文件为 create_pretraining_data.sh。

以第一份数据为例，脚本如下所示：

nohup python create_pretraining_data.py
--input_file=./data/train_0
--output_file=./records/train_0.tfrecord
--vocab_file=./bert-mini/vocab.txt
--max_seq_length=256
--max_predictions_per_seq=32
--do_lower_case=True
--masked_lm_prob=0.15
--random_seed=12345
--dupe_factor=5
> create/0.log 2>&1 &

调用的 python 脚本是 create_pretraining_data.py，后面携带了大量参数，参数说明如下：

• input_file：输入文件的路径
• output_file：输出文件的路径，保存了经过 mask 的文本
• vocab_file：字典文件
• max_seq_length：每一条训练数据的最大长度，这里是指每句话的最大长度。（如有使用 NSP，那么表示两句话拼接后的最大长度）
• max_predictions_per_seq：每一条训练数据的 mask 的最大数量
• do_lower_case：True 表示忽略大小写，False 表示不忽略大小写
• do_whole_word_mask：True 表示使用 word tokenization，False 表示使用其他 word tokenization
• masked_lm_prob：产生 mask 的概率，上面设置了 max_predictions_per_seq 个 mask，而实际上的 mask 的数量是 max_predictions_per_seq 和 $ 句子长度 \times masked_lm_prob$ 的较小值。
• random_seed：随机种子
• dupe_factor：在代码会对文档多次重复随机产生训练集，这个参数是指重复的次数

在 create_pretraining_data.py 中，首先使用 tf.flags 来接收参数。

flags = tf.flags

FLAGS = flags.FLAGS

flags.DEFINE_string("input_file", None,
                    "Input raw text file (or comma-separated list of files).")

flags.DEFINE_string(
    "output_file", None,
    "Output TF example file (or comma-separated list of files).")

flags.DEFINE_string("vocab_file", None,
                    "The vocabulary file that the BERT model was trained on.")

flags.DEFINE_bool(
    "do_lower_case", True,
    "Whether to lower case the input text. Should be True for uncased "
    "models and False for cased models.")

flags.DEFINE_bool(
    "do_whole_word_mask", False,
    "Whether to use whole word masking rather than per-WordPiece masking.")

flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")

flags.DEFINE_integer("max_predictions_per_seq", 20,
                     "Maximum number of masked LM predictions per sequence.")

flags.DEFINE_integer("random_seed", 12345, "Random seed for data generation.")

flags.DEFINE_integer(
    "dupe_factor", 10,
    "Number of times to duplicate the input data (with different masks).")

flags.DEFINE_float("masked_lm_prob", 0.15, "Masked LM probability.")

flags.DEFINE_float(
    "short_seq_prob", 0.1,
    "Probability of creating sequences which are shorter than the "
    "maximum length.")

接着我们看 main() 函数，主要流程如下：

• 首先创建 WhitespaceTokenizer，作用是根据词典，将词转化为该词对应的 id。

• 然后调用 create_training_instances()，对每句话进行 mask，返回 instances，是一个 list，每个元素是 TrainingInstance，表示一个经过 mask 后的句子。

  TrainingInstance 的内容包括：

  • tokens：经过 mask 的 tokens
  • segment_ids：句子 id
  • is_random_next：True 表示 tokens 的第二句是随机查找，False 表示第二句为第一句的下文。这里不使用 NSP，只有一句话，因此为 false
  • masked_lm_positions：mask 的位置
  • masked_lm_labels：mask 对应的真实 token

• 最后调用 write_instance_to_example_files()，将 instances 保存到文件中。

def main(_):
    tf.logging.set_verbosity(tf.logging.INFO)

    # tokenizer以vocab_file为词典，将词转化为该词对应的id。
    tokenizer = tokenization.WhitespaceTokenizer(
        vocab_file=FLAGS.vocab_file)

    input_files = []
    for input_pattern in FLAGS.input_file.split(","):
        # tf.gfile.Glob: 查找匹配 filename 的文件并以列表的形式返回，
        # filename 可以是一个具体的文件名，也可以是包含通配符的正则表达式。
        input_files.extend(tf.gfile.Glob(input_pattern))

    tf.logging.info("*** Reading from input files ***")
    for input_file in input_files:
        tf.logging.info("  %s", input_file)

    rng = random.Random(FLAGS.random_seed)
    # instances 是一个 list，每个元素是 TrainingInstance，表示一个经过mask 后的句子，
    # TrainingInstance 的内容包括：
    # tokens: 经过 mask 的 tokens
    # segment_ids: 句子 id
    # is_random_next: True 表示 tokens 的第二句是随机查找，False 表示第二句为第一句的下文。
    # 这里不使用 NSP，只有一句话，因此为 false
    # masked_lm_positions: mask 的位置
    # masked_lm_labels: mask 的真实 token
    instances = create_training_instances(
        input_files, tokenizer, FLAGS.max_seq_length, FLAGS.dupe_factor,
        FLAGS.short_seq_prob, FLAGS.masked_lm_prob, FLAGS.max_predictions_per_seq,
        rng)

    output_files = FLAGS.output_file.split(",")
    tf.logging.info("*** Writing to output files ***")
    for output_file in output_files:
        tf.logging.info("  %s", output_file)

    write_instance_to_example_files(instances, tokenizer, FLAGS.max_seq_length,
                                    FLAGS.max_predictions_per_seq, output_files)

create_training_instances()

我们看来 create_training_instances() 的代码。

主要流程是：

• 首先读取文件，把每篇文章进行 tokenization（把每篇文章转换为 list），并添加到 all_documents 中，然后去除空行。
• for 循环，把 all_documents 中的每个元素传入 create_instances_from_document_nsp() 方法，进行真正的 mask。

# 输入是文件列表，输出是 instance 的列表
def create_training_instances(input_files, tokenizer, max_seq_length,
                              dupe_factor, short_seq_prob, masked_lm_prob,
                              max_predictions_per_seq, rng):
    """Create `TrainingInstance`s from raw text."""
    # all_documents 是一个 list，每个元素是一篇文章
    # 文章中的是一个 list，list中每个元素是一个单词
    all_documents = [[]]

    # Input file format:
    # (1) One sentence per line. These should ideally be actual sentences, not
    # entire paragraphs or arbitrary spans of text. (Because we use the
    # sentence boundaries for the "next sentence prediction" task).
    # (2) Blank lines between documents. Document boundaries are needed so
    # that the "next sentence prediction" task doesn't span between documents.
    for input_file in input_files:
        with tf.gfile.GFile(input_file, "r") as reader:
            while True:
                line = tokenization.convert_to_unicode(reader.readline())
                if not line:
                    break
                line = line.strip()

                # Empty lines are used as document delimiters
                # 如果是空行，作为文章之间的分隔符，那么使用一个新的 list 来存储文章
                if not line:
                    all_documents.append([])
                tokens = tokenizer.tokenize(line)
                if tokens:
                    all_documents[-1].append(tokens)

    # Remove empty documents
    # 去除空行
    all_documents = [x for x in all_documents if x]
    rng.shuffle(all_documents)
    # vocab_words 是一个 list，每个元素是单词 (word)
    vocab_words = list(tokenizer.vocab.keys())
    instances = []
    # 对文档重复 dupe_factor 次，随机产生训练集
    for _ in range(dupe_factor):
        for document_index in range(len(all_documents)):
            instances.extend(
                create_instances_from_document_nsp(
                    all_documents, document_index, max_seq_length, short_seq_prob,
                    masked_lm_prob, max_predictions_per_seq, vocab_words, rng))

    rng.shuffle(instances)
    return instances

create_instances_from_document_nsp()

create_instances_from_document_nsp() 函数把每篇文章分割为句子。

主要流程如下：

• 首先根据 document_index 取出文章document，document 是一个 list，只有一个元素，表示一篇文章。

• 然后把 document 传入 create_segments_from_document()，分割为多个句子，返回的 segments 是一个 list，里面包含多个句子。

• 接着 for 循环，对每个句子进行操作：

  • 在句子前后添加 [CLS] 和 [SEP]。
  • 调用 create_masked_lm_predictions() ，对句子进行 mask 操作。
  • 返回的内容，再组成一个 TrainingInstance。

• 将所有的 TrainingInstance 组成一个 instances 列表，并返回 instances。

# NSP: next sentence prediction
# 返回一篇文章的 TrainingInstance 的列表
# TrainingInstance 是一个 class，表示一个句子的
def create_instances_from_document_nsp(
        all_documents, document_index, max_seq_length, short_seq_prob,
        masked_lm_prob, max_predictions_per_seq, vocab_words, rng):
    """Creates `TrainingInstance`s for a single document. Remove NSP."""
    # 取出 index 对应的文章
    document = all_documents[document_index]

    # 留出 2 个位置给 [CLS], [SEP]
    # Account for [CLS], [SEP]
    max_segment_length = max_seq_length - 2

    instances = []
    # document 是一个 list，只有一个元素，表示一篇文章
    segments = create_segments_from_document(document, max_segment_length)
    for j, segment in enumerate(segments):
        is_random_next = False
        # tokens 是添加了 [CLS] 和 [SEP] 的句子
        tokens = []
        # 表示句子的 id，都是 0，表示只有 1 个句子
        segment_ids = []
        # 在开头添加 [CLS]
        tokens.append("[CLS]")
        segment_ids.append(0)
        for token in segment:
            tokens.append(token)
            segment_ids.append(0)
        # 在开头添加 [SEP]
        tokens.append("[SEP]")
        # 添加句子 id:0
        segment_ids.append(0)
        # 对一个句子进行 mask ，获得 mask 后的 tokens，mask 的位置以及真实的 token(label)
        (tokens, masked_lm_positions,
         masked_lm_labels) = create_masked_lm_predictions(
            tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng)
        # TrainingInstance 表示一个句子的内容
        instance = TrainingInstance(
            tokens=tokens, # 经过 mask 的 token 列表
            segment_ids=segment_ids, # 句子 id
            is_random_next=is_random_next, # False
            masked_lm_positions=masked_lm_positions, # mask 的位置列表
            masked_lm_labels=masked_lm_labels) # mask 的真实 token
        instances.append(instance)

    return instances

create_masked_lm_predictions()

create_masked_lm_predictions() 是核心函数，在这个函数里，对每个句子进行真正的 mask，并返回 mask 后的内容。

输入的参数是 tokens，表示一个句子。

主要流程如下：

• 首先把每个 token 的索引添加到 cand_indexes 中，并打乱，这是为了下面进行 mask 操作。

• 然后计算实际的 mask 数量num_to_predict，计算方法是取 max_predictions_per_seq 和 $ 句子长度 \times masked_lm_prob$ 的较小值。

• 进入 for 循环，遍历每个 token：

  • 首先判断已经生成的 mask 数量是否已经达到 num_to_predict，如果达到了这个数量，则退出循环。
  • 80% 的概率替换为 [MASK]。
  • 10% 的概率保留原来的 token。
  • 10% 的概率替换为随机的一个 token。
  • 保存 mask 的位置和真实的 token

• mask 完成后，再把所有 masked_lms 根据 index 进行排序。

• 返回 3 个参数：

  • output_tokens：经过 mask 的 tokens，表示一个句子
  • masked_lm_positions：记录 mask 的索引
  • masked_lm_labels：记录 mask 对应的真实 token，也就是标签

# tokens 是一个句子
# 返回 mask 后的 tokens，mask 的位置以及真实的 token(label)
def create_masked_lm_predictions(tokens, masked_lm_prob,
                                 max_predictions_per_seq, vocab_words, rng):
    """Creates the predictions for the masked LM objective."""

    cand_indexes = []
    for (i, token) in enumerate(tokens):
        if token == "[CLS]" or token == "[SEP]":
            continue
        # Whole Word Masking means that if we mask all of the wordpieces
        # corresponding to an original word. When a word has been split into
        # WordPieces, the first token does not have any marker and any subsequence
        # tokens are prefixed with ##. So whenever we see the ## token, we
        # append it to the previous set of word indexes.
        #
        # Note that Whole Word Masking does *not* change the training code
        # at all -- we still predict each WordPiece independently, softmaxed
        # over the entire vocabulary.
        if (FLAGS.do_whole_word_mask and len(cand_indexes) >= 1 and
                token.startswith("##")):
            cand_indexes[-1].append(i)
        else: # 只会执行这个条件，添加的是一个 list，里面只有一个元素
            cand_indexes.append([i])

    rng.shuffle(cand_indexes)

    output_tokens = list(tokens)
    # 得到 mask 的数量
    num_to_predict = min(max_predictions_per_seq,
                         max(1, int(round(len(tokens) * masked_lm_prob))))

    masked_lms = []
    covered_indexes = set()
    for index_set in cand_indexes:
        # 如果 mask 的数量大于 num_to_predict，就停止
        if len(masked_lms) >= num_to_predict:
            break
        # If adding a whole-word mask would exceed the maximum number of
        # predictions, then just skip this candidate.
        # index_set 是一个 list，里面只有一个元素
        if len(masked_lms) + len(index_set) > num_to_predict:
            continue
        is_any_index_covered = False
        for index in index_set:
            if index in covered_indexes:
                is_any_index_covered = True
                break
        # 如果已经包含了这些 index，那么就跳过
        if is_any_index_covered:
            continue

        for index in index_set:
            # covered_indexes 是一个 set
            covered_indexes.add(index)

            masked_token = None
            # 80% of the time, replace with [MASK]
            # 80% 的概率替换为 mask
            if rng.random() < 0.8:
                masked_token = "[MASK]"
            else:
                # 剩下的 20%，再分为两半
                # 10% of the time, keep original
                # 20% *0.5 =10% 的概率保留原来的token
                if rng.random() < 0.5:
                    masked_token = tokens[index]
                # 10% of the time, replace with random word
                else:
                    # 20% *0.5 =10% 的概率替换为随机的一个 token
                    masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)]

            output_tokens[index] = masked_token

            masked_lms.append(MaskedLmInstance(index=index, label=tokens[index]))
    # 执行完循环后，mask 的数量小于 num_to_predict
    assert len(masked_lms) <= num_to_predict
    # 根据 index 排序
    masked_lms = sorted(masked_lms, key=lambda x: x.index)

    # masked_lm_positions 保存 mask 的位置
    masked_lm_positions = []
    # masked_lm_labels 保存 mask 的真实 token
    masked_lm_labels = []
    for p in masked_lms:
        masked_lm_positions.append(p.index)
        masked_lm_labels.append(p.label)

    return (output_tokens, masked_lm_positions, masked_lm_labels)

write_instance_to_example_files()

现在我们再回到最外层的代码，我们得到经过 mask 的所有句子的数据 instances（instances是一个 list，每个元素是 TrainingInstance），然后调用 write_instance_to_example_files() ，将所有不到最大长度的部分用 0 补齐，保存到文件中。

代码主要流程如下：

• for 循环，遍历每个 TrainingInstance：
  • 把 token 转换为 id。
  • 创建 input_ids（token 的 id）、input_mask（默认为 1，当句子长度小于 max_seq_length，就补 0）、segment_ids（句子 id），并补充长度到 max_seq_length。
  • 创建 masked_lm_positions（mask 的位置）、masked_lm_ids（mask 对应的真实 token）、masked_lm_weights（默认为 1，当 mask 数量小于 max_predictions_per_seq，就补 0），并补充长度到 max_predictions_per_seq。
  • 创建 tf.train.Example，保存到 output_file 中。

# 把所有的 instance 保存到 output_files 中
def write_instance_to_example_files(instances, tokenizer, max_seq_length,
                                    max_predictions_per_seq, output_files):
    """Create TF example files from `TrainingInstance`s."""
    writers = []
    for output_file in output_files:
        writers.append(tf.python_io.TFRecordWriter(output_file))

    writer_index = 0

    total_written = 0
    for (inst_index, instance) in enumerate(instances):
        # 把 token 转为 id
        input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)
        input_mask = [1] * len(input_ids)
        segment_ids = list(instance.segment_ids)
        assert len(input_ids) <= max_seq_length
        # 补全为 0，补到长度为 max_seq_length
        while len(input_ids) < max_seq_length:
            input_ids.append(0)
            input_mask.append(0)
            segment_ids.append(0)

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length

        masked_lm_positions = list(instance.masked_lm_positions)
        masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels)
        # weight 默认都是 1
        masked_lm_weights = [1.0] * len(masked_lm_ids)
        # 补全为 0，补到长度为 max_predictions_per_seq，表示每句话 mask 的数量
        while len(masked_lm_positions) < max_predictions_per_seq:
            masked_lm_positions.append(0)
            masked_lm_ids.append(0)
            masked_lm_weights.append(0.0)

        # 这里是 0
        next_sentence_label = 1 if instance.is_random_next else 0

        features = collections.OrderedDict()
        # 这是为了保存到文件中
        features["input_ids"] = create_int_feature(input_ids)
        features["input_mask"] = create_int_feature(input_mask)
        features["segment_ids"] = create_int_feature(segment_ids)
        features["masked_lm_positions"] = create_int_feature(masked_lm_positions)
        features["masked_lm_ids"] = create_int_feature(masked_lm_ids)
        features["masked_lm_weights"] = create_float_feature(masked_lm_weights)
        features["next_sentence_labels"] = create_int_feature([next_sentence_label])

        tf_example = tf.train.Example(features=tf.train.Features(feature=features))

        writers[writer_index].write(tf_example.SerializeToString())
        writer_index = (writer_index + 1) % len(writers)

        total_written += 1
        # 打印前 20 句话的内容
        if inst_index < 20:
            tf.logging.info("*** Example ***")
            tf.logging.info("tokens: %s" % " ".join(
                [tokenization.printable_text(x) for x in instance.tokens]))

            for feature_name in features.keys():
                feature = features[feature_name]
                values = []
                if feature.int64_list.value:
                    values = feature.int64_list.value
                elif feature.float_list.value:
                    values = feature.float_list.value
                tf.logging.info(
                    "%s: %s" % (feature_name, " ".join([str(x) for x in values])))

    for writer in writers:
        writer.close()

    tf.logging.info("Wrote %d total instances", total_written)

下图是一个经过 mask 的句子的输出：

至此，数据预处理部分就讲完了。

在下一篇文章中，我们会进入到 Bert 源码讲解。

如果你有疑问，欢迎留言。

参考

• 谷歌 BERT 预训练源码解析
• 阿里天池 NLP 入门赛 TextCNN 方案代码详细注释和流程讲解

如果你觉得这篇文章对你有帮助，不妨点个赞，让我有更多动力写出好文章。