经典神经网络(14)T5模型原理详解及其微调(文本摘要)

• 2018 年，谷歌发布基于双向 Transformer 的大规模预训练语言模型 BERT，而后一系列基于 BERT 的研究工作如春笋般涌现，预训练模型也成为了业内解决 NLP 问题的标配。

• 2019年，谷歌又提出预训练模型 T5(Text-to-Text Transfer Transformer)，T5模型本质上来说是一个基于Transformer架构的encoder-decoder模型。T5模型将各种NLP任务都视为Text-to-Text任务，也就是输入为Text，输出也为Text的任务。

• 我们知道BERT相关的预训练语言模型，在下游任务微调过程中都需要添加非线性层，将模型的输出转化为任务指定的输出格式。但是，T5不需要对模型做任何改动，不需要添加任何非线性层，唯一需要做的就是在输入数据前加上任务声明前缀。

• T5模型刚发布时，刷新了 Glue 榜单和 SuperGLUE 榜单，直至今日还是这两个榜单的前10名。

  • https://gluebenchmark.com/leaderboard

  • https://super.gluebenchmark.com/leaderboard

• 今天，我们来了解下T5这个经典的模型。

  • 论文链接：https://arxiv.org/abs/1910.10683
  • Github 链接：https://github.com/google-research/text-to-text-transfer-transformer

1 T5模型简介

• 如下图所示，T5(Text-to-Text Transfer Transformer)模型将翻译、分类、回归、摘要生成等任务都统一转成Text-to-Text任务，从而使得这些任务在训练(pre-train和fine-tune)时能够使用相同的目标函数，在测试时也能使用相同的解码过程。

• T5模型在NLU和NLG上都具有出色表现，能够完成翻译任务、文本分类、阅读理解、摘要生成任务等多种下游任务。

• 然而，T5刚出来的时候，我们可能没有什么存在感，原因很简单：没有中文版T5可用。

  • 不过Google后面放出了多国语言版的T5（mT5），里边包含了中文语言。

    • 论文链接：mT5: A massively multilingual pre-trained text-to-text transformer
    • Hugging face链接：https://huggingface.co/collections/google/mt5-release-65005f1a520f8d7b4d039509

  • 另外，国内还有一些公司，利用T5模型使用了大量中文数据进行训练。

    • 孟子T5预训练生成模型与T5结构相同，但是不包含下游任务，需要在特定任务上 Finetune 后使用。孟子T5预训练生成模型-中文-base

    • iic在mt5模型基础上使用了大量中文数据进行训练，并引入了零样本分类增强的技术。全任务零样本学习-mT5分类增强版-中文-base

1.1 T5模型网络架构

1.1.1 Encoder-Decoder结构

• 如下图所示，目前基于Transformer的模型架构主要有Encoder-Decoder结构(传统的Transformer结构)、Language model结构 (GPT的结构)和Prefix LM结构(UniLM的结构)。
  • Encoder-Decoder结构：Seq2Seq常用模型，编码器输入中可以看到序列中包括自己的全部字符，解码器的输出只能看到当前字符及之前的字符；
  • LM模型：Encoder-Decoder中的Decoder部分，单向结构，每次只能看到当前及之前的部分；
  • 基于前缀的语言模型Prefix LM：前面一部分文本可以看到前缀部分所有内容，后面剩下的内容只能看到自己及之前的内容。

• 如下图所示，作者通过实验发现Encoder-decoder架构的模型效果最好，所以T5模型本质上来说是一个基于Transformer的Encoder-decoder模型。

1.1.2 SentencePiece

• 把一个句子看作一个整体，再拆成片段，而没有保留天然的词语的概念。

• SentencePiece不将空格视为分隔符，而是将字符串作为其原始格式的输入，使用BPE或ULM作为其分词器来构建词汇表。

  • 下划线被引入，代替了空格和句子开头特殊符号；

  from transformers import T5Tokenizer
  
  model_dir = r'D:\\python\\models\\model-download\\iic\\nlp_mt5_zero-shot-augment_chinese-base'
  tokenizer = T5Tokenizer.from_pretrained(model_dir, legacy=False)
  print(tokenizer.tokenize("Don't make the user feel stupid"))
  
  # ['▁Don', "'", 't', '▁make', '▁the', '▁user', '▁feel', '▁stupid']
  

  • 中文可以看到一些多字词，但有些词其实不符合一般的分词习惯

print(tokenizer.tokenize("笔画最多的汉字是龘(da)字"))

# 可以看到"龘"字经过tokenize变为：'<0xE9>', '<0xBE>', '<0x98>'
# ['▁', '笔', '画', '最多', '的', '汉', '字', '是', '<0xE9>', '<0xBE>', '<0x98>', '(', 'da', ')', '字']

1.2 相对位置编码

不同于RNN、CNN等模型，对于Transformer模型来说，位置编码的加入是必不可少的，因为纯粹的Attention模块是无法捕捉输入顺序的，即无法区分不同位置的Token。为此我们大体有两个选择：

• 1、将位置信息融入到输入中，这构成了绝对位置编码的一般做法；
• 2、微调一下Attention结构，使得它有能力分辨不同位置的Token，这构成了相对位置编码的一般做法。

