使用 Hugging Face Transformers 預訓練 BERT
作者： Sreyan Ghosh
建立日期 2022/07/01
上次修改日期 2022/08/27
描述： 使用 Hugging Face Transformers 在 NSP 和 MLM 上預訓練 BERT。
簡介
BERT (Bidirectional Encoder Representations from Transformers，來自 Transformer 的雙向編碼器表示)
在電腦視覺領域，研究人員已多次展示遷移學習的價值 - 在已知的任務/資料集上預訓練神經網路模型，例如 ImageNet 分類，然後執行微調 - 使用訓練的神經網路作為新的特定用途模型基礎。近年來，研究人員已表明，類似的技術在許多自然語言任務中也很有用。
BERT 利用 Transformer，一種學習文字中單詞（或子詞）之間上下文關係的注意力機制。在其原始形式中，Transformer 包含兩種不同的機制 - 一種讀取文字輸入的編碼器，以及一種產生任務預測的解碼器。由於 BERT 的目標是產生語言模型，因此只需要編碼器機制。Transformer 的詳細工作原理在 Google 的一篇論文中描述。
與按順序（由左至右或由右至左）讀取文字輸入的方向模型相反，Transformer 編碼器一次讀取整個單詞序列。因此，它被認為是雙向的，儘管更準確地說它是非方向的。此特性允許模型根據其所有周圍環境（單詞的左邊和右邊）來學習單詞的上下文。
訓練語言模型時，一個挑戰是定義預測目標。許多模型預測序列中的下一個單詞（例如 "The child came home from _"），這是一種方向方法，會固有地限制上下文學習。為了克服這個挑戰，BERT 使用兩種訓練策略
遮罩語言建模 (Masked Language Modeling, MLM)
在將單詞序列輸入 BERT 之前，每個序列中有 15% 的單詞會被替換為 [MASK] 標記。然後，模型會嘗試根據序列中其他未遮罩的單詞提供的上下文來預測遮罩單詞的原始值。
下一個句子預測 (Next Sentence Prediction, NSP)
在 BERT 訓練過程中，模型接收成對的句子作為輸入，並學習預測配對中的第二個句子是否是原始文件中隨後的句子。在訓練期間，50% 的輸入是一對，其中第二個句子是原始文件中隨後的句子，而在其他 50% 中，則從語料庫中隨機選擇一個句子作為第二個句子。假設隨機句子將代表與第一個句子的斷連。
儘管 Google 為英語提供了預訓練的 BERT 檢查點，但您可能經常需要從頭開始預訓練模型以適應不同的語言，或者執行持續預訓練以使模型適應新的領域。在本筆記本中，我們從頭開始預訓練 BERT，使用 🤗 Transformers 在從 🤗 Datasets 載入的 WikiText 英語資料集上優化 MLM 和 NSP 目標。
設定
安裝需求
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")
# Only log error messages
tf.get_logger().setLevel(logging.ERROR)
# Set random seed
tf.keras.utils.set_random_seed(42)
[nltk_data] Downloading package punkt to /speech/sreyan/nltk_data...
[nltk_data] Package punkt is already up-to-date!
定義某些變數
TOKENIZER_BATCH_SIZE = 256 # Batch-size to train the tokenizer on
TOKENIZER_VOCABULARY = 25000 # Total number of unique subwords the tokenizer can have
BLOCK_SIZE = 128 # Maximum number of tokens in an input sample
NSP_PROB = 0.50 # Probability that the next sentence is the actual next sentence in NSP
SHORT_SEQ_PROB = 0.1 # Probability of generating shorter sequences to minimize the mismatch between pretraining and fine-tuning.
MAX_LENGTH = 512 # Maximum number of tokens in an input sample after padding
MLM_PROB = 0.2 # Probability with which tokens are masked in MLM
TRAIN_BATCH_SIZE = 2 # Batch-size for pretraining the model on
MAX_EPOCHS = 1 # Maximum number of epochs to train the model for
LEARNING_RATE = 1e-4 # Learning rate for training the model
MODEL_CHECKPOINT = "bert-base-cased" # Name of pretrained model from 🤗 Model Hub
載入 WikiText 資料集
我們現在下載 WikiText 語言建模資料集。它是從 Wikipedia 上一組經過驗證的「良好」和「精選」文章中提取的超過 1 億個標記的集合。
我們從 🤗 Datasets 載入資料集。為了在本筆記本中進行示範，我們僅使用資料集的 train 分割。這可以使用 load_dataset 函數輕鬆完成。
from datasets import load_dataset
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
Downloading and preparing dataset wikitext/wikitext-2-raw-v1 (download: 4.50 MiB, generated: 12.90 MiB, post-processed: Unknown size, total: 17.40 MiB) to /speech/sreyan/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126...
Downloading data: 0%| | 0.00/4.72M [00:00<?, ?B/s]
Generating test split: 0%| | 0/4358 [00:00<?, ? examples/s]
Generating train split: 0%| | 0/36718 [00:00<?, ? examples/s]
Generating validation split: 0%| | 0/3760 [00:00<?, ? examples/s]
Dataset wikitext downloaded and prepared to /speech/sreyan/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/a241db52902eaf2c6aa732210bead40c090019a499ceb13bcbfa3f8ab646a126. Subsequent calls will reuse this data.
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
})
})
訓練新的 Tokenizer
首先，我們從頭開始在我們的語料庫上訓練我們自己的 tokenizer，以便我們可以使用它從頭開始訓練我們的語言模型。
但是，為什麼您需要訓練 tokenizer 呢？這是因為 Transformer 模型經常使用子詞標記化演算法，並且需要訓練它們來識別您正在使用的語料庫中經常出現的單詞部分。
🤗 Transformers Tokenizer（顧名思義）將標記輸入（包括將標記轉換為預訓練詞彙表中對應的 ID），並將其放入模型期望的格式中，以及產生模型需要的其他輸入。
首先，我們從 WikiText 語料庫中建立所有原始文件的清單
all_texts = [
doc for doc in dataset["train"]["text"] if len(doc) > 0 and not doc.startswith(" =")
]
接下來，我們建立一個 batch_iterator 函數，它將幫助我們訓練 tokenizer。
def batch_iterator():
for i in range(0, len(all_texts), TOKENIZER_BATCH_SIZE):
yield all_texts[i : i + TOKENIZER_BATCH_SIZE]
在本筆記本中，我們使用與現有 tokenizer 完全相同的演算法和參數訓練 tokenizer。例如，我們使用相同的標記化演算法在 Wikitext-2 上訓練 BERT-CASED tokenizer 的新版本。
首先，我們需要載入我們想要用作模型的 tokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.
Moving 52 files to the new cache system
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 = tokenizer.train_new_from_iterator(
batch_iterator(), vocab_size=TOKENIZER_VOCABULARY
)
現在我們完成了訓練新的 tokenizer！接下來，我們繼續進行資料預處理步驟。
資料預處理
為了示範工作流程，在本筆記本中，我們只取整個 WikiText train 和 test 分割的小子集。
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 相容。總而言之，我們的預處理函數應執行以下操作
- 透過建立成對的句子 (A,B) 來準備用於 NSP 任務的資料集，其中 B 實際上是 A 的後續句子，或者 B 是從語料庫中的其他地方隨機採樣的。它還應為每一對產生對應的標籤，如果 B 實際上是 A 的後續句子，則為 1，如果不是則為 0。
- 將文字資料集標記為其對應的標記 ID，這些 ID 將用於 BERT 中的嵌入查詢
- 為模型建立額外的輸入，例如
token_type_ids、attention_mask等。
# We define the maximum number of tokens after tokenization that each training sample
# will have
max_num_tokens = BLOCK_SIZE - tokenizer.num_special_tokens_to_add(pair=True)
def prepare_train_features(examples):
"""Function to prepare features for NSP task
Arguments:
examples: A dictionary with 1 key ("text")
text: List of raw documents (str)
Returns:
examples: A dictionary with 4 keys
input_ids: List of tokenized, concatnated, and batched
sentences from the individual raw documents (int)
token_type_ids: List of integers (0 or 1) corresponding
to: 0 for senetence no. 1 and padding, 1 for sentence
no. 2
attention_mask: List of integers (0 or 1) corresponding
to: 1 for non-padded tokens, 0 for padded
next_sentence_label: List of integers (0 or 1) corresponding
to: 1 if the second sentence actually follows the first,
0 if the senetence is sampled from somewhere else in the corpus
"""
# Remove un-wanted samples from the training set
examples["document"] = [
d.strip() for d in examples["text"] if len(d) > 0 and not d.startswith(" =")
]
# Split the documents from the dataset into it's individual sentences
examples["sentences"] = [
nltk.tokenize.sent_tokenize(document) for document in examples["document"]
]
# Convert the tokens into ids using the trained tokenizer
examples["tokenized_sentences"] = [
[tokenizer.convert_tokens_to_ids(tokenizer.tokenize(sent)) for sent in doc]
for doc in examples["sentences"]
]
# Define the outputs
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 = [] # a buffer stored current working segments
current_length = 0
i = 0
# We *usually* want to fill up the entire sequence since we are padding
# to `block_size` anyways, so short sequences are generally wasted
# computation. However, we *sometimes*
# (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
# sequences to minimize the mismatch between pretraining and fine-tuning.
# The `target_seq_length` is just a rough target however, whereas
# `block_size` is a hard limit.
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` is how many segments from `current_chunk` go into the `A`
# (first) sentence.
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)
# This should rarely go for more than one iteration for large
# corpora. However, just to be careful, we try to make sure that
# the random document is not the same as the document
# we're processing.
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
# We didn't actually use these segments so we "put them back" so
# they don't go to waste.
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
)
# add token type ids, 0 for sentence a, 1 for sentence b
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
# We delete all the un-necessary columns from our dataset
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,
)
Parameter 'function'=<function prepare_train_features at 0x7fd4a214cb90> of the transform datasets.arrow_dataset.Dataset._map_single couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.
0%| | 0/5 [00:00<?, ?ba/s]
0%| | 0/1 [00:00<?, ?ba/s]
0%| | 0/1 [00:00<?, ?ba/s]
對於 MLM，我們將對我們的資料集使用與之前相同的預處理，但有一個額外的步驟：我們隨機遮罩一些標記（將其替換為 [MASK]），並且標籤將被調整為僅包含遮罩的標記（我們不必預測未遮罩的標記）。如果您使用自己訓練的 tokenizer，請確保 [MASK] 標記在您在訓練期間傳遞的特殊標記中！
為了準備用於 MLM 的資料，我們只需在已經為 NSP 任務準備好的資料集上使用 🤗 Transformers 程式庫提供的 collator，稱為 DataCollatorForLanguageModeling。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：我們的 collator 函數
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,
)
定義模型
為了定義我們的模型，首先我們需要定義一個設定檔，該設定檔將幫助我們定義模型架構的某些參數。這包括諸如 Transformer 層數、注意力頭數、隱藏維度等參數。對於本筆記本，我們嘗試定義原始 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)
No loss specified in compile() - the model's internal loss computation will be used as the loss. Don't panic - this is a common way to train TensorFlow models in Transformers! To disable this behaviour please pass a loss argument, or explicitly pass `loss=None` if you do not want your model to compute a loss.
最後，所有步驟都已完成，現在我們可以開始訓練我們的模型了！
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 個 epoch，以獲得良好的效能。預訓練模型現在充當語言模型，目的是在下游任務上進行微調。因此，它現在可以在任何下游任務上進行微調，例如問答、文本分類等等！
現在您可以將此模型推送到 🤗 Model Hub，並與您所有的朋友、家人、最喜歡的寵物分享：他們都可以使用識別碼 "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")
在這種情況下，預訓練頭將被丟棄，並且模型將僅使用 transformer 層進行初始化。將會添加一個具有隨機權重的新任務特定頭。