代码示例 / 自然语言处理 / 用Hugging Face Transformers预训练BERT

用Hugging Face Transformers预训练BERT

作者: Sreyan Ghosh
创建日期: 2022/07/01
最后修改日期: 2022/08/27
描述: 使用Hugging Face Transformers在NSP和MLM上预训练BERT。

在Colab中查看 GitHub源代码


介绍

BERT(双向编码器表示来自Transformers)

在计算机视觉领域,研究人员反复证明了迁移学习的价值——在已知任务/数据集上对神经网络模型进行预训练,例如ImageNet分类,然后进行微调——使用训练好的神经网络作为新的专用模型的基础。近年来,研究人员表明,类似的技术在许多自然语言任务中也可以发挥作用。

BERT利用了Transformer,一种学习文本中单词(或子词)之间上下文关系的注意机制。在其原始形式中,Transformer包含两个独立的机制——一个读取文本输入的编码器和一个为任务生成预测的解码器。由于BERT的目标是生成语言模型,因此仅需要编码器机制。Transformer的详细工作机制在谷歌的一篇论文中进行了描述。

与方向性模型不同,方向性模型顺序读取文本输入(从左到右或从右到左),Transformer编码器一次读取整个单词序列。因此,它被认为是双向的,尽管更准确地说它是无方向的。这一特性使得模型能够根据单词的所有周围环境(单词的左侧和右侧)学习单词的上下文。

在训练语言模型时,一个挑战是定义预测目标。许多模型预测序列中的下一个单词(例如,"小孩从_回到家"),这是一种方向性方法,固有地限制了上下文学习。为了克服这个挑战,BERT使用了两种训练策略:

掩蔽语言建模(MLM)

在将单词序列输入BERT之前,每个序列中的15%单词会被替换为[MASK]标记。模型然后尝试根据序列中其他未掩蔽单词提供的上下文来预测被掩蔽单词的原始值。

下一句预测(NSP)

在BERT的训练过程中,模型接收成对的句子作为输入,并学习预测第二个句子是否是原始文档中的后续句子。在训练期间,50%的输入是一对,其中第二个句子是原始文档中的后续句子,而另外50%则随机选择语料库中的一个句子作为第二个句子。假设随机句子将与第一个句子产生断裂。

尽管谷歌为英语提供了一个预训练的BERT检查点,您可能经常需要针对其他语言从头预训练模型,或者进行继续预训练以使模型适应新的领域。在这个笔记本中,我们从头开始对BERT进行预训练,优化MLM和NSP目标,使用🤗 Transformers在来自🤗 Datasets的WikiText英语数据集上进行训练。


设置

安装要求

pip install git+https://github.com/huggingface/transformers.git
pip install datasets
pip install huggingface-hub
pip install nltk

导入必要的库

import nltk
import random
import logging

import tensorflow as tf
from tensorflow import keras

nltk.download("punkt")
# 仅记录错误消息
tf.get_logger().setLevel(logging.ERROR)
# 设置随机种子
tf.keras.utils.set_random_seed(42)
[nltk_data] 下载包punkt到/speech/sreyan/nltk_data...
[nltk_data]   包punkt已经是最新的!

定义某些变量

TOKENIZER_BATCH_SIZE = 256  # 训练tokenizer的批大小
TOKENIZER_VOCABULARY = 25000  # tokenizer可以拥有的唯一子词的总数

BLOCK_SIZE = 128  # 输入样本中最大token数量
NSP_PROB = 0.50  # 下一句是实际下一句的概率
SHORT_SEQ_PROB = 0.1  # 生成更短序列的概率,以减少预训练与微调之间的不匹配
MAX_LENGTH = 512  # 填充后输入样本中的最大token数量

MLM_PROB = 0.2  # 在MLM中token被遮蔽的概率

TRAIN_BATCH_SIZE = 2  # 预训练模型的批大小
MAX_EPOCHS = 1  # 训练模型的最大轮数
LEARNING_RATE = 1e-4  # 训练模型的学习率

MODEL_CHECKPOINT = "bert-base-cased"  # 来自🤗 Model Hub的预训练模型名称

加载 WikiText 数据集

我们现在下载 WikiText 语言建模数据集。它是从维基百科的经过验证的“优秀”和“特色”文章集中提取的超过 1 亿个标记的集合。

我们从 🤗 Datasets 加载数据集。为了在本笔记本中演示,我们仅使用数据集的 train 切分。这可以通过 load_dataset 函数轻松完成。

from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
正在下载和准备数据集 wikitext/wikitext-2-raw-v1(下载:4.50 MiB,生成:12.90 MiB,后处理:未知大小,总计:17.40 MiB)到 /speech/sreyan/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126...

正在下载数据:   0%|          | 0.00/4.72M [00:00<?, ?B/s]