1.2.1 常规相对位置编码的可视化解释

Transformer中有两种常用的位置编码，分别为绝对位置编码和相对位置编码。

我们先看常规相对位置编码的思路：

论文链接：https://arxiv.org/pdf/1803.02155

视频解释：Self-Attention with Relative Position Representations – Paper explained

• 如下图，假如有5个token，其中一个token与其他所有位置包括自己在内的token之间存在一个权重。

• 如下图，$w_0$表示$x_4$与自己的位置关系，0表示与自己的距离，$w_1$表示向右移动一个位置，$w_{-1}$表示向左移动一个位置。

• $x_3$可以表示为下图所示：

• 那么，第一个到最后一个就可以分别表示为下图所示：

• 如下图所示，一共有9个不同的位置编码，分别为$w_{-4},w_{-2},w_{-3},w_{-1},w_0,w_1,w_2,w_3,w_4$。

• 我们可以用用标识对表示

• 我们可以使用一个阈值k，例如k=2，当超过这个特定的阈值(就是下图中红色背景的部分)

• 即其他的position_embedding距离自身超过2个位置，那么这些位置的position_embedding就和距离最近的position_embedding值一样。例如下图中$x_1$的$w_3$和$w_4$就会变成$w_2$，其他同理。

1.2.2 常规相对位置编码的公式解释

• 下图是论文(https://arxiv.org/pdf/1803.02155)中给出的自注意力机制的公式
• 其中$e_{ij}$的计算方式采用的是Scaled Dot-Product

• 我们知道，相对位置编码的做法就是：微调一下Attention结构，使得它有能力分辨不同位置的Token
• 一般认为，相对位置编码是由绝对位置编码启发而来，考虑一般的带绝对位置编码的Attention(下面推导公式来源于苏神博客)：

• Google论文(https://arxiv.org/pdf/1803.02155)中，对上式进行了修改：

• 通过上面的解释，我们就很容易理解论文中下面的公式了：

• 如下图左边所示，是论文中提出的具体的截断方式。
• 如下图右边所示，通过在每个注意头之间共享相对位置表示来降低存储相对位置表示的空间复杂度。
  • 分子第一项中，我们的输入$x_i$的tensor的Shape为:(B, h, seq_length, d)，它计算的是query和key的关系，所以第一项的输出为(B, h, seq_length, seq_length)，第二项的输出shape必须跟第一项一致。
  • 第二项中，$a_{ij}^K$表示的是$ij$的相对位置编码，从位置编码的Embeding向量table中去lookup得到的，lookup后的shape为(seq_length, seq_length, da)，转换下维度得到(seq_length, da, seq_length)，其中原始位置编码lookup后的向量table我们用A来表示，转换维度后我们用$A^T$表示。
  • $x_i$跟$W^Q$相乘后得到tensor其shape为(B, h, seq_length, dz)，转换下维度得到(seq_length, B, h, dz)，再转换下得到(seq_length, B×h, dz)，再跟$a_{ij}^K$来相乘，实质是跟$A^T$相乘，所以(seq_length, B×h, dz)和矩阵(seq_length, da, seq_length)相乘，因此需要dz=da，得到(seq_length, B×h, seq_length)后reshape下得到(seq_length, B, h, seq_length)，转置后shape为(B, h, seq_length, seq_length)这样就跟第一项对应起来了。

1.2.3 T5模型中的位置编码

我们先看下苏神博客中的内容，分析下T5模型中相对位置编码公式的由来：

• T5采用了一个长距离不敏感的相对位置编码，这一设计是考虑到远距离的单词依赖往往比较稀疏且不精细，因此我们需要对周围单词的位置做精确的区分，而远距离单词的位置变化则相对缓慢。
• 如下图所示，T5模型对相对位置进行了一个“分桶”处理，将原始的relative position当成一个个小方块放置在顺序排列的桶中，最后用方块所属的桶号来代替相对距离：
  • 在T5中num_buckets=32，max_distance=128源码中将num_buckets/2的距离定义为近的分割线(对于双向attention是8，对单向attention是16)
  • 低于这个数值的距离被认为是近的，高于这个数值的距离被认为是远的。
  • 这个设计的思路其实也很直观，就是比较邻近的位置(0-7)，我们需要比较得精细一些，所以给它们都分配一个独立的位置编码，至于稍远的位置(比如8~11)，我们不用区分得太清楚，所以它们可以共用一个位置编码。距离越远，共用的范围就可以越大，直到达到指定范围再clip。

• 我们来看下transformers库中T5模型相对位置编码的实现：

    # transformers/models/t5/modeling_t5.py中的T5Attention类
    
    def compute_bias(self, query_length, key_length, device=None):
        """Compute binned relative position bias"""
        if device is None:
            device = self.relative_attention_bias.weight.device
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]    
        # 计算相对位置
        relative_position = memory_position - context_position  # shape (query_length, key_length)
        # 分桶处理
        relative_position_bucket = self._relative_position_bucket(
            relative_position,  # shape (query_length, key_length)
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        # Embedding矩阵为：self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
        # look-up查找，并进行维度转换
        values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
        values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
        return values

# 这里我们假设query_length=key_length=128
# 那么相对位置矩阵为
>>> relative_position
tensor([[   0,    1,    2,  ...,  125,  126,  127],
        [  -1,    0,    1,  ...,  124,  125,  126],
        [  -2,   -1,    0,  ...,  123,  124,  125],
        ...,
        [-125, -124, -123,  ...,    0,    1,    2],
        [-126, -125, -124,  ...,   -1,    0,    1],
        [-127, -126, -125,  ...,   -2,   -1,    0]])
# 分桶后
>>> relative_position_bucket
tensor([[ 0, 17, 18,  ..., 31, 31, 31],
        [ 1,  0, 17,  ..., 31, 31, 31],
        [ 2,  1,  0,  ..., 31, 31, 31],
        ...,
        [15, 15, 15,  ...,  0, 17, 18],
        [15, 15, 15,  ...,  1,  0, 17],
        [15, 15, 15,  ...,  2,  1,  0]])

>>> relative_position_bucket[0]
# 查看第一个，双向attention近的分割线为8
tensor([ 0, 17, 18, 19, 20, 21, 22, 23, 24, 24, 24, 24, 25, 25, 25, 25, 26, 26,
        26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 27, 27, 27, 27, 28, 28, 28, 28,
        28, 28, 28, 28, 28, 28, 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 29, 29,
        29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30, 30, 30, 30,
        30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30,
        30, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31,
        31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31,
        31, 31])

>>> relative_position_bucket[-1]
# 查看最后一个，双向attention近的分割线为8
# 就是比较邻近的位置(0~7)，我们需要比较得精细一些，所以给它们都分配一个独立的位置编码
# 至于稍远的位置(比如8~11)，我们不用区分得太清楚，所以它们可以共用一个位置编码。距离越远，共用的范围就可以越大
tensor([15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
        15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
        15, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14,
        14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 13, 13, 13, 13, 13, 13, 13, 13,
        13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 12, 12, 12, 12, 12, 12, 12, 12,
        12, 12, 12, 12, 12, 12, 11, 11, 11, 11, 11, 11, 11, 11, 11, 10, 10, 10,
        10, 10, 10, 10,  9,  9,  9,  9,  8,  8,  8,  8,  7,  6,  5,  4,  3,  2,
         1,  0])

# look-up查找后
values shape=(128, 128, 12)
# 维度转换后
values shape=(1, 12, 128, 128)

    # transformers/models/t5/modeling_t5.py中的T5Attention类
    def forward(
        self,
        hidden_states,
        mask=None,
        key_value_states=None,
        position_bias=None,
        past_key_value=None,
        layer_head_mask=None,
        query_length=None,
        use_cache=False,
        output_attentions=False,
    ):
        ......
        # 计算注意力分数
        scores = torch.matmul(
            query_states, key_states.transpose(3, 2)
        )  # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9

        if position_bias is None:
            if not self.has_relative_attention_bias:
               ......
            else:
                # 计算相对位置编码 shape = (1, num_heads, seq_length, key_length)
                position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)

            if mask is not None:
                # seq_length = key_length
                # 加上mask信息 mask shape = (batch_size, 1, 1, seq_length)
                # position_bias shape = (batch_size, n_heads, seq_length, seq_length)
                position_bias = position_bias + mask  

        if self.pruned_heads:
            mask = torch.ones(position_bias.shape[1])
            mask[list(self.pruned_heads)] = 0
            position_bias_masked = position_bias[:, mask.bool()]
        else:
            position_bias_masked = position_bias
        
        # Note: 位置编码加在了注意力分数scores上
        scores += position_bias_masked
        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
            scores
        )  # (batch_size, n_heads, seq_length, key_length)
        ......

