10.5 Transformer对联模型#

经过前面几节内容的介绍,相信各位读者对于Transformer的基本原理以及实现过程已经有了一个较为清晰地认识。在接下来的这节内容中我们将通过搭建一个完整的对联生成模型来再次理解Transformer模型的整个工作流程,包括数据的预处理、数据集的构建以及模型的推理部分的实现等。

从整体上来看,基于Transformer的对联生成模型在结构上类似于我们在9.9节内容中介绍到的神经翻译模型,编码器对上联进行编码得到记忆向量;然后解码器在逐时刻对其进行解码得到下联。同时,由于对联的上下联长度都是一致的,所以解码器在解码时的最大长度等于编码器的输入长度。进一步,由于在这一场景下编码器和解码器的输入均为汉字所以还可以共用同一个字符嵌入层来进行处理。

本节内容完整示例代码可参见Code/Chapter10/C03_Transformer文件。下面,我们首先开始介绍整个数据集的构建过程。

10.5.1 数据预处理#

经过前面多节内容的介绍,对于文本数据预处理这部分内容相信各位读者已经比较熟悉了,所以在下面的介绍过程中对于之前已经介绍过的内容就不再赘述,直接参见相关引用即可。

1. 数据集介绍

本次所使用到的数据集是一个网上公开的对联数据集,在GitHub中检索"couplet-dataset”即可找到。整个数据集中一共包含有770491条训练样本和4000条测试样本。与翻译数据集类似,对联数据集也包含上下两句,原文件如下所示:

1 # 上联  in.txt
2 晚 风 摇 树 树 还 挺 
3 愿 景 天 成 无 墨 迹 
4 丹 枫 江 冷 人 初 去 
5 
6 # 下联  out.txt
7 晨 露 润 花 花 更 红 
8 万 方 乐 奏 有 于 阗 
9 绿 柳 堤 新 燕 复 来 

如上所示便是3条样本,分别存放在in.txtout.txt这两个文件中。可以看出,原始数据已经做了分字处理这一步,所以后续我们只需要进行简单的分割操作即可。

2. 建立词表

首先根据原始文本构建词表,我们这里直接分别延续使用7.2.4节中介绍的tokenize()函数来将文本切分成字的形式,Vocab()类来构建整个词表。最终,根据原始语料便可以得到类似如下所示词表:

