GPT 非常适合生成类任务，下面使用对联语料从零训练一个用于对联生成的模型。该模型：

输入一个字，自动生成上联和下联
输入上联，自动生成下联

语料文件结构如下：

couplet
├── test
│   ├── in.txt
│   └── out.txt
└── train
    ├── in.txt
    └── out.txt

语料部分截图：

导入需要的模块：

import pandas as pd
from transformers import BertTokenizer
from datasets import Dataset

1. 语料清洗

对语料进行简单的处理：

去除一些重复的数据，一些包含了 “妓女” 敏感词的语料，
去除一些包含问号的语料，例如：”万险何辞？为圆大漠科研梦” 等
去除一些长度小于5的语料
去除一些上联和下联长度不同的语料
…等等

此时数据量从 74 万减少到近 31+ 万。我们这里使用全部 31 万数据进行训练。

def clean_corpus():

    # 读取语料
    corpus1 = open('couplet/train/in.txt')
    corpus2 = open('couplet/train/out.txt')

    # 语料清晰
    lines = []
    for line1, line2 in zip(corpus1, corpus2):

        # 去除语料多余空格
        line1 = ''.join(line1.strip().split())
        line2 = ''.join(line2.strip().split())

        # 去除包含问号语料
        if line1.find('？') != -1:
            continue

        if line1.find('，') != -1:
            continue
            
        if line1.find('、') != -1:
            continue
            
        if line1.find('！') != -1:
            continue
            
        if line1.find('：') != -1:
            continue

        # 去除包含敏感词汇
        if line1.find('妓女') != -1 or line2.find('妓女') != -1:
            continue

        if line1.find('政治') != -1 or line2.find('政治') != -1:
            continue

        # 上下联长度不对等
        if len(line1) != len(line2):
            continue

        # 长度限制语料
        # 上联一般 5、7、10 个字
        if len(line1) < 5 or len(line1) > 9:
            continue


        lines.append([line1, line2])

    # 上联去重
    lines = pd.DataFrame(data=lines, columns=['single', 'double'])
    lines = lines[~lines['single'].duplicated()]

    # 存储语料
    part_number = None
    lines[:part_number].to_csv('data/data-clean-corpus.csv', index=False)

    # 关闭文件
    corpus1.close()
    corpus2.close()

    print('1. 语料清洗完毕!', '数据总量:', len(lines), '使用总量:', 'ALL' if part_number is None else part_number)

修改上述代码中 part_number 变量可以改变使用的语料数量，None 表示使用所有的语料，执行完该函数会在 data 目录下创建 data-clean-corpus.csv 文件，该文件中存储了 part_number 条上下联数据，截图如下：

2. 构建词表

我们后面会使用到 BertTokenizer 来对语料进行编码，而构建该对象，需要一个包含了字的文本文件，该文件如下图所示：

实现代码如下：

def build_vocab():

    # 读取数据
    train_data = pd.read_csv('data/data-clean-corpus.csv')
    train_data = train_data[['single', 'double']].values.tolist()

    # 语料分词
    words = []
    for line1, line2 in train_data:
        words.extend(line1)
        words.extend(line2)

    words = list(set(words))
    words.insert(0, '[BRK]')
    words.insert(0, '[END]')
    words.insert(0, '[UNK]')
    words.insert(0, '[PAD]')

    # 词表存储
    with open('data/data-build-vocab.txt', 'w') as file:
        for word in words:
            file.write(word + '\n')

    print('2. 构建词表完毕!', '词表大小:', len(words))

注意: 上面我们添加了 4 个特殊标记，[BRK] 作为两个输入的分割标记， [END] 表示输入序列的结束标记，[PAD] 用于对齐填充，[UNK] 用于表示 OOV 词。

函数执行之后，会在 data 目录下创建 data-build-vocab.txt 文件，文件内容如下图所示:

3. 训练语料

为了能够让模型实现根据一个字生成上联和下联，也能够根据上联生成下联，我们将上联和下联拼接起来送入模型，并使用自定义的 [BRK] 标记来分割，最后以 [END] 表示结束。

晚风摇树树还挺[BRK]晨露润花花更红[END]

根据上面的语料，构建模型的输入和目标如下：

