NLP学习1 - 使用Huggingface Transformers框架从头训练语言模型

摘要

由于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