1 [('[UNK]', 0), ('[PAD]', 1), ('[BOS]', 2), ('[EOS]', 3), (',', 4), ('风', 5), ('春', 6), ('一', 7), ('人', 
2 8), ('月', 9), ('山', 10), ('心', 11), ('花', 12), ('天', 13), ('水', 14), ('千', 15), 

3. 构建类初始化方法

此时,我们需要定义一个类,并在类的初始化过程中根据训练语料完成词表的构建,示例代码如下所示:

1 class LoadCoupletDataset():
2     def __init__(self, train_file_paths=None, batch_size=2, top_k=5000):
3         raw_data = self.load_raw_data(train_file_paths)
4         self.vocab = Vocab(top_k, raw_data[0] + raw_data[1], False)
5         self.PAD_IDX = self.vocab['<PAD>']
6         self.BOS_IDX = self.vocab['<BOS>']
7         self.EOS_IDX = self.vocab['<EOS>']
8         self.batch_size = batch_size

在上述代码中,第2行中train_file_paths是指定上联和下联文本数据所在的路径,top_k表示指定用前top_k个字来构建词表。第3~4行是读取原始的上联和下联语料,并根据语料构建得到字典。

4. 转换为词表索引

在构建完成词表后,进一步可根据词表将原始语料样本中的每个词转换为词表索引,示例代码如下所示:

 1     def data_process(self, filepaths):
 2         results = self.load_raw_data(filepaths)
 3         data = []
 4         for (raw_in, raw_out) in zip(results[0], results[1]):
 5             in_tensor_ = torch.tensor([self.vocab[token] for token in
 6                           tokenize(raw_in.rstrip("\n"))], dtype=torch.long)
 7             out_tensor_ = torch.tensor([self.vocab[token] for token in
 8                           tokenize(raw_out.rstrip("\n"))], dtype=torch.long)
 9             data.append((in_tensor_, out_tensor_))
10         return data

在上述代码中,第4行代码是开始遍历上联和下联对应的每一行文本。第5~8行则是分别将上联和下联转换为词表索引。第9行是保存每条处理完成的样本。

最终,根据原始语料便可以得到类似如下所示的样本输出:

1 原始上联: 八 面 逶 迤  岭 嶂 峰 峦 争 翠 秀
2 token id: tensor([ 200,  267, 3444, 3091, 4,  361, 2133,  303, 1754,  334,  235,  215])
3 原始下联: 九 重 激 荡  霓 霞 雾 霭 竞 琼 瑶
4 token id: tensor([ 108,  126,  945,  447, 4, 1976,  280,  579, 1542,  870, 1025, 1129])
5 原始上联: 静 水 西 流 鱼 读 月
6 token id: tensor([312,  14, 188,  50, 389, 414,   9])
7 原始下联: 闲 云 北 去 鸟 谈 天
8 token id: tensor([164,  18, 217, 162, 301, 687,  13])

5. 序列填充

进一步,根据7.6.2节内容的介绍可知,我们需要对原始样本进行填充以保证每个小批量数据中的样本长度一致;同时,根据生成模型的原理我们还需要在解码器目标输入序列的首尾分别加上开始和结束标志,示例代码如下所示:

 1     def generate_batch(self, data_batch):
 2         in_batch, out_batch = [], []
 3         for (in_item, out_item) in data_batch: 
 4             in_batch.append(in_item) 
 5             out = torch.cat([torch.tensor([self.BOS_IDX]), out_item, 
 6                                 torch.tensor([self.EOS_IDX])], dim=0)
 7             out_batch.append(out)
 8         in_batch = pad_sequence(in_batch, padding_value=self.PAD_IDX)
 9         out_batch = pad_sequence(out_batch, padding_value=self.PAD_IDX)
10         return in_batch, out_batch

在上述代码中,第3行是开始遍历data_process()方法返回的每一个样本。第5~6行是在目标序列的首尾分别加上开始和结束符号。第8~9行是分别对与序列和目标序列按照每个小批量样本中的最长样本为标准进行填充,关于pad_sequence()函数的介绍可参见7.2.4节内容。

6. 构造掩码向量

在处理完成前面几个步骤后,进一步需要根据源输入和目标输入来构造相关的掩码向量,示例代码如下所示:

1     def create_mask(self, src, tgt, device='cpu'):
2         src_seq_len, tgt_seq_len = src.shape[0], tgt.shape[0]
3         tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device)
4         src_mask = torch.zeros((src_seq_len, src_seq_len), device=device)
5         src_padding_mask = (src == self.PAD_IDX).transpose(0, 1)
6         tgt_padding_mask = (tgt == self.PAD_IDX).transpose(0, 1)
7         return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

在上述代码中,第3行同10.3.4节中介绍的一致,用于生成掩码注意力矩阵。第4行生成的是一个全零的掩码注意力矩阵用于编码器中,实际上没有任何作用只是作为一个预留的接口使用。第5~6行用于生成源输入和目标输入对应的填充掩码向量,并转换成[batch_size, seq_len]的形状。

7. 构建迭代器

经过前面5步的操作,整个数据集的构建就算是基本完成了,只需要再构造一个DataLoader迭代器即可,示例代码如下所示:

1     def load_train_val_test_data(self, train_file_paths, test_file_paths):
2         train_data = self.data_process(train_file_paths)
3         test_data = self.data_process(test_file_paths)
4         train_iter = DataLoader(train_data, batch_size=self.batch_size,
5                                 shuffle=True, collate_fn=self.generate_batch)
6         test_iter = DataLoader(test_data, batch_size=self.batch_size,
7                                shuffle=True, collate_fn=self.generate_batch)
8         return train_iter, test_iter

