【预训练语言模型】使用Transformers库进行GPT2预训练

基于 HuggingFace的Transformer库，在Colab或Kaggle进行预训练。

本教程提供：英文数据集wikitext-2和代码数据集的预训练。
注：可以自行上传数据集进行训练

目的：跑通自回归语言模型的预训练流程

一、准备

1.1 安装依赖

!pip install -U datasets
!pip install accelerate -U

注意：在Colab上训练时，最好将datasets更新到最新版(再重启kernel)，避免版本低报错

colab和kaggle已经预安装transformers库

1.2 数据准备

加载数据

from datasets import load_dataset

datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')

当然你也可使用huggingface上任何公开的文本数据集，或使用自己构造的数据，并将路径替换为指定路径：

# datasets = load_dataset("text", data_files={"train": path_to_train.txt, "validation": path_to_validation.txt}

要访问一个数据中实际的元素，您需要先选择一个key，然后给出一个索引:
看一下数据的格式

datasets["train"][10].keys()

可以看到该数据集的每个元素就是一个仅包含文本的字典

dict_keys(['text'])

查看例子

datasets["train"][1]

{‘text': ' =Valkyria Chronicles III = \n'}

训练集和测试集数量

print(len(datasets["train"]), len(datasets["test"]))

36718 4358

通过如下的函数来随机展示数据集中的一些样本：

from datasets import ClassLabel
import random
import pandas as pd
from IPython.display import display, HTML

def show_random_elements(dataset, num_examples=10):
    assert num_examples <= len(dataset), "Can't pick more elements than there are in the dataset."
    picks = []
    for _ in range(num_examples):
        pick = random.randint(0, len(dataset)-1)
        while pick in picks:
            pick = random.randint(0, len(dataset)-1)
        picks.append(pick)
    
    df = pd.DataFrame(dataset[picks])
    for column, typ in dataset.features.items():
        if isinstance(typ, ClassLabel):
            df[column] = df[column].transform(lambda i: typ.name