基于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 npdef 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和summarydataset = dataset.map(read_map, 'text', ['article', 'summary'])# 将列名修改为input_ids和labelsdataset = 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 LearningRateScheduleclass LinearWithWarmUp(LearningRateSchedule):"""Warmup-decay learning rate."""def __init__(self, learning_rate, num_warmup_steps, num_training_steps):super().__init__()self.learning_rate = learning_rateself.num_warmup_steps = num_warmup_stepsself.num_training_steps = num_training_stepsdef construct(self, global_step):if global_step < self.num_warmup_steps:return global_step / float(max(1, self.num_warmup_steps)) * self.learning_ratereturn 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, GPT2LMHeadModelconfig = 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 datasettest_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 = 0for (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进行文本总结任务的基本流程,包括数据导入、数据预处理、模型训练、和模型推理。