10.5.2 网络结构#

在介绍完整个数据集的构建流程后,下面就开始正式进入到对联模型的构建中。总体来说基于Transformer的对联生成模型,其网络结构如图10-18所示,只是在10.4节内容中我们并没有考虑字符嵌入部分的实现,这是因为对于不同的文本生成模型字符嵌入的实现并不一样。例如在本场景中编码器和解码器可以共用一个字符嵌入层,而在翻译模型中则需要两个。

1. 训练结构定义

首先,我们定义一个名为CoupletModel的类来完成对联模型的构建,同时需要定义单独的编码器和解码器以便后续推理时使用,示例代码如下所示:

 1 class CoupletModel(nn.Module):
 2     def __init__(self, vocab_size, d_model=512, nhead=8, num_encoder_layers=6,
 3                  num_decoder_layers=6, dim_feedforward=2048,dropout=0.1):
 4         super(CoupletModel, self).__init__()
 5         self.my_transformer = MyTransformer(d_model,nhead,num_encoder_layers,
 6                                 num_decoder_layers,dim_feedforward,dropout)
 7         self.pos_embedding = PositionalEncoding(d_model, dropout)
 8         self.token_embedding = TokenEmbedding(vocab_size, d_model)
 9         self.classification = nn.Linear(d_model, vocab_size)
10 
11     def forward(self, src=None, tgt=None, src_mask=None,
12                 tgt_mask=None, mem_mask=None, src_key_padding_mask=None,
13                 tgt_key_padding_mask=None, mem_key_padding_mask=None):
14         src_embed = self.pos_embedding(self.token_embedding(src))
15         tgt_embed = self.pos_embedding(self.token_embedding(tgt))
16         outs = self.my_transformer(src_embed, tgt_embed, src_mask,tgt_mask, mem_mask,
17                     src_key_padding_mask,tgt_key_padding_mask,mem_key_padding_mask)
18         logits = self.classification(outs)
19         return logits

在上述代码中,第5~6行是实例化一个Transformer网络结构用于模型的训练过程。第7~8行是分别实例化一个位置编码层和字符嵌入层。第9行是实例化解码器最后的分类层。第11~13行是指定模型各部分的输入,其中src为编码器输入形状为[src_len,batch_size]tgt为训练时的解码器输入形状为[tgt_len,batch_size]tgt_mask为解码器中的注意力掩码矩阵形状为[tgt_len,tgt_len]src_key_padding_mask为编码器输入的填充掩码形状为[batch_size, src_len]tgt_key_padding_mask为解码器输入对应的填充掩码形状为[batch_size, tgt_len]memory_key_padding_mask用于掩盖掉编码器输出中对应填充部分的信息,形状为[batch_size, src_len]。第14~15行是利用同一个嵌入层和位置编码层对源输入序列和目标输入序列进行字符嵌入和位置编码。第16~17行是计算整个结构的前向传播过程。第18~19行是对解码器的输出进行分类并返回。

2. 推理结构定义

在完成训练时的网络结构定义后,我们需要再定义一个分离的编码器和解码器在模型推理时进行使用,示例代码如下:

 1     def encoder(self, src):
 2         src_embed = self.token_embedding(src)
 3         src_embed = self.pos_embedding(src_embed)
 4         mem = self.my_transformer.encoder(src_embed)
 5         return mem
 6 
 7     def decoder(self, tgt, mem):
 8         tgt_embed = self.token_embedding(tgt)
 9         tgt_embed = self.pos_embedding(tgt_embed)
10         outs = self.my_transformer.decoder(tgt_embed, mem=mem)
11         return outs

在上述代码中,第1~5行用于对源输入序列进行编码得到编码向量。第7~11行则是根据当前时刻的输入以及编码器的输出来逐时刻进行解码。

52