总结一下：

1.3 T5模型的训练策略

1.3.1 训练数据集探索

作者对公开爬取的网页数据集Common Crawl进行了过滤：

• 仅保留以终端标点符号（即句号、感叹号、问号或结束引号）结尾的文本行；
• 删除任何包含「污秽、下流或其他脏话字眼」的页面；
• 由于爬取到的很多页面包含「应启用 Javascript」的警告信息，所以删除含有 Javascript 一词的所有文本行；
• 有些页面包含占位符「乱数假文」（lorem ipsum），所以删除出现「乱数假文」短语的所有页面；
• 有些页面会无意中含有代码。由于花括号「{」出现在很多编程语言中（如网上广泛使用的 Javascript），但不会出现在自然文本中，所以删除所有含有花括号的页面；
• 为了删除数据集中的重复数据，删除数据集中多次出现的任何三个句子中的两个。
• 此外，由于大多数下游任务都集中在英文文本上，因此研究者使用 langdetect 来过滤掉所有未归类为英文的页面（概率至少为 0.99）。
• 为了汇编基础数据集，他们下载了自 2019 年 4 月开始网络爬取的文本并进行了过滤。这样产生的文本集合不仅比用于预训练的大多数数据集大几个数量级（大约 750GB），而且还包含非常干净自然的英文文本。
• 研究者将此数据集称为「Colossal Clean Crawled Corpus」（或简称 C4 语料库），并将其作为 TensorFlow 数据集的一部分发布。

从上图实验结果可以看出：

(1) C4比unfiltered C4效果好，说明数据清洗的重要性；

(2) Wikipedia+TBC在SGLUE上的效果比C4好，主要是因为在SGLUE中的MultiRC任务得分很高，MultiRC是一个阅读理解数据集，其中的数据主要是小说书籍，和TBC属于同一领域数据。由此说明预训练的数据集中包含一定的领域数据对下游该领域任务的性能提升有效；

从上图中可以看出，随着数据集不断缩小，模型的性能逐渐下降，说明大的模型很可能在小的数据集上发生了过拟合，因此建议预训练模型还是尽可能使用大数据集。

1.3.2 无监督预训练目标探索

论文对无监督目标过程中所做的选择进行了探索，从下面四个方面进行实验：

• 1、自监督的预训练方法
  • 语言模型式：单向的从左到右依次预测，典型代表为GPT-2模型
  • BERT-style式：像BERT模型一样随机破坏掉一部分内容，然后进行还原
  • 顺序还原式：将文本打乱，然后进行还原

如上图所示，给定句子“Thank you for inviting me to your party last week .”，图中展示了针对不同的预训练目标，模型的Input和Target样本的形式。

• 2、对文本一部分进行破坏时的策略
  • Mask法：将被破坏token换成特殊符如[M];
  • Replace span(小段替换)法：可以把它当作是把上面 Mask 法中相[M]都合成了一个特殊符，每一小段替换一个特殊符，提高计算效率;
  • Drop法：没有替换操作，直接随机丢弃一些字符。

其中Replace corrupted spans就是上上图中的noise replace spans目标

Drop corrupted tokens就是上上图中的noise, drop tokens目标。

结果表明，这几种BERT-style预训练目标的变种效果差不多，但是后两种方法不需要预测整个输入序列，而仅需要预测被污染的部分，因此预测的序列长度更短，训练速度也更快。

• 3、对文本百分之多少进行破坏

• 4、对大概多长的文本段进行破坏

因此，经过上述实验，确定了T5模型的预训练方式如下：

1.3.3 T5模型的微调

微调模型的所有参数可能会导致结果欠佳，尤其是在资源匮乏的情况下。论文专注于两种替代的微调方法，这些方法仅更新编码器-解码器模型的参数的子集。

• Adapter layers。
  • 在微调时保持大多数原始模型固定不变。
  • 在transformer每个块中前馈神经网络后添加dense-ReLU-denseblocks。
  • 新的前馈网络使得输出可以与输入维度匹配。 (这样就可以将它们插入网络，而无需更改结构或参数)
  • 进行微调时，仅更新适配器层和层归一化参数。
  • 这种方法的主要超参数是前馈网络的内部维数 d，它改变了添加到模型中的新参数的数量。
• Gradual unfreezing。
  • 初始微调时，只有最后一层的参数被更新，训练一段时间后，倒数第2层及其之后层的参数被更新，直至整个网络的参数都被更新。

如下图所示，所有参数一起更新效果是最好的，但是缺点就是慢。

adapter layers可能是一种在较少参数上进行微调的有前途的技术，只要将维度适当地缩放到任务大小即可，假如任务数据量小的话，d取小一些，任务数据量大的话，d取大一些。

Gradual unfreezing尽管在微调过程中确实提供了一定的加速，但全局解冻会在所有任务中造成轻微的性能下降。因此通过更仔细地调整解冻时间表，可以获得更好的结果。

1.3.4 T5模型的多任务学习

1) 多任务学习如何取样

