基于 PEGASUS 生成中文文本摘要

预训练模型好不好，最终得落到自己的数据上才看得出来。PEGASUS 预训练时练的就是读全文写概括，但那是在海量英文网页和新闻上练的，真要拿来处理中文新闻短文本，还得拿中文语料继续微调。这一节我们走一遍完整流程，用 IDEA 开源的中文 Pegasus 预训练权重，在 LCSTS 中文短文本摘要数据集上继续微调，最后跑出一条输入新闻、输出一句话摘要的完整链路。这里要感谢 IDEA-CCNL 在 https://huggingface.co/IDEA-CCNL 开源的预训练模型。

整个流程分四块，先搭数据输入，再写训练循环，然后做评估，最后单独喂一条新闻生成摘要看效果。读之前需要会用 PyTorch 写普通训练循环，知道 DataLoader 和交叉熵怎么回事，也明白 Pegasus 预训练大致在做什么。

1. 构造输入

先把数据长什么样看清楚，输入怎么拼自然就知道了。我们用的是 LCSTS，一个从微博里抓出来的中文短文本加对应摘要的数据集，目录结构如下。

LCSTS
|-- test.src.txt
|-- test.tgt.txt
|-- train.src.txt
|-- train.tgt.txt
|-- valid.src.txt
|-- valid.tgt.txt

从文件名能看出来，数据分成训练集、验证集、测试集三份，每份各有一个源文件和一个目标文件。训练集里两行是这样对应的。

train.src.txt 是正文。

新华社受权于18日全文播发修改后的《中华人民共和国立法法》，修改后的立法法分为“总则”“法律”“行政法规”“地方性法规、自治条例和单行条例、规章”“适用与备案审查”“附则”等6章，共计105条。
一辆小轿车，一名女司机，竟造成9死24伤。日前，深圳市交警局对事故进行通报：从目前证据看，事故系司机超速行驶且操作不当导致。目前24名伤员已有6名治愈出院，其余正接受治疗，预计事故赔偿费或超一千万元。
...

train.tgt.txt 是对应的人工摘要。

修改后的立法法全文公布
深圳机场9死24伤续：司机全责赔偿或超千万
...

正文每一行短文，对应目标文件里一行摘要。Pegasus 是编码器解码器结构，训练时本来要分别准备编码器输入和解码器输入。好在用 transformers 库的 PegasusForConditionalGeneration 时，我们只需要喂编码器输入和标签，解码器输入由库内部根据标签自动构造。它的做法是把标签整体往右挪一位，首位补一个起始符，这样解码器看到的永远是正确答案的前半段，逼它去预测下一个词。这段内部逻辑长这样。

if labels is not None:
        if use_cache:
            logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
        use_cache = False
        if decoder_input_ids is None:
            decoder_input_ids = shift_tokens_right(
                labels, self.config.pad_token_id, self.config.decoder_start_token_id
            )


# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
    """
    Shift input ids one token to the right.
    """
    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
    shifted_input_ids[:, 0] = decoder_start_token_id

    if pad_token_id is None:
        raise ValueError("self.model.config.pad_token_id has to be defined.")
    # replace possible -100 values in labels by `pad_token_id`
    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

    return shifted_input_ids

知道了这层，我们要做的就简单了。先把两个文件按行读成字符串列表，包成一个数据集类，真正转成 token 索引留到训练时的批处理函数里做。

def data_read(type='train'):

    all_inputs, all_labels = [], []
    for inputs, labels in zip(open('LCSTS/' + type + '.src.txt'), open('LCSTS/' + type + '.tgt.txt')):
        all_inputs.append(inputs.strip())
        all_labels.append(labels.strip())

    return all_inputs, all_labels


class SummaryDataset:

    def __init__(self, inputs, labels):
        self.inputs = inputs
        self.labels = labels
        self.counts = len(self.labels)

    def __len__(self):
        return self.counts

    def __getitem__(self, index):
        return self.inputs[index], self.labels[index]