生成测试切分:   0%|          | 0/4358 [00:00<?, ? examples/s]

生成训练切分:   0%|          | 0/36718 [00:00<?, ? examples/s]

生成验证切分:   0%|          | 0/3760 [00:00<?, ? examples/s]

数据集 wikitext 已下载并准备好至 /speech/sreyan/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126。后续调用将重用此数据。

  0%|          | 0/3 [00:00<?, ?it/s]

该数据集仅包含一列,即原始文本,而这正是我们进行 BERT 预训练所需的所有内容!

print(dataset)
DatasetDict({
    test: Dataset({
        features: ['text'],
        num_rows: 4358
    })
    train: Dataset({
        features: ['text'],
        num_rows: 36718
    })
    validation: Dataset({
        features: ['text'],
        num_rows: 3760
    })
})

训练新的分词器

首先,我们从头开始在我们的语料库上训练自己的分词器,以便能够使用它从头训练我们的语言模型。

但为什么需要训练一个分词器呢?这是因为 Transformer 模型通常使用子词分词算法,且需要训练以识别在你使用的语料库中经常出现的单词部分。

🤗 Transformers Tokenizer(顾名思义)将对输入进行分词(包括将令牌转换为预训练词汇中的对应 ID)并将其放入模型所需的格式中,同时生成模型所需的其他输入。

首先,我们列出 WikiText 语料库中的所有原始文档:

all_texts = [
    doc for doc in dataset["train"]["text"] if len(doc) > 0 and not doc.startswith(" =")
]

接下来,我们创建一个 batch_iterator 函数,以帮助我们训练分词器。

def batch_iterator():
    for i in range(0, len(all_texts), TOKENIZER_BATCH_SIZE):
        yield all_texts[i : i + TOKENIZER_BATCH_SIZE]

在本笔记本中,我们使用与现有分词器完全相同的算法和参数训练一个分词器。例如,我们在 Wikitext-2 上使用相同的分词算法训练一个新的 BERT-CASED 分词器。

首先,我们需要加载我们想要作为模型使用的分词器:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
Transformers v4.22.0 的模型文件缓存已更新。正在迁移您的旧缓存。这仅仅是一项一次性操作。您可以中断此操作并稍后通过调用 `transformers.utils.move_cache()` 恢复迁移。

正在将 52 个文件移动到新缓存系统

  0%|          | 0/52 [00:00<?, ?it/s]

vocab_file vocab.txt
tokenizer_file tokenizer.json
added_tokens_file added_tokens.json
special_tokens_map_file special_tokens_map.json
tokenizer_config_file tokenizer_config.json

现在,我们使用整个 Wikitext-2 数据集的 train 切分来训练我们的分词器。

tokenizer = tokenizer.train_new_from_iterator(
    batch_iterator(), vocab_size=TOKENIZER_VOCABULARY
)

所以现在我们已经完成训练新的分词器!接下来我们进入数据预处理步骤。


数据预处理

为了演示工作流程,在本笔记本中我们仅取整个 WikiText traintest 切分的小子集。

dataset["train"] = dataset["train"].select([i for i in range(1000)])
dataset["validation"] = dataset["validation"].select([i for i in range(1000)])

在我们将这些文本输入到模型之前,我们需要对它们进行预处理,并为任务做好准备。如前所述,BERT 预训练任务总共包含两个任务,即 NSP 任务和 MLM 任务。🤗 Transformers 有一个容易实现的 collator,称为 DataCollatorForLanguageModeling。然而,我们需要手动为 NSP 准备数据。 接下来我们编写一个简单的函数,称为 prepare_train_features,它帮助我们进行预处理,并与 🤗 Datasets 兼容。总结一下,我们的预处理函数应该:

  • 准备好 NSP 任务的数据集,通过创建句子对 (A,B),其中 B 要么确实跟在 A 后面,要么 B 是从语料库中的其他地方随机抽样的。它还应该为每对生成一个相应的标签,如果 B 实际上跟在 A 后面则为 1,如果不跟则为 0。
  • 将文本数据集标记化为相应的令牌 id,这些令牌 id 将用于 BERT 的嵌入查找。
  • 为模型创建额外的输入,如 token_type_idsattention_mask 等。

我们定义了经过标记化后每个训练样本的最大 token 数量

max_num_tokens = BLOCK_SIZE - tokenizer.num_special_tokens_to_add(pair=True)

def prepare_train_features(examples):

"""准备 NSP 任务的特征函数

参数:
  examples: 一个字典,包含 1 个键 ("text")
    text: 原始文档列表 (str)
返回:
  examples: 一个包含 4 个键的字典
    input_ids: 来自单个原始文档的 tokenized、连接和批处理
      句子的列表 (int)
    token_type_ids: 对应的整数列表 (0 或 1)
      : 0 表示句子编号 1 和填充,1 表示句子
      编号 2
    attention_mask: 对应的整数列表 (0 或 1)
      : 1 表示非填充 token,0 表示填充
    next_sentence_label: 对应的整数列表 (0 或 1)
      : 1 如果第二个句子实际上跟在第一个句子后面,
      0 如果句子是从语料库中的其他地方采样的
"""

