摘要
由于huaggingface放出了Tokenizers工具，结合之前的transformers，因此预训练模型就变得非常的容易，本文以学习官方example为目的，由于huggingface目前给出的run_language_modeling.py中尚未集成Albert（目前有 GPT, GPT-2, BERT, DistilBERT and RoBERTa，具体可以点开前面的链接），这是由于目前对于Albert的支持，在分词时，AlbertTokenizer仅仅支持SentencePiece file 类型的分词，所以不能直接用Tokenizers分词。实际实验时，发现可以直接用Tokenizers.BertWordPieceTokenizer进行分词。因此，本文在example的基础上，对中文文本《论语》进行预训练，使用Tokenizers和Transformers，环境在Google Colab GPU下。
训练Tokenizer
本文选择训练一个BertWordPieceTokenizer的分词器，由于Bert和Albert大致相似，因此分词器上选择BertWordPieceTokenizer不会有问题。首先在colab装一下包：
# We won't need TensorFlow here
!pip uninstall -y tensorflow
# Install `transformers` from master
!pip install git+https://github.com/huggingface/transformers
!pip list | grep -E 'transformers|tokenizers'
# transformers version at notebook update --- 2.9.1
# tokenizers version at notebook update --- 0.7.0
接着，建立分词器：
%%time
from tokenizers import BertWordPieceTokenizer
files = "./lunyu.txt" # 训练文本文件
vocab_size = 10000
min_frequency = 2
limit_alphabet = 10000
special_tokens = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] #适用于Bert和Albert
# Initialize a tokenizer
tokenizer = BertWordPieceTokenizer(
clean_text=True, handle_chinese_chars=True, strip_accents=True, lowercase=True,
)
# Customize training
tokenizer.train(
files,
vocab_size = vocab_size,
min_frequency=min_frequency,
show_progress=True,
special_tokens=special_tokens,
limit_alphabet=limit_alphabet,
wordpieces_prefix="##"
)
然后把分词器保存在硬盘中：
!mkdir tokenizer
tokenizer.save("tokenizer")
这样，我们就拥有了一个针对《论语》文本的分词器，这会对于论语类文本的应用更加合适，这里还可以利用tokenizers包来测试一下分词效果：
from tokenizers.implementations import BertWordPieceTokenizer
from tokenizers.processors import BertProcessing
tokenizer = BertWordPieceTokenizer(
"./tokenizer/vocab.txt",
)
tokenizer._tokenizer.post_processor = BertProcessing(
("[CLS]", tokenizer.token_to_id("[SEP]")),
("[SEP]", tokenizer.token_to_id("[CLS]")),
)
tokenizer.enable_truncation(max_length=512)
tokenizer.encode("子曰：学而时习之。").tokens
这样，可以对子曰：学而时习之。这句话进行分词，由于我们用的是BertWordPieceTokenizer，对于中文来说，就是对每一个字进行分词。得到结果：
['[SEP]', '子', '曰', '：', '学', '而', '时', '习', '之', '。', '[CLS]']
从头训练一个语言模型
我们的目标是训练一个Albert模型，因为Albert模型与Bert相似，因此可以进行MLM（ Masked language modeling）的任务，首先，我们需要定义一下模型的配置：
from transformers import AlbertConfig
config = AlbertConfig(
vocab_size = 1359,
embedding_size = 256,
hidden_size = 768,
num_hidden_layers = 6,
num_attention_heads = 12,
intermediate_size = 3072,
hidden_act = "gelu",
hidden_dropout_prob = 0.1,
attention_probs_dropout_prob = 0.1,
max_position_embeddings = 512,
type_vocab_size = 2,
initializer_range = 0.02,
layer_norm_eps = 1e-12,
)
这里比较重要的超参数有vocab_size, embedding_size, hidden_size, num_hidden_layers, num_attention_heads涉及到网络的size，由于我也是个菜鸡，等下次学了再仔细写写，这里是抄的。
然后，我们还需要在Transformers中重新建立一个分词器：
from transformers import BertTokenizerFast
tokenizer = BertTokenizerFast.from_pretrained("./tokenizer", max_len=512)
接下来我们需要初始化我们的albert模型，这里需要注意的是，在配置模型参数时，因为我们从头开始训练，所以我们只从配置中进行初始化，而不需要从checkponit文件中初始化。
from transformers import AlbertForMaskedLM
model = AlbertForMaskedLM(config=config)
model.num_parameters()
# => 8554575个参数
接着，我们将使用建立好的分词器对文本数据进行分词，构建训练所需要的数据集。这里，由于我们只有一个文本文件，我们甚至不需要定制我们的数据集，直接使用LineByLineDataset。（还没仔细研究这个是干嘛的）
%%time
from transformers import LineByLineTextDataset
dataset = LineByLineTextDataset(
tokenizer=tokenizer,
file_path="./lunyu.txt",
block_size=256,
)
然后，我们需要定义一个data_collator，这只是一个helper，它将帮助我们将数据集的不同sample批处理到一个object中，PyTorch就知道如何在这个object上执行backprop。（这里没有看懂，大概知道是转换成Pytorch能够处理的数据格式）
from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)
所有配置完成后，我们就可以初始化我们的模型了：
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="./lunyuAlbert",
overwrite_output_dir=True,
num_train_epochs=20,
per_gpu_train_batch_size=16,
save_steps=2000,
save_total_limit=2,
)
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=dataset,
prediction_loss_only=True,
)
最后，跑起来就完事儿了：
%%time
trainer.train()
由于文本数据比较小，所以结果还挺快的：
CPU times: user 2min 25s, sys: 1min 22s, total: 3min 47s
Wall time: 3min 47s
TrainOutput(global_step=680, training_loss=5.1874437787953545)
这里不评价好坏，只是学习一下如何使用Huggingface的Tokenizers包和Transformers包对Albert进行预训练。
参考：
Github的个人集成脚本：easy_bert_pretrain
huggingface的官方example：01_how_to_train.ipynb