2. 训练函数

数据读进来了，下一步就是把模型装上开始训练。这里有几个环境上的坑得先说清楚，照着配才能跑通。

第一，我们从 huggingface 下载的 Randeng-Pegasus-238M-Summary-Chinese，已经是作者在摘要数据上调过一版的权重，我们做的是在它基础上继续练，不是从零训。第二，这个模型的词表被作者改过，他把 jieba 分词里的词级 token 加进了词表，编码时能把一个有意义的词整体当成一个 id，默认的 PegasusTokenizer 认不出这些。所以要把模型目录下的 data_utils.py 和 tokenizers_pegasus.py 拷到自己工程目录，跑起来如果报错，按提示把 data_utils 里报 jieba 错的那段注释掉就行。

第三是数据量和显存的现实问题。原始 LCSTS 训练集有 240 多万条，编码器加解码器一起算，训练慢，显存也扛不住。这里做了两个妥协，数据集只取前 8 万条，batch 设成 4，再用梯度累加，每攒够 32 个小批量才更新一次参数，等效 batch 是 128。照这个配置，一张 6G 显存的 RTX2060 上每轮大约跑半小时。完整训练代码如下。

def train():

    estimator = PegasusForConditionalGeneration.from_pretrained('Randeng-Pegasus-238M-Summary-Chinese').to(device)
    tokenizer = PegasusTokenizer.from_pretrained('Randeng-Pegasus-238M-Summary-Chinese')
    optimizer = optim.Adam(estimator.parameters(), lr=3e-5)
    criterion = nn.CrossEntropyLoss()

    train_inputs, train_labels = data_read(type='train')
    train_number = 80000
    train_inputs, train_labels = train_inputs[:train_number], train_labels[:train_number]
    train_data = SummaryDataset(train_inputs, train_labels)

    def collate_function(batch_data):

        encoder_inputs, decoder_inputs = [], []
        for einputs, dinputs in batch_data:
            encoder_inputs.append(einputs)
            decoder_inputs.append(dinputs)

        # 将摘要内容转换为目标值
        summary_labels = tokenizer.batch_encode_plus(decoder_inputs, return_attention_mask=False)
        summary_labels = [torch.tensor(input_ids, device=device) for input_ids in summary_labels['input_ids']]
        summary_labels = pad_sequence(summary_labels, padding_value=criterion.ignore_index, batch_first=True)

        # 将文本输入转换索引表示
        encoder_inputs = tokenizer(encoder_inputs, max_length=1024, truncation=True, padding='longest', return_tensors="pt")
        encoder_inputs = { key: value.to(device) for key, value in encoder_inputs.items()}
        encoder_inputs['labels'] = summary_labels

        return encoder_inputs

    dataloader = DataLoader(train_data, batch_size=4, shuffle=True, collate_fn=collate_function)

    for epoch in range(10):

        estimator.train()
        progress = tqdm(range(len(dataloader)))
        total_loss = 0.0
        acumulate_start = 1

        optimizer.zero_grad()
        for encoder_inputs in dataloader:
            outputs = estimator(**encoder_inputs)

            (outputs.loss/32).backward()

            acumulate_start += 1
            if acumulate_start % 32 == 0 or acumulate_start == len(dataloader):
                optimizer.step()
                optimizer.zero_grad()

            total_loss += outputs.loss.item()
            progress.set_description('epoch %d loss %8.2f' % (epoch + 1, total_loss))
            progress.update()
        progress.close()

        # 模型存储
        checkpoint = 'model/%d-pegasus-summary-%.2f' % (epoch + 1, total_loss)
        estimator.save_pretrained(checkpoint)
        tokenizer.save_pretrained(checkpoint)
        # 模型评估
        train_evaluate(estimator, tokenizer)
        print('-' * 100)

跑起来的输出长这样。8 万条除以 batch 4，正好每轮 2 万个小批量，和日志里的 20000 步对得上。

