基于MindSpore的GPT2文本摘要
数据集加载
使用nlpcc2017摘要数据，共包含5万个样本，内容是新闻正文及其摘要。
from mindnlp.utils import http_get
from mindspore.dataset import TextFileDataset
# 下载数据集
url = 'https://download.mindspore.cn/toolkits/mindnlp/dataset/text_generation/nlpcc2017/train_with_summ.txt'
path = http_get(url, './')
# 加载数据集
dataset = TextFileDataset(str(path), shuffle=False)
数据预处理
train_dataset, test_dataset = dataset.split([0.9, 0.1], randomize=False)
按9：1划分测试集与训练集，randomize表示不对数据进行随机排序，按照原顺序直接拆分。
原始数据格式是：
article: [CLS] article_context [SEP]
summary: [CLS] summary_context [SEP]
期望的预处理后的数据格式是：
[CLS] article_context [SEP] summary_context [SEP]
import json
import numpy as np
def process_dataset(dataset, tokenizer, batch_size=6, max_seq_len=1024, shuffle=False):
# 加载json格式的数据并转成numpy数组
def read_map(text):
data = json.loads(text.tobytes())
return np.array(data['article']), np.array(data['summarization'])
# 使用分词器处理artical和summary。
# text=article表示主文本输入，text_pair=summary表示辅助文本
# padding指将输入序列填充或截断的最大长度
# truncation指定截断策略，only_first表示指仅截断主文本（article）。
# 通常是主文本较长需要截断，而辅助文本较短并且需要完整保留。
def merge_and_pad(article, summary):
tokenized = tokenizer(text=article, text_pair=summary,
padding='max_length', truncation='only_first', max_length=max_seq_len)
return tokenized['input_ids'], tokenized['input_ids']
# 提取article和summary
dataset = dataset.map(read_map, 'text', ['article', 'summary'])
# 将列名修改为input_ids和labels
dataset = dataset.map(merge_and_pad, ['article', 'summary'], ['input_ids', 'labels'])
# 将数据进行分批
dataset = dataset.batch(batch_size)
# 如果shuffle是true，则打乱数据
if shuffle:
dataset = dataset.shuffle(batch_size)
return dataset
这里的tokenizer使用BertTokenizer，因为GPT2没有中文的分词器。
from mindnlp.transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
# 使用刚刚定义的函数分成四批进行处理。
train_dataset = process_dataset(train_dataset, tokenizer, batch_size=4)
模型构建
from mindspore import ops
from mindnlp.transformers import GPT2LMHeadModel
class GPT2ForSummarization(GPT2LMHeadModel):
def construct(
self,
input_ids = None,
attention_mask = None,
labels = None,
):
outputs=super().construct(input_ids=input_ids, attention_mask=attention_mask)
shift_logits=outputs.logits[..., :-1, :]
shift_labels=labels[..., 1:]
loss=ops.cross_entropy(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1), ignore_index=tokenizer.pad_token_id)
return loss
具体解释一下以上代码：
shift 的作用使为了对齐预测和标签，使模型输出和标签对应，从而得到每个位置的预测误差。
举例 I love program
在自回归模型中，模型会根据前文逐步预测（不熟悉的盆友可以看一下上一篇文章：自回归模型与文本生成方法）
- 输入 “I”，输出 “love”
- 输入 “I love”，输出 “programming”
也就是位置1的输出对应的是位置2（的标签），位置2的输出对应位置3（的标签）
对应到实际的数据就是
- outputs.logits[…, :-1, :]
去除 logits 的最后一个时间步，因为没有标签与之对应。继续以上面的为例，就是去掉programming，因为programing没有后面输出了。 - labels[…, 1:]
去除 labels 的第一个时间步，因为没有预测值与之对应。去掉I，因为I前面没有输入，因此也不是输出的一部分。
由此完成了shift错位操作。
搞好数据结构之后再计算损失
- 将shift_logits形状调整成二维张量，使每一行对应一个token的预测分布。
- 将shift_labels形状调整为一维张量，使每个元素对应一个标签。
- 使用cross_entropy()计算交叉熵损失，忽略填充token的损失。
定义学习率warmup
from mindspore import ops
from mindspore.nn.learning_rate_schedule import LearningRateSchedule
class LinearWithWarmUp(LearningRateSchedule):
"""
Warmup-decay learning rate.
"""
def __init__(self, learning_rate, num_warmup_steps, num_training_steps):
super().__init__()
self.learning_rate = learning_rate
self.num_warmup_steps = num_warmup_steps
self.num_training_steps = num_training_steps
def construct(self, global_step):
if global_step < self.num_warmup_steps:
return global_step / float(max(1, self.num_warmup_steps)) * self.learning_rate
return ops.maximum(
0.0, (self.num_training_steps - global_step) / (max(1, self.num_training_steps - self.num_warmup_steps))
) * self.learning_rate
这里定义了LinearWithWarmUp作为自定义学习率调度类。它可以在训练的初始阶段进行学习率的线性预热，再在剩余的训练步骤中线性衰减。
初始化 预热步数、训练步数、学习率。
构建时如果步数小于预热步数则开始进行线性增长，从0增长到learning rate。
如果步数大于等于时，则进行线性衰减，由learning rate变回0。maximum保证不会降低到0以下。
模型训练
内容详见注释内容
# 初始化参数
num_epochs = 1
warmup_steps = 2000
learning_rate = 1.5e-4
#训练的步数=数据集的大小 乘以 准备完整遍历数据集的次数（epoch）
num_training_steps = num_epochs * train_dataset.get_dataset_size()
from mindspore import nn
from mindnlp.transformers import GPT2Config, GPT2LMHeadModel
config = GPT2Config(vocab_size=len(tokenizer))
# 初始化一个用于文本摘要的GPT2模型
model = GPT2ForSummarization(config)
# 初始化学习率调度器和adam优化器
lr_scheduler = LinearWithWarmUp(learning_rate=learning_rate, num_warmup_steps=warmup_steps, num_training_steps=num_training_steps)
optimizer = nn.AdamWeightDecay(model.trainable_params(), learning_rate=lr_scheduler)
from mindnlp._legacy.engine import Trainer
from mindnlp._legacy.engine.callbacks import CheckpointCallback
# 设置检查点回调，用于在训练过程中保存模型检查点
ckpoint_cb = CheckpointCallback(save_path='checkpoint', ckpt_name='gpt2_summarization',
epochs=1, keep_checkpoint_max=2)
# 初始化训练器，设置模型，训练集，训练轮次，优化器和回调函数
trainer = Trainer(network=model, train_dataset=train_dataset,
epochs=1, optimizer=optimizer, callbacks=ckpoint_cb)
# 开启混合精度训练，以提高训练速度和节省显存
trainer.set_amp(level='O1')
# 开始训练，并指定目标列（labels）作为标签。
trainer.run(tgt_columns="labels")
模型推理
def process_test_dataset(dataset, tokenizer, batch_size=1, max_seq_len=1024, max_summary_len=100):
# 处理测试集的过程和训练集差不多
# 依然是提取出article和sumarization
# 再进行分词，制定最大长度和截断
def read_map(text):
data = json.loads(text.tobytes())
return np.array(data['article']), np.array(data['summarization'])
def pad(article):
tokenized = tokenizer(text=article, truncation=True, max_length=max_seq_len-max_summary_len)
return tokenized['input_ids']
dataset = dataset.map(read_map, 'text', ['article', 'summary'])
dataset = dataset.map(pad, 'article', ['input_ids'])
dataset = dataset.batch(batch_size)
return dataset
test_dataset = process_test_dataset(test_dataset, tokenizer, batch_size=1)
# 加载已经训练好的模型
model = GPT2LMHeadModel.from_pretrained('./checkpoint/gpt2_summarization_epoch_0.ckpt', config=config)
# 设为非训练模式，可以禁用掉一些训练相关的操作如dropout
model.set_train(False)
# 模型的结束标记设置成分隔符标记，这样生成的文本遇到分隔符就会终止
model.config.eos_token_id = model.config.sep_token_id
i = 0
for (input_ids, raw_summary) in test_dataset.create_tuple_iterator():
output_ids = model.generate(input_ids, max_new_tokens=50, num_beams=5, no_repeat_ngram_size=2)
output_text = tokenizer.decode(output_ids[0].tolist())
i += 1
总结
本章介绍了使用GPT2进行文本总结任务的基本流程，包括数据导入、数据预处理、模型训练、和模型推理。
打卡凭证