• 大多数将多任务学习应用于NLP的应用都会添加特定于任务的分类网络，或者为每个任务使用不同的损失函数。本论文的多任务学习仅将不同任务的数据集混合在一起。

• 若现在有无监督任务、有监督任务1、有监督任务2、有监督任务3，一起训练时如何采样数据?

• 作者设置了三种采样方法：

  • Examples-proportional mixing：设任务的数据集大小为$e_n,n\in 1,...,N$，采样时，采样自第m个任务数据的概率为$r_m=\frac{min(e_m,K)}{\sum min(e_n,K)}$，这里K是提前设定的参数；
  • Temperature-scaled mixing：在第一个实验的基础上，对求得的$r_m$再求$\frac{1}{T}$方根，当T=1时，即Examples-proportional mixing。T越大，各个任务数据集采样越均衡；
  • Equal mixing：各个任务数据采样概率相同。

2) 多任务学习+微调

多任务学习+微调：先用多个任务进行预训练，再对具体任务进行微调

• 实验一：在Examples-proportional mixing 的人工混合数据集上预训练模型，然后在每个单独的下游任务上对其进行微调；
• 实验二：在相同的混合数据集上对模型进行预训练，只是从该预训练混合物中省略了一项下游任务。然后，我们在预训练中遗漏的任务上对模型进行微调。对于考虑的每个下游任务，都会重复此步骤。这种方法为“leave-one-out”多任务训练；
• 实验三：把无监督目标(即baseline的预训练目标)剔除，对所有考虑的监督任务进行预训练。