epoch 1 loss 59105.36: 100%|██████████████| 20000/20000 [31:08<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:40<00:00,  5.48it/s]
Rouge-1: 29.16
Rouge-2: 14.5
Rouge-L: 27.19
BLEU-4: 0.1161
----------------------------------------------------------------------------------------------------
epoch 2 loss 53440.37: 100%|██████████████| 20000/20000 [31:07<00:00, 10.71it/s]
100%|█████████████████████████████████████████| 553/553 [01:40<00:00,  5.50it/s]
Rouge-1: 29.4
Rouge-2: 14.85
Rouge-L: 27.54
BLEU-4: 0.1185
----------------------------------------------------------------------------------------------------
epoch 3 loss 49058.87: 100%|██████████████| 20000/20000 [31:08<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:41<00:00,  5.42it/s]
Rouge-1: 29.93
Rouge-2: 14.97
Rouge-L: 27.9
BLEU-4: 0.1191
----------------------------------------------------------------------------------------------------
epoch 4 loss 45132.87: 100%|██████████████| 20000/20000 [31:08<00:00, 10.71it/s]
100%|█████████████████████████████████████████| 553/553 [01:39<00:00,  5.58it/s]
Rouge-1: 30.22
Rouge-2: 15.0
Rouge-L: 28.22
BLEU-4: 0.1192
----------------------------------------------------------------------------------------------------
epoch 5 loss 41556.33: 100%|██████████████| 20000/20000 [31:09<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:38<00:00,  5.60it/s]
Rouge-1: 30.12
Rouge-2: 15.12
Rouge-L: 28.06
BLEU-4: 0.1192
----------------------------------------------------------------------------------------------------
epoch 6 loss 38239.26: 100%|██████████████| 20000/20000 [31:08<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:37<00:00,  5.68it/s]
Rouge-1: 30.01
Rouge-2: 14.86
Rouge-L: 28.14
BLEU-4: 0.1159
----------------------------------------------------------------------------------------------------
epoch 7 loss 35088.80: 100%|██████████████| 20000/20000 [31:09<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:37<00:00,  5.67it/s]
Rouge-1: 30.57
Rouge-2: 15.54
Rouge-L: 28.66
BLEU-4: 0.1232
----------------------------------------------------------------------------------------------------
epoch 8 loss 32208.71: 100%|██████████████| 20000/20000 [31:08<00:00, 10.70it/s]
100%|█████████████████████████████████████████| 553/553 [01:36<00:00,  5.71it/s]
Rouge-1: 30.08
Rouge-2: 15.37
Rouge-L: 28.23
BLEU-4: 0.1232
----------------------------------------------------------------------------------------------------
epoch 9 loss 29451.46: 100%|██████████████| 20000/20000 [31:02<00:00, 10.74it/s]
100%|█████████████████████████████████████████| 553/553 [01:34<00:00,  5.87it/s]
Rouge-1: 30.22
Rouge-2: 15.18
Rouge-L: 28.34
BLEU-4: 0.1183
----------------------------------------------------------------------------------------------------
epoch 10 loss 26907.70: 100%|█████████████| 20000/20000 [31:02<00:00, 10.74it/s]
100%|█████████████████████████████████████████| 553/553 [01:34<00:00,  5.83it/s]
Rouge-1: 30.58
Rouge-2: 15.64
Rouge-L: 28.6
BLEU-4: 0.1229
----------------------------------------------------------------------------------------------------

我们看损失和分数怎么走。累计损失从第一轮的 59105.36 降到第十轮的 26907.70，摊到每一步大约从 2.96 降到 1.35，下降很稳。生成指标里 Rouge-1 从 29.16 爬到 30.5 上下，第七轮摸到最高的 30.57，之后几轮在原地波动，说明再往下练边际收益已经很小，十轮是个划算的停点。

3. 评估函数