输入: 晚 风 摇 树 树 还 挺    [BRK] 晨 露 润 花 花 更 红 
目标: 风 摇 树 树 还 挺 [BRK] 晨    露 润 花 花 更 红 [END]

实现代码如下：

def build_corpus():

    # 数据
    train_data = pd.read_csv('data/data-clean-corpus.csv')
    train_data = train_data[['single', 'double']].values.tolist()

    # 词表
    BRK = '[BRK]'
    END = '[END]'

    # 语料
    train_inputs = []  # 输入
    train_labels = []  # 标签
    for line1, line2 in train_data:
        line = line1 + BRK + line2 + END
        train_inputs.append(line[:-5])
        train_labels.append(line[1:])

    train_data = pd.DataFrame()
    train_data['inputs'] = train_inputs
    train_data['labels'] = train_labels

    train_data.to_csv('data/data-build-corpus.csv')

    print('3. 训练语料构建完毕!')

函数执行之后，会在 data 目录下创建 data-build-corpus.csv 文件，该文件内容如下图所示：

4. 语料编码

为了能够语料送入 GPT 模型计算，我们需要将其转换为 id 表示。GPT Tokenizer 使用 BPE 分词，我们使用的中文，以字粒度来拆分句子。所以，我们使用 Bert Tokenizer，而不是 GPT Tokenizer 来构建分词器。

另外，由于我们添加了 [BRK]、[END] 特殊标记，为了能够让分词器识别它，需要使用 add_special_tokens 函数将其添加到对象中。

另外，需要注意的是：

前面我们已经在构建训练语料时，添加了需要的特殊标记，所以在此处使用 encode 编码时，需要设置 add_special_tokens=False；
输入的语料是批次送入模型计算，需要对其进行填充补齐，而目标值则不需要填充，这点额外注意。

完整的实现代码如下：

def encode_corpus():

    # 标记器
    tokenizer = BertTokenizer('data/data-build-vocab.txt', pad_token='[PAD]')
    tokenizer.add_special_tokens({'additional_special_tokens': ['[END]', '[BRK]']})
    tokenizer.save_pretrained('data/tokenizer-encode-tokenizer')

    # 读取数据
    data = pd.read_csv('data/data-build-corpus.csv', )
    data = Dataset.from_pandas(data)


    def handle(inputs, labels):
        encode_inputs = []
        encode_labels = []
        for label, input in zip(labels, inputs):
            encode_labels.append(tokenizer.encode(label, add_special_tokens=False))
            encode_inputs.append(tokenizer.encode(input, add_special_tokens=False,
                                                  max_length=20,
                                                  padding='max_length',
                                                  truncation=True))

        return {'labels': encode_labels, 'inputs': encode_inputs}

    # 标记语料
    data = data.shuffle(seed=42).map(lambda x: handle(x['inputs'], x['labels']),
                                     batched=True,
                                     batch_size=1000, desc='开始训练语料编码')
    # 存储语料
    train_data = data.remove_columns(['Unnamed: 0'])
    train_data.save_to_disk('data/data-encode-corpus')

    print('4. 训练语料编码完毕!')

上述函数执行之后，会在 data 目录下创建 data-encode-corpus 目录，该目录结构如下：

data/data-encode-corpus
├── dataset.arrow
├── dataset_info.json
└── state.json

下次使用该语料时候，只需要 datasets.load_from_disk 加载即可拿到 Dataset 对象。

5. 入口函数

def start():

    # 1. 语料清洗
    clean_corpus()
    # 2. 构建词表
    build_vocab()
    # 3. 训练语料
    build_corpus()
    # 4. 语料编码
    encode_corpus()


if __name__ == '__main__':
    start()

执行该函数之后，将会处理 31万+ 数据，大概需要 2 分左右。会得到以下输出：

1. 语料清洗完毕! 数据总量: 312831 使用总量: ALL
2. 构建词表完毕! 词表大小: 7397
3. 训练语料构建完毕!
开始训练语料编码: 100%|███████████████████████| 313/313 [02:01<00:00,  2.58ba/s]
4. 训练语料编码完毕!

基于 GPT 训练对联模型 – 数据处理

1. 语料清洗

2. 构建词表

3. 训练语料

4. 语料编码

5. 入口函数

文章目录