如上图所示，我们发现Multi-task pretraining+fine-tuning的效果和Unsupervised pre-training + fine-tuning的效果差不多，但是前者在预训练过程还能够监控下游任务的性能，因此作者最后采用Multi-task pre-training。

1.4 最终选择

通过对各种对比实验的结果进行分析，作者最终确定了训练T5模型的较优方案：

• 无监督训练目标：采用span-corruption目标，类似SpanBERT的做法。
• 预训练策略：采用multi-task预训练方式(即无监督任务和有监督任务一起预训练)。

• 把前面的最佳方案组合在一起，作者训练了Small、Base、Large、3B、11B五种T5模型。

模型	Layers	Hidden Size	Attention Head	参数量
Small	6	512	8	60M
Base	12	768	12	220M
Large	24	1024	16	770M
3B	24	1024	32	3B
11B	24	2028	128	11B

1.5 mT5和Flan-T5模型

1.5.1 mT5模型

mT5仍是T5数据构造方式(C4数据集)，但语料不再只限于英语，而是扩大到101种语言，其中就包括了中文、俄语等。模型结构是使用的T5.1.1版本，相比于原始T5版本主要有如下升级：

• 激活函数变更: Gated-GELU activation替代ReLU；
• 无标签数据不做dropout；
• embedding和分类层不做参数共享；
• 更大的d_model，更小的num_heads和d_ff；

1.5.2 Flan-T5模型

• 这里的Flan指的是（Instruction finetuning），即"基于指令的微调"。论文的核心贡献是提出一套多任务的微调方案（Flan），来极大提升语言模型的泛化性。
• 如下图所示，引入Flan微调方案，可以很好提高语言模型在超大规模任务上的整体效果。

• Flan微调的过程：

  • 第一步是收集一系列监督的数据；

  • 第二步将任务都转换成相同的“输入格式”喂给模型训练，同时这些任务的输出也需要是统一的“输出格式”，这样是为了使用单个语言模型来完成超过1800+种不同的任务；

    • 如下图，根据 “是否需要进行思维链推理 （CoT）” 以及 “是否需要提供示例（Few-shot）” 可将输入、输出划分成四种类型。
    • 

  • 第三步就是微调过程。

    • 将多个训练样本“打包”成一个训练样本，这些训练样本直接会通过一个特殊的“结束token”进行分割。
    • 训练时候在每个指定的步数会在“保留任务”上进行模型评估，保存最佳的checkpoint。
    • 尽管微调的任务数量很多，但是相比于语言模型本身的预训练过程，计算量小了非常多，只有0.2%。

• 例如下面文章中的例子，模型训练好之后，可直接让模型做问答:

  「模型输入」是："Geoffrey Hinton和George Washington这两个人有没有交谈过？在回答之前想一想原因。“

  「模型返回」是：Geoffrey Hinton是一个计算机科学家，出生在1947年；而George Washington在1799年去世。所以这两个不可能有过交谈。所以答案时“没有”。

• 论文: Scaling Instruction-Finetuned Language Models
  • 地址：https://arxiv.org/abs/2210.11416
  • 模型：Flan-T5 release - a google Collection

2 基于T5模型的微调(文本摘要)

• 我们这里基于澜舟科技开源的mengzi-t5-base做摘要生成。
• 孟子T5与T5结构相同，但是不包含下游任务，需要在特定任务上 finetune 后使用。
• 这里是自己实现文本摘要，如果借助transformers库的组件会更加简单。

2.1 加载数据集

• 数据集地址：https://huggingface.co/datasets/supremezxc/nlpcc_2017

• 我们看一条数据集：

  {
  	'title': '组图:黑河边防军人零下30℃户外训练,冰霜沾满眉毛和睫毛,防寒服上满是冰霜。', 
  	'content': '中国军网2014-12-1709:08:0412月16日,黑龙江省军区驻黑河某边防团机动步兵连官兵,冒着-30℃严寒气温进行体能训练,挑战极寒,锻造钢筋铁骨。该连素有“世界冠军的摇篮”之称,曾有5人24人次登上世界军事五项冠军的领奖台。(魏建顺摄)黑龙江省军区驻黑河某边防团机动步兵连官兵冒着-30℃严寒气温进行体能训练驻黑河某边防团机动步兵连官兵严寒中户外训练,防寒服上满是冰霜驻黑河某边防团机动步兵连官兵严寒中户外训练,防寒服上满是冰霜官兵睫毛上都被冻上了冰霜官兵们睫毛上都被冻上了冰霜驻黑河某边防团机动步兵连官兵严寒中进行户外体能训练驻黑河某边防团机动步兵连官兵严寒中进行户外体能训练驻黑河某边防团机动步兵连官兵严寒中进行户外体能训练'
  }
  

• 我们可以通过num参数控制数据集的数量进行训练。

"""
 微调澜舟科技开源的mengzi-t5-base做摘要生成
    数据集(只取5000条)：nlpcc_2017: https://huggingface.co/datasets/supremezxc/nlpcc_2017
    model link：https://modelscope.cn/models/langboat/mengzi-t5-base
          link: https://huggingface.co/Langboat/mengzi-t5-base

    孟子中文T5预训练生成模型与T5结构相同，只有无监督数据训练，不包含下游任务，需要在特定任务上finetune后使用
"""
import os
import platform
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from datasets import load_from_disk
from transformers import AdamW
from transformers import T5Tokenizer, T5ForConditionalGeneration
from logging_util import get_logger
from rouge import Rouge