训练时每轮结束都会自动跑一次评估，这段评估代码具体在算什么，得单独拆开看。生成式摘要不能只盯着交叉熵，得真的把句子生成出来，再和人工摘要比。这里用了四个指标，Rouge-1 和 Rouge-2 分别看一元、二元词组的重合，Rouge-L 看最长公共子序列，另外再加一个 BLEU-4。生成时用束搜索，束宽设成 3。完整代码如下。

@torch.no_grad()
def train_evaluate(estimator):
    estimator.eval()

    tokenizer = PegasusTokenizer.from_pretrained('Randeng-Pegasus-238M-Summary-Chinese')

    def collate_function(batch_data):
        encoder_inputs, decoder_inputs = [], []
        for einputs, dinputs in batch_data:
            encoder_inputs.append(einputs)
            decoder_inputs.append(dinputs)

        # 将文本输入转换索引表示
        encoder_inputs = tokenizer(encoder_inputs, max_length=1024, truncation=True, padding='longest', return_tensors="pt")
        # 将张量移动到 device 设备上
        encoder_inputs = {key: value.to(device) for key, value in encoder_inputs.items()}
        return encoder_inputs, decoder_inputs

    # 评估数据集
    valid_inputs, valid_labels = data_read(type='valid')
    # valid_inputs, valid_labels = valid_inputs[:100], valid_labels[:100]
    valid_data = SummaryDataset(valid_inputs, valid_labels)
    dataloader = DataLoader(valid_data, shuffle=False, batch_size=2, collate_fn=collate_function)

    true_labels, pred_labels = [], []
    proress = tqdm(range(len(dataloader)))
    for encoder_inputs, decoder_inputs in dataloader:
        outputs = estimator.generate(encoder_inputs['input_ids'], max_length=512, num_beams=3, do_sample=False)
        pred_label = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        true_label = decoder_inputs
        true_labels.extend(true_label)
        pred_labels.extend(pred_label)
        proress.update()
    proress.close()

    true_labels = [' '.join(tokenizer.tokenize(true_label)) for true_label in true_labels]
    pred_labels = [' '.join(tokenizer.tokenize(pred_label)) for pred_label in pred_labels]

    # 计算评估分数
    rouge = Rouge(metrics=['rouge-1', 'rouge-2', 'rouge-l'], stats=['f'])
    rouge_scores = rouge.get_scores(pred_labels, true_labels)

    rouge_1 = [score['rouge-1']['f'] for score in rouge_scores]
    rouge_2 = [score['rouge-2']['f'] for score in rouge_scores]
    rouge_L = [score['rouge-l']['f'] for score in rouge_scores]

    # 将 rouge score 表示成百分比的形式，所以分数乘以 100
    print('Rouge-1:', np.round(np.mean(rouge_1) * 100, decimals=2))
    print('Rouge-2:', np.round(np.mean(rouge_2) * 100, decimals=2))
    print('Rouge-L:', np.round(np.mean(rouge_L) * 100, decimals=2))

    # 计算 BLEU 值
    # bleu_score 函数要求我们输入的摘要以词列表的形式
    # 候选摘要: [['候', '选', '摘', '要', '1'], ['候', '选', '摘', '要', '2']]
    # 参考摘要: [[['参考摘要1-1', '参考摘要1-2']], [['参考摘要2-1']]]
    # 一个候选摘要可以对应多个参考摘要
    true_labels = [[tokenizer.tokenize(true_label)] for true_label in true_labels]
    pred_labels = [tokenizer.tokenize(pred_label)   for pred_label in pred_labels]

    bleu = bleu_score(candidate_corpus=pred_labels, references_corpus=true_labels)
    print('BLEU-4:', np.round(bleu, decimals=4))


