昇思25天学习打卡营第11天|应用实践之基于MindSpore的GPT2文本摘要

基本介绍

        今天的应用实践是基于MindSpore的GPT2文本摘要，其实就是使用MindSpore下载模型及其权重，准备相关数据集，然后模型构建，接着进行训练，最后进行推理。具体如下：

数据集准备

        本次实验使用的是nlpcc2017摘要数据，内容为新闻正文及其摘要，总计50000个样本。首先，可以通过MindSpore提供的API接口进行数据集下载，并进行加载，这步很简单，不做展示；接着是进行数据集预处理，数据集的格式如下：

我们需要将其处理成GPT2输入的数据集格式，目标数据集格式如下：

关于数据集处理的方法，可使用下面代码：

# preprocess dataset
def process_dataset(dataset, tokenizer, batch_size=6, max_seq_len=1024, shuffle=False):
    def read_map(text):
        data = json.loads(text.tobytes())
        return np.array(data['article']), np.array(data['summarization'])

    def merge_and_pad(article, summary):
        # tokenization
        # pad to max_seq_length, only truncate the article
        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']
    
    dataset = dataset.map(read_map, 'text', ['article', 'summary'])
    # change column names to input_ids and labels for the following training
    dataset = dataset.map(merge_and_pad, ['article', 'summary'], ['input_ids', 'labels'])

    dataset = dataset.batch(batch_size)
    if shuffle:
        dataset = dataset.shuffle(batch_size)

    return dataset

处理完数据集，我们的大部分工作就完成，但是还有最后一步，也是最重要的一步，就是选择一个Tokenizer器，这是因为GPT系列的模型并不是直接输入词向量，而是需要先进行Token化。由于数据集是中文数据集，而GPT2并没有中文Tokenizer器，所以采用中文BERT的Tokenizer，使用如下代码待用Tokenizer，并用于数据集处理：

from mindnlp.transformers import BertTokenizer

# We use BertTokenizer for tokenizing chinese context.
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
len(tokenizer)

from mindnlp.transformers import BertTokenizer

# We use BertTokenizer for tokenizing chinese context.
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
len(tokenizer)

到此，数据集准备工作基本完成

模型构建

        今日所需构建的模型是GPT2，更准确的说是GPT2ForSummarization，我们可以借助MindNLP所提供的GPT2LMHeadModel等API，轻而易举的构建出我们想要的模型，具体代码如下：

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:]
        # Flatten the tokens
        loss = ops.cross_entropy(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1), ignore_index=tokenizer.pad_token_id)
        return loss

模型训练

        准备好模型和数据集后，自然就可以进行模型训练。模型训练需要配置好训练轮次，学习率，模型超参数等等。具体代码如下：

# 我们的学习率是动态学习率，可以借助相关API实现
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

# 训练相关参数配置如下
num_epochs = 1
warmup_steps = 2000
learning_rate = 1.5e-4

num_training_steps = num_epochs * train_dataset.get_dataset_size()

config = GPT2Config(vocab_size=len(tokenizer))
model = GPT2ForSummarization(config)

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)

        一切准备妥当后便可使用 trainer.run(tgt_columns="labels") 进行训练。

模型推理

        模型参数大概有1B，相比于昨天的ChatGLM-6B小了不少，但是仍然很多。训练完大概需要8个小时。

而官方提供的在线环境运行2小时后就会自动释放资源，所以没法训练好，自然也没法进行推理，不过可以大致看看推理的代码，其实就是简单的API调用而已

# 输入文本预处理
def process_test_dataset(dataset, tokenizer, batch_size=1, max_seq_len=1024, max_summary_len=100):
    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)

# 调用API进行推理
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())
    print(output_text)
    i += 1
    if i == 1:
        break

Jupyter在线运行