device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 获取当前操作系统的名称
os_name = platform.system()
logger = get_logger(model_name='mengzi-t5-base')

# 设置模型路径及数据集路径
if os_name == "Windows":
    model_dir = r'D:\python\models\langboat\meng_zi_t5'
    data_dir = r'D:\python\datas\nlp_seq2seq\nlpcc_2017'
    logger.info("当前执行环境是 Windows...")
elif os_name == "Linux":
    model_dir = r'/root/autodl-fs/models/meng_zi_t5'
    data_dir = r'/root/autodl-fs/data/nlp_ai/nlp_seq2seq/nlpcc_2017'
    logger.info("当前执行环境是 Linux...")
else:
    raise ValueError("当前执行环境不是 Windows 也不是 Linux")


class Dataset(Dataset):
    def __init__(self, split, num=None):

        # 在线加载数据集(需外网环境)
        # dataset = load_dataset(path='nlpcc_2017', split=split)
        # 我们可以将数据保存到本地磁盘，下次利用load_from_disk直接从本地加载即可
        # dataset.save_to_disk("./nlpcc_2017")

        # 划分数据集为：训练集和测试集
        dataset = load_from_disk(dataset_path=data_dir)
        # 选取4900条训练集、100条测试集
        split_dataset = dataset.train_test_split(100, seed=42)
        if num:
            # 离线加载数据集，cpu环境取少量数据用来训练、测试模型
            dataset = split_dataset[split].select(range(num))
        else:
            # 离线加载数据集
            dataset = split_dataset[split]

        # 过滤掉太长的句子，需要去掉CLS、SEP
        def f(data):
            return len(data['content']) <= 512 - 2

        dataset = dataset.filter(f)

        self.dataset = dataset

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, i):
        content = self.dataset[i]['content']
        title = self.dataset[i]['title']
        return content, title

2.2 对数据集进行组装

• 主要就是在content前面加上特定的提示词
• 然后调用tokenizer进行批处理

def get_collate_fn(tokenizer):
    def collate_fn(batch):
        contents = ["摘要生成: \n" + tup2[0] for tup2 in batch]
        original_labels = [tup2[1] for tup2 in batch]
        # 特殊字符
        # 0 -> <pad>
        # 1 -> </s>
        # 2 -> <unk>
        inputs = tokenizer(contents, max_length=384, truncation=True, return_tensors='pt', padding=True)
        labels = tokenizer(text_target=original_labels, max_length=64, truncation=True, return_tensors='pt', padding=True)
        return inputs, labels
    return collate_fn

2.3 构建模型

• 利用transformers库中的T5ForConditionalGeneration加载摘要生成的预训练模型

• T5ForConditionalGeneration源码中实现了KV-cache，有兴趣的可以阅读源码(主要是T5Attention中实现相对位置编码、KV-cache、以及门结构的前馈神经网络T5DenseGatedActDense的实现)。

• T5模型中可以调用generate函数(transformers\generation\utils.py)，这个函数中实现了8种采样的方法(如下代码所示)，每种采样都有不同的配置参数，具体可以看transformers\generation\configuration_utils.py中GenerationConfig类的参数解释。

  # transformers\generation\configuration_utils.py
  class GenerationMode(ExplicitEnum):
      """
      Possible generation modes, downstream of the [`~generation.GenerationMixin.generate`] method.
      """
      # Non-beam methods
      CONTRASTIVE_SEARCH = "contrastive_search"
      GREEDY_SEARCH = "greedy_search"
      SAMPLE = "sample"
      ASSISTED_GENERATION = "assisted_generation"
      # Beam methods
      BEAM_SEARCH = "beam_search"
      BEAM_SAMPLE = "beam_sample"
      CONSTRAINED_BEAM_SEARCH = "constrained_beam_search"
      GROUP_BEAM_SEARCH = "group_beam_search"
  

class MengZiT5Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 加载预训练模型
        self.model = T5ForConditionalGeneration.from_pretrained(model_dir)


    def forward(self, inputs, labels=None):
        # 1、encoder的input_ids和attention_mask
        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']

        if labels is not None:
            # 2、decoder 的labels
            train_labels = labels['input_ids'].contiguous()
            train_labels_mask = labels['attention_mask']

            # 3、decoder 的input_ids和attention_mask
            decoder_input_ids = train_labels.new_zeros(train_labels.shape)
            decoder_input_ids[..., 1:] = train_labels[..., :-1].clone()

            decoder_attention_mask = train_labels_mask.new_zeros(train_labels_mask.shape)
            decoder_attention_mask[..., 1:] = train_labels_mask[..., :-1].clone()
            decoder_attention_mask[..., 0] = 1
            # 4、送入模型进行预测
            outputs = self.model(input_ids=input_ids
                                 , attention_mask=attention_mask
                                 , decoder_input_ids=decoder_input_ids
                                 , decoder_attention_mask=decoder_attention_mask
                                 , labels=train_labels)
            # 5、返回训练时候的Loss值
            return outputs.loss
        else:
            # 模型生成
            summary_ids = self.model.generate(input_ids
                                              , num_beams=4 # 束搜索法
                                              , no_repeat_ngram_size=2 # 确保不重复
                                              , min_length=10 # 长度限制
                                              , max_length=64
                                              , early_stopping=True
            )
            # 将id转换为输出 summary_ids.shape = [bs, length]
            outputs = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
            return outputs