# 从文件加载模型
@torch.no_grad()
def evaluate(checkpoint):

    estimator = PegasusForConditionalGeneration.from_pretrained(checkpoint).to(device)
    estimator.eval()
    tokenizer = PegasusTokenizer.from_pretrained('Randeng-Pegasus-238M-Summary-Chinese')

    def collate_function(batch_data):
        encoder_inputs, decoder_inputs = [], []
        for einputs, dinputs in batch_data:
            encoder_inputs.append(einputs)
            decoder_inputs.append(dinputs)

        # 将文本输入转换索引表示
        encoder_inputs = tokenizer(encoder_inputs, max_length=1024, truncation=True, padding='longest', return_tensors="pt")
        # 将张量移动到 device 设备上
        encoder_inputs = {key: value.to(device) for key, value in encoder_inputs.items()}
        return encoder_inputs, decoder_inputs

    # 评估数据集
    valid_inputs, valid_labels = data_read(type='test')
    # valid_inputs, valid_labels = valid_inputs[:100], valid_labels[:100]
    valid_data = SummaryDataset(valid_inputs, valid_labels)
    dataloader = DataLoader(valid_data, shuffle=False, batch_size=4, collate_fn=collate_function)

    true_labels, pred_labels = [], []
    proress = tqdm(range(len(dataloader)))
    for encoder_inputs, decoder_inputs in dataloader:
        outputs = estimator.generate(encoder_inputs['input_ids'], max_length=512, num_beams=3, do_sample=False)
        pred_label = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        true_label = decoder_inputs
        true_labels.extend(true_label)
        pred_labels.extend(pred_label)
        proress.update()
    proress.close()

    true_labels = [' '.join(tokenizer.tokenize(true_label)) for true_label in true_labels]
    pred_labels = [' '.join(tokenizer.tokenize(pred_label)) for pred_label in pred_labels]

    # 计算评估分数
    rouge = Rouge(metrics=['rouge-1', 'rouge-2', 'rouge-l'], stats=['f'])
    rouge_scores = rouge.get_scores(pred_labels, true_labels)

    rouge_1 = [score['rouge-1']['f'] for score in rouge_scores]
    rouge_2 = [score['rouge-2']['f'] for score in rouge_scores]
    rouge_L = [score['rouge-l']['f'] for score in rouge_scores]

    # 将 rouge score 表示成百分比的形式，所以分数乘以 100
    print('Rouge-1:', np.round(np.mean(rouge_1) * 100, decimals=2))
    print('Rouge-2:', np.round(np.mean(rouge_2) * 100, decimals=2))
    print('Rouge-L:', np.round(np.mean(rouge_L) * 100, decimals=2))

    # 计算 BLEU 值
    # bleu_score 函数要求我们输入的摘要以词列表的形式
    # 候选摘要: [['候', '选', '摘', '要', '1'], ['候', '选', '摘', '要', '2']]
    # 参考摘要: [[['参考摘要1-1', '参考摘要1-2']], [['参考摘要2-1']]]
    # 一个候选摘要可以对应多个参考摘要
    true_labels = [[tokenizer.tokenize(true_label)] for true_label in true_labels]
    pred_labels = [tokenizer.tokenize(pred_label)   for pred_label in pred_labels]

    bleu = bleu_score(candidate_corpus=pred_labels, references_corpus=true_labels)
    print('BLEU-4:', np.round(bleu, decimals=4))

这段代码里有两个函数。train_evaluate 在训练循环里每轮结束自动调用，读 valid 文件逐条生成打分。evaluate 是独立用的，从存好的 checkpoint 重新加载模型，读 test 文件再整体评一遍。有个小细节，rouge 库吐出来的是 0 到 1 的小数，代码里乘了 100，才变成我们习惯的百分数。

独立评估在 test 文件上的结果如下。2667 个批量乘 batch 4，对应约 1.07 万条样本。

100%|███████████████████████████████████████| 2667/2667 [09:04<00:00,  4.90it/s]
Rouge-1: 30.42
Rouge-2: 15.8
Rouge-L: 28.28
BLEU-4: 0.1286

4. 摘要生成