# 从训练集中删除不需要的样本
examples["document"] = [
    d.strip() for d in examples["text"] if len(d) > 0 and not d.startswith(" =")
]
# 将数据集中的文档拆分为单个句子
examples["sentences"] = [
    nltk.tokenize.sent_tokenize(document) for document in examples["document"]
]
# 使用已训练的分词器将 token 转换为 ids
examples["tokenized_sentences"] = [
    [tokenizer.convert_tokens_to_ids(tokenizer.tokenize(sent)) for sent in doc]
    for doc in examples["sentences"]
]

# 定义输出
examples["input_ids"] = []
examples["token_type_ids"] = []
examples["attention_mask"] = []
examples["next_sentence_label"] = []

for doc_index, document in enumerate(examples["tokenized_sentences"]):

    current_chunk = []  # 存储当前工作段的缓冲区
    current_length = 0
    i = 0

    # 我们 *通常* 希望填满整个序列,因为我们无论如何都会填充到 `block_size`,
    # 所以短序列通常会浪费计算。不过,我们 *有时*
    # (即,short_seq_prob == 0.1 == 10% 的时间) 想使用更短的
    # 序列以最小化预训练和微调之间的不匹配。
    # `target_seq_length` 只是一个粗略的目标,而
    # `block_size` 是一个硬限制。
    target_seq_length = max_num_tokens

    if random.random() < SHORT_SEQ_PROB:
        target_seq_length = random.randint(2, max_num_tokens)

    while i < len(document):
        segment = document[i]
        current_chunk.append(segment)
        current_length += len(segment)
        if i == len(document) - 1 or current_length >= target_seq_length:
            if current_chunk:
                # `a_end` 是从 `current_chunk` 中有多少段进入了 `A`
                # (第一个) 句子。
                a_end = 1
                if len(current_chunk) >= 2:
                    a_end = random.randint(1, len(current_chunk) - 1)

                tokens_a = []
                for j in range(a_end):
                    tokens_a.extend(current_chunk[j])

                tokens_b = []

                if len(current_chunk) == 1 or random.random() < NSP_PROB:
                    is_random_next = True
                    target_b_length = target_seq_length - len(tokens_a)

                    # 对于大语料库,这种情况通常不会超过一个迭代。
                    # 但是,为了小心,我们尝试确保
                    # 随机文档与我们正在处理的文档不同。
                    for _ in range(10):
                        random_document_index = random.randint(
                            0, len(examples["tokenized_sentences"]) - 1
                        )
                        if random_document_index != doc_index:
                            break

                    random_document = examples["tokenized_sentences"][
                        random_document_index
                    ]
                    random_start = random.randint(0, len(random_document) - 1)
                    for j in range(random_start, len(random_document)):
                        tokens_b.extend(random_document[j])
                        if len(tokens_b) >= target_b_length:
                            break
                    # 我们实际上并没有使用这些段,所以我们 "放回" 它们,以防它们浪费。
                    num_unused_segments = len(current_chunk) - a_end
                    i -= num_unused_segments
                else:
                    is_random_next = False
                    for j in range(a_end, len(current_chunk)):
                        tokens_b.extend(current_chunk[j])

                input_ids = tokenizer.build_inputs_with_special_tokens(
                    tokens_a, tokens_b
                )
                # 添加 token 类型 id,句子 a 为 0,句子 b 为 1
                token_type_ids = tokenizer.create_token_type_ids_from_sequences(
                    tokens_a, tokens_b
                )

                padded = tokenizer.pad(
                    {"input_ids": input_ids, "token_type_ids": token_type_ids},
                    padding="max_length",
                    max_length=MAX_LENGTH,
                )

                examples["input_ids"].append(padded["input_ids"])
                examples["token_type_ids"].append(padded["token_type_ids"])
                examples["attention_mask"].append(padded["attention_mask"])
                examples["next_sentence_label"].append(1 if is_random_next else 0)
                current_chunk = []
                current_length = 0
        i += 1

# 我们删除数据集中所有不必要的列
del examples["document"]
del examples["sentences"]
del examples["text"]
del examples["tokenized_sentences"]

return examples

tokenized_dataset = dataset.map( prepare_train_features, batched=True, remove_columns=["text"], num_proc=1, )