2.4 构造训练及测试函数

• 这里简单实现了模型的训练及评估的函数
• 评估主要使用了指标rouge-1、rouge-2、rouge-l

# 模型训练
def train(epochs, model, loader):
    model.to(device)
    lr = 2e-5
    # 训练
    optimizer = AdamW(model.parameters(), lr=lr)

    model.train()
    for epoch in range(epochs):
        for step, (inputs, labels) in enumerate(loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            # 模型计算
            # [b, lens] -> [b, lens, 8]
            loss = model(inputs, labels)

            # 梯度下降
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            if step % 10 == 0:
                print(f'epoch = {epoch}, step = {step}, loss = {loss:.4f}')

    os.makedirs('./output', exist_ok=True)
    torch.save(model, './output/meng_zi_t5_sft.pt')


# 模型评估
def test():
    # 1、加载模型
    model_load = torch.load('output/meng_zi_t5_sft.pt').to(device)
    model_load.eval()

    rouge = Rouge()

    # 2、加载测试集
    test_loader = DataLoader(dataset=Dataset('test'),
                              batch_size=32,
                              collate_fn=get_collate_fn(tokenizer=tokenizer),
                              shuffle=False,
                              drop_last=True)

    for step, (inputs, labels) in enumerate(test_loader):
        if step == 2:
            break
        with torch.no_grad():
            # [b, lens] -> [b, lens, 8]
            decode_preds = model_load(inputs.to(device))
            decode_labels = tokenizer.batch_decode(labels['input_ids'].to(device), skip_special_tokens=True)
            decode_preds = [" ".join(p) for p in decode_preds]
            decode_labels = [" ".join(l) for l in decode_labels]
            scores = rouge.get_scores(decode_preds, decode_labels, avg=True)
            r = {
                "rouge-1": scores["rouge-1"]["f"],
                "rouge-2": scores["rouge-2"]["f"],
                "rouge-l": scores["rouge-l"]["f"],
            }
            logger.info(f'setp = {step}, 评估结果：\n{r}')
            return r



if __name__ == '__main__':
    # 1、加载分词器
    tokenizer = T5Tokenizer.from_pretrained(model_dir, legacy=False)

    # 2、加载训练数据
    train_loader = DataLoader(dataset=Dataset('train'),
                              batch_size=16,
                              collate_fn=get_collate_fn(tokenizer=tokenizer),
                              shuffle=True,
                              drop_last=True)

    # 3、创建模型
    model = MengZiT5Model()

    # 4、模型训练及评估
    train(epochs=1, model=model, loader=train_loader)

    test() # 模型评估

    # 5、对模型进行预测
    text = """摘要生成: \n在经历了那段惊心动魄但又充满人情味的艰难时刻后，32岁的埃里克森时隔1100天再次为国征战欧洲杯，而且奉献了进球。
    丹麦队对垒斯洛文尼亚，这场热度并不算高的小组赛首轮争夺因为一个人的出现得到了外界的关注，他就是埃里克森。
    曼联中场在在第17分钟的进球帮助祖国球队取得了领先，他也在经历上届欧洲杯的心脏骤停的遭遇之后，实现了“王者归来”。
    尽管这次破门遗憾没能帮助丹麦队最终获得胜利，但绰号“爱神”的埃里克森依然得到了全场乃至全世界球迷的祝福。
    """
    inputs = tokenizer(text, return_tensors='pt')
    model_load = torch.load('output/meng_zi_t5_sft.pt')
    model_load.eval()
    print('摘要内容：\n', model_load(inputs))

# 训练好的模型可以进行摘要生成
摘要内容：
 ['曼联32岁埃里克森时隔1100天再次为国征战欧洲杯,并奉献了进球。']

3 mT5模型的推理

• 这里的mT5模型在下游任务上做了微调，可以直接使用

"""
    model: iic/nlp_mt5_zero-shot-augment_chinese-base
    link: https://modelscope.cn/models/iic/nlp_mt5_zero-shot-augment_chinese-base

该模型在mt5模型基础上使用了大量中文数据进行训练，并引入了零样本分类增强的技术，使模型输出稳定性大幅提升。
    支持任务包含：
        文本分类：给定一段文本和候选标签，模型可输出文本所属的标签。
        自然语言推理：给定两段文本，判断两者关系。
        阅读理解：给定问题和参考文本，输出问题的答案。
        问题生成：给定答案和参考文本，生成该答案对应的问题。
        摘要生成：给定一段文本，生成该文本的摘要。
        标题生成：给定一段文本，为其生成标题。
        评价对象抽取：给定一段文本，抽取该段文本的评价对象。
        翻译：给定一段文本，将其翻译成另一种语言。
"""
import torch
import platform
from transformers import T5Tokenizer, T5ForConditionalGeneration, T5Config


# 获取当前操作系统的名称
os_name = platform.system()
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

# 设置模型路径及数据集路径
if os_name == "Windows":
    model_dir = r'D:\\python\\models\\model-download\\iic\\nlp_mt5_zero-shot-augment_chinese-base'
    print("当前执行环境是 Windows...")
elif os_name == "Linux":
    model_dir = r'/root/autodl-fs/models/nlp_mt5_zero-shot-augment_chinese-base'
    print("当前执行环境是 Linux...")
else:
    raise ValueError("当前执行环境不是 Windows 也不是 Linux")


# 1、加载tokenizer及预训练模型
tokenizer = T5Tokenizer.from_pretrained(model_dir, legacy=False)
model = T5ForConditionalGeneration.from_pretrained(model_dir)
model.to(device)


def t5inference(text, task_prefix, max_length=None, min_length=1):
    """
      :param text: 要生成摘要的文本
      :param task_prefix: 执行的任务
      :param max_length: 摘要的最大长度
      :return:
      """
    # 准备前缀+文本
    t5_prepared_text = f'{task_prefix}: ' + text

    print("input text: \n", t5_prepared_text)

    # 分词
    tokenized_text = tokenizer.encode(t5_prepared_text, return_tensors="pt").to(device)

    # 进行文本摘要
    # prepare_inputs_for_generation
    summary_ids = model.generate(tokenized_text,
                                 num_beams=4, # 束搜索法
                                 no_repeat_ngram_size=2, # 确保不重复
                                 min_length=min_length, # 长度限制
                                 max_length=max_length,
                                 early_stopping=True # 提前停止
    )
    # 将id转换为输出 summary_ids.shape = [1, length]
    output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return output

def t0():
    text = """
          “足球从未高于生死”，这是3年前欧洲杯赛场上丹麦球员埃里克森心脏骤停时，各路媒体报道该事件用的最多的表达。
          而在经历了那段惊心动魄但又充满人情味的艰难时刻后，32岁的埃里克森时隔1100天再次为国征战欧洲杯，而且奉献了进球。
          17日凌晨的欧洲杯小组赛，埃里克森进球的那一刻，感动和欣慰扑面而来。最终丹麦和斯洛文尼亚队1比1战平，各取1分。

          丹麦队对垒斯洛文尼亚，这场热度并不算高的小组赛首轮争夺因为一个人的出现得到了外界的关注，他就是埃里克森。
          曼联中场在在第17分钟的进球帮助祖国球队取得了领先，他也在经历上届欧洲杯的心脏骤停的遭遇之后，实现了“王者归来”。
          尽管这次破门遗憾没能帮助丹麦队最终获得胜利，但绰号“爱神”的埃里克森依然得到了全场乃至全世界球迷的祝福。
        """
    output = t5inference(text, task_prefix='文本摘要', min_length=20, max_length=64)
    print('模型结果:\n', output)

    text = """
        候选标签：故事,房产,娱乐,文化,游戏,国际,股票,科技,军事,教育。
        文本内容：他们的故事平静而闪光，一代人奠定沉默的基石，让中国走向繁荣。
        """
    output = t5inference(text, task_prefix='文本分类', max_length=64)
    print('模型结果:\n', output)

    text = """
            如果日本沉没，中国会接收日本难民吗？
        """
    output = t5inference(text, task_prefix='翻译成英文', max_length=512)
    print('模型结果:\n', output)


if __name__ == '__main__':
    # print(model.config)
    # print(model)
    t0()

input text: 
文本摘要: 
          “足球从未高于生死”，这是3年前欧洲杯赛场上丹麦球员埃里克森心脏骤停时，各路媒体报道该事件用的最多的表达。
          而在经历了那段惊心动魄但又充满人情味的艰难时刻后，32岁的埃里克森时隔1100天再次为国征战欧洲杯，而且奉献了进球。
          17日凌晨的欧洲杯小组赛，埃里克森进球的那一刻，感动和欣慰扑面而来。最终丹麦和斯洛文尼亚队1比1战平，各取1分。

          丹麦队对垒斯洛文尼亚，这场热度并不算高的小组赛首轮争夺因为一个人的出现得到了外界的关注，他就是埃里克森。
          曼联中场在在第17分钟的进球帮助祖国球队取得了领先，他也在经历上届欧洲杯的心脏骤停的遭遇之后，实现了“王者归来”。
          尽管这次破门遗憾没能帮助丹麦队最终获得胜利，但绰号“爱神”的埃里克森依然得到了全场乃至全世界球迷的祝福。 
模型结果:
 埃里克森:足球从未高于生死,爱神归来,为国征战
 
 
input text: 
文本分类: 
        候选标签：故事,房产,娱乐,文化,游戏,国际,股票,科技,军事,教育。
        文本内容：他们的故事平静而闪光，一代人奠定沉默的基石，让中国走向繁荣。     
模型结果:
 文化
 
 
input text: 
翻译成英文: 
            如果日本沉没，中国会接收日本难民吗？     
模型结果:
 will China accept Japanese refugees if Japan sinks?