指标终究是数字，最后还得拿一条真实新闻让模型现场生成一句看看。这里还有一个词表的坑要提前说。保存模型时，tokenizer 加载词表会把几个特殊 token 从 unused 槽位里换出来，存盘之后词表已经被换过一次，下次直接加载存下来的 vocab.txt 会报 KeyError。代码里图省事，干脆一直用原始发布目录 Randeng-Pegasus-238M-Summary-Chinese 里的 vocab.txt，不去碰存出来的那份。

@torch.no_grad()
def predict():

    checkpoint = glob.glob('model/10-pegasus-summary-*')[0]
    print(checkpoint)

    estimator = PegasusForConditionalGeneration.from_pretrained(checkpoint)
    estimator.eval()

    # 加载 vocab.txt 词表时，修改后的 pegssus tokenizer 会进行如下的替换
    #         self.vocab[self.eos_token] = self.vocab.pop("[unused1]")
    #         # self.vocab[self.eos_token] = self.vocab.pop("[unused2]")
    #         self.vocab[self.pad_token] = self.vocab.pop("[PAD]")
    #         self.vocab[self.unk_token] = self.vocab.pop("[UNK]")
    # 但是存储模型时，如果直接 tokenizer.save_pretrained 之后，单词已经被替换过了
    # 所以，直接加载词表会报错，提示缺少 unused1 KeyError 错误
    # 解决方法：要不直接用模型的 vocab 进行替换，要不注释掉上面这部分代码
    # 我们这里就将 Randeng-Pegasus-238M-Summary-Chinese 中的 vocab.txt 替换掉 save_pretrained 之后的 vocab.txt
    tokenizer = PegasusTokenizer.from_pretrained('Randeng-Pegasus-238M-Summary-Chinese')

    text = '2007年乔布斯向人们展示iPhone并宣称“它将会改变世界”，还有人认为他在夸大其词，然而在8年后，以iPhone为代表的触屏智能手机已经席卷全球各个角落。未来，智能手机将会成为“真正的个人电脑”，为人类发展做出更大的贡献。'

    print(text)
    inputs = tokenizer([text], max_length=1024, truncation=True, return_attention_mask=False, return_tensors='pt')
    print(inputs)
    # num_beams=1 and do_sample=False 表示使用贪心解码
    # num_beams>1 and do_sample=False 表示使用束搜索解码
    y_pred = estimator.generate(inputs['input_ids'], max_length=512, num_beams=3, do_sample=False)
    print(tokenizer.batch_decode(y_pred, skip_special_tokens=True))

程序输出如下。

model/10-pegasus-summary-26907.70
2007年乔布斯向人们展示iPhone并宣称“它将会改变世界”，还有人认为他在夸大其词，然而在8年后，以iPhone为代表的触屏智能手机已经席卷全球各个角落。未来，智能手机将会成为“真正的个人电脑”，为人类发展做出更大的贡献。
{'input_ids': tensor([[43156,  1625, 36562,   810,  8311, 16475, 48063,  1626, 15849,   175,
          1397,  1463,   346, 20039,  7136,   176,  5661, 32012,   297, 30447,
           314,  1101,  1240,  1230,   527,  4628,  5661, 24631,  1101,   129,
          1625,   807,  5661,   323, 48063,   230,  8486,  3399,  4585,  1497,
          2305,  4054,  1934,  2355, 16764, 16900,  9715, 12721, 30342,   179,
         21415,  5661,  2305,  4054,  1934,  2355,  1463,   346, 18526,   175,
         26223,  3399, 43867,   176,  5661,   230,  8355, 12286,  9354,  2328,
          1230,  3399, 31045,   179,     1]])}
['“触屏”智能手机的未来之路']

输入是一段讲 iPhone 改变世界的长文，模型压出一句“触屏智能手机的未来之路”。句子很短，抓住了触屏智能手机这个主题，但信息量有限，这也符合只拿 8 万条继续微调的预期。想再往上提分数，要么加大训练数据，要么调束搜索的长度惩罚、重复惩罚这些解码参数，或者干脆换更大的预训练权重。到这里，整条链路从读数据、拼输入、梯度累加训练、束搜索评估到现场生成，就全部跑通了。