参数 'function'=<function prepare_train_features at 0x7fd4a214cb90> 无法正确哈希,使用了随机哈希。确保您的转换和参数可序列化以便数据集指纹和缓存功能正常工作。如果您重用此转换,缓存机制会将其视为与之前调用不同,并重新计算所有内容。此警告仅显示一次。后续的哈希失败不会显示。

  0%|          | 0/5 [00:00<?, ?ba/s]

  0%|          | 0/1 [00:00<?, ?ba/s]

  0%|          | 0/1 [00:00<?, ?ba/s]

对于MLM,我们将对数据集使用与之前相同的预处理,增加一个步骤:我们随机掩盖一些标记(通过将它们替换为[MASK]),标签将只调整为包括被掩盖的标记(我们不需要预测未掩盖的标记)。如果您使用的是自己训练的分词器,请确保[MASK]标记是在您训练时传递的特殊标记之一!

为了准备数据以用于MLM,我们简单地使用🤗 Transformers库提供的DataCollatorForLanguageModeling作为collator,该数据集已经为NSP任务做好了准备。collator期望某些参数。我们在此笔记本中使用了来自原始BERT论文的默认参数。return_tensors='tf'确保我们获得tf.Tensor对象。

from transformers import DataCollatorForLanguageModeling

collater = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=MLM_PROB, return_tensors="tf"
)

接下来,我们定义了我们的训练集来训练模型。同样,🤗 Datasets为我们提供了to_tf_dataset方法,这将帮助我们将数据集与上述定义的collator集成。该方法期望某些参数:

  • columns: 将作为自变量的列
  • label_cols: 将作为标签或因变量的列
  • batch_size: 训练时的批次大小
  • shuffle: 是否希望打乱训练数据集
  • collate_fn: 我们的汇总功能
train = tokenized_dataset["train"].to_tf_dataset(
    columns=["input_ids", "token_type_ids", "attention_mask"],
    label_cols=["labels", "next_sentence_label"],
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=True,
    collate_fn=collater,
)

validation = tokenized_dataset["validation"].to_tf_dataset(
    columns=["input_ids", "token_type_ids", "attention_mask"],
    label_cols=["labels", "next_sentence_label"],
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=True,
    collate_fn=collater,
)

定义模型

要定义我们的模型,首先需要定义一个配置,它将帮助我们定义模型架构的某些参数。这包括变压器层数、注意力头数、隐藏维度等参数。在此笔记本中,我们尝试定义原始BERT论文中定义的确切配置。

我们可以轻松地使用🤗 Transformers库中的BertConfig类来实现这一点。from_pretrained()方法期望一个模型的名称。这里我们定义最简单的模型,即我们也训练了的模型,bert-base-cased

from transformers import BertConfig

config = BertConfig.from_pretrained(MODEL_CHECKPOINT)

为了定义我们的模型,我们使用🤗 Transformers库中的TFBertForPreTraining类。该类内部处理一切,从定义我们的模型到解包我们的输入和计算损失。因此,我们只需用我们想要的正确config定义模型即可!

from transformers import TFBertForPreTraining

model = TFBertForPreTraining(config)

现在我们定义优化器并编译模型。损失计算在内部处理,因此我们无需担心!

optimizer = keras.optimizers.Adam(learning_rate=LEARNING_RATE)

model.compile(optimizer=optimizer)
在compile()中没有指定损失 - 模型将使用其内部损失计算作为损失。不要恐慌 - 这是在Transformers中训练TensorFlow模型的常见方法!要禁用此行为,请传递损失参数,或者如果您不希望模型计算损失,请明确传递`loss=None`。

最后,所有步骤都完成了,现在我们可以开始训练模型!

model.fit(train, validation_data=validation, epochs=MAX_EPOCHS)
483/483 [==============================] - 96s 141ms/step - loss: 8.3765 - val_loss: 8.5572

<keras.callbacks.History at 0x7fd27c219790>

我们的模型已经训练完成!我们建议请在完整数据集上训练模型至少50个周期以获得良好性能。预训练模型现在作为 一个语言模型,旨在进行下游任务的微调。因此,它现在可以在任何下游任务上进行微调,例如问答、文本分类等!

现在你可以将这个模型推送到 🤗 模型库,并与你的朋友、家人、宠物分享:他们都可以加载该模型,使用标识符 "your-username/the-name-you-picked",例如:

model.push_to_hub("pretrained-bert", organization="keras-io")
tokenizer.push_to_hub("pretrained-bert", organization="keras-io")

在你推送模型后,未来可以这样加载它!

from transformers import TFBertForPreTraining

model = TFBertForPreTraining.from_pretrained("your-username/my-awesome-model")

或者,因为这是一个预训练模型,你通常会将其用于下游任务的微调,你也可以为其他任务加载它,例如:

from transformers import TFBertForSequenceClassification

model = TFBertForSequenceClassification.from_pretrained("your-username/my-awesome-model")

在这种情况下,预训练头将被丢弃,模型将仅使用变换器层进行初始化。将随机权重添加一个新的特定于任务的头。