第7节 基于 Transformer 的对联生成模型#
经过前面6节内容的介绍，对于Transformer相信大家都应该理解得差不多了。不过要想做到灵活运用Transformer结构，那就还得多看看其它场景下的运用。在接下来的这部分内容中，将会以一个含有70余万条的对联数据集为例，来搭建一个基于Transformer结构的对联生成模型。
如图7-1所示便是一个基于Transformer的对联生成模型。可以看出，其实它与前面介绍的基于Transformer结构的翻译模型没有特别的变化，唯一不同的可能就是在对联生成模型中解码器和编码器共用同一个词表（因为两者都是中文）。
7.1 数据预处理#
7.1.1 语料介绍#
按老规矩，在正式介绍模型搭建之前我们还是先来看看数据集都长什么样。本次所使用到的数据集是一个网上公开的对联数据集，在github中搜索"couplet-dataset”就能找到，其一共包含有770491条训练样本，4000条测试样本。和翻译数据集类似，对联数据集也包含上下两句：
1 # 上联 in.txt
2 风 弦 未 拨 心 先 乱
3 花 梦 粘 于 春 袖 口
4 晋 世 文 章 昌 二 陆
1 # 下联 out.txt
2 夜 幕 已 沉 梦 更 闲
3 莺 声 溅 落 柳 枝 头
4 魏 家 词 赋 重 三 曹如上所示便是3条样本，分别存放在in.txt和out.txt这两个文件中。可以看出，原始数据已经做了分字这步操作，所有后续我们只需要进行简单的split操作即可。
7.1.2 数据集构建#
总体上来说对联生成模型的数据集构建过程和翻译模型的数据集构建过程基本上没有太大差别，主要步骤同样也是：①构建字典；②将文本中的每一个词（字）转换为Token序列；③对不同长度的样本序列按照某个标准进行padding处理；④构建DataLoader类。
第1步：定义tokenize
由于原始数据每个字已经被空格隔开了，所以这里tokenizer的定义只需要进行split操作即可，代码如下：
1 def my_tokenizer(s):
2 return s.split()可以看到，其实也非常简单。例如对于如下文本来说
腾 飞 上 铁 ， 锐 意 改 革 谋 发 展 ， 勇 当 千 里 马 其tokenize后的结果为：
['腾', '飞', '上', '铁', '，', '锐', '意', '改', '革', '谋', '发', '展', '，
', '勇', '当', '千', '里', '马']第2步：建立词表
在介绍完tokenize的实现方法后，我们就可以正式通过 torchtext.vocab 中的 Vocab 方法来构建词典了，代码如下：
1 def build_vocab(tokenizer, filepath, min_freq=1, specials=None):
2 if specials is None:
3 specials = ['<unk>', '<pad>', '<bos>', '<eos>']
4 counter = Counter()
5 with open(filepath[0], encoding='utf8') as f:
6 for string_ in f:
7 counter.update(tokenizer(string_))
8 with open(filepath[1], encoding='utf8') as f:
9 for string_ in f:
10 counter.update(tokenizer(string_))
11 return Vocab(counter, specials=specials, min_freq=min_freq)在上述代码中，第3行代码用来指定特殊的字符；第5~10行分别用来遍历in.txt文件和out.txt文件中的每一个样本（每行一个）并进行tokenize和计数，其中对于 counter.update 进行介绍可以参考[10]；第8行则是返回最后得到词典。值得注意的是，由于在对联生成这一场景中编码器和解码器共用的是一个词表，所以这里同时对in.txt和out.txt文件进行了遍历。
在完成上述过程后，我们将得到一个Vocab类的实例化对象，即：
{'<unk>': 0, '<pad>': 1, '<bos>': 2, '<eos>': 3, '，': 4, '风': 5, '春': 6, '
一': 7, '人': 8, '月': 9, '山': 10, '心': 11, '花': 12, '天': 13, ...}此时，我们就需要定义一个类，并在类的初始化过程中根据训练语料完成字典的构建，代码如下：
1 class LoadCoupletDataset():
2 def __init__(self, train_file_paths=None, tokenizer=None,
3 batch_size=2, min_freq=1):
4 # 根据训练预料建立字典，由于都是中文，所以共用一个即可
5 self.tokenizer = tokenizer
6 self.vocab = build_vocab(self.tokenizer, filepath=train_file_paths,min_freq=min_freq)
7 self.specials = ['<unk>', '<pad>', '<bos>', '<eos>']
8 self.PAD_IDX = self.vocab['<pad>']
9 self.BOS_IDX = self.vocab['<bos>']
10 self.EOS_IDX = self.vocab['<eos>']
11 self.batch_size = batch_size第3步：转换为Token序列
在得到构建的字典后，便可以通过如下函数来将训练集和测试集转换成Token序列：
1 def data_process(self, filepaths):
2 """
3 将每一句话中的每一个词根据字典转换成索引的形式
4 :param filepaths:
5 :return:
6 """
7 raw_in_iter = iter(open(filepaths[0], encoding="utf8"))
8 raw_out_iter = iter(open(filepaths[1], encoding="utf8"))
9 data = []
10 for (raw_in, raw_out) in zip(raw_in_iter, raw_out_iter):
11 in_tensor_ = torch.tensor([self.vocab[token] for token in
12 self.tokenizer(raw_in.rstrip("\n"))], dtype=torch.long)
13 out_tensor_ = torch.tensor([self.vocab[token] for token in
14 self.tokenizer(raw_out.rstrip("\n"))], dtype=torch.long)
15 data.append((in_tensor_, out_tensor_))
16 return data在上述代码中，第11~4行分别用来将原始序列上联和目标序列下联转换为对应词表中的Token形式。在处理完成后，就会得到类似如下的结果：
[(tensor([5,549,250,1758,11,228,651]),tensor([154,1420,310,598,29,206,164])),(tensor([12,29,3218,262,6,628,419]),tensor([ 441,62,2049,93,66,304,111])),
(tensor([1137,40,47,286,819,364,1383]),tensor([1803,49,586,556,126,25,1830])),(tensor([7,291,138,115,216,151,9]),tensor([15,311,107,57,80,5,21]))]其中左边的一列就是原始序列上联的Token形式，右边一列就是目标序列下联的Token形式，每一行构成一个样本。
第4步：padding处理
从上面的输出结果可以看到，无论是对于原始序列来说还是目标序列来说，在不同的样本中其对应长度都不尽相同。但是在将数据输入到相应模型时却需要保持同样的长度，因此在这里我们就需要对Token序列化后的样本进行padding处理。同时需要注意的是，一般在这种生成模型中，模型在训练过程中只需要保证同一个batch中所有的原始序列等长，所有的目标序列等长即可，也就是说不需要在整个数据集中所有样本都保证等长。
因此，在实际处理过程中无论是原始序列还是目标序列都会以每个batch中最长的样本为标准对其它样本进行padding，具体代码如下：
1 def generate_batch(self, data_batch):
2 in_batch, out_batch = [], []
3 for (in_item, out_item) in data_batch: # 开始对一个batch中的每一个样本进行处理。
4 in_batch.append(in_item) # 编码器输入序列不需要加起止符
5 # 在每个idx序列的首位加上 起始token 和 结束 token
6 out = torch.cat([torch.tensor([self.BOS_IDX]),
7 out_item, torch.tensor([self.EOS_IDX])], dim=0)
8 out_batch. append(out)
9 # 以最长的序列为标准进行填充
10 in_batch = pad_sequence(in_batch, padding_value=self.PAD_IDX) # [de_len,batch_size]
11 out_batch = pad_sequence(out_batch, padding_value=self.PAD_IDX) # [en_len,batch_size]
12 return in_batch, out_batch在上述代码中，第6~7行用来在目标序列的首尾加上特定的起止符；第10~11行则是分别对一个batch中的原始序列和目标序列以各自当中最长的样本为标准进行padding（这里的pad_sequence导入自torch.nn.utils.rnn）。
第5步：构造mask向量
在处理完成前面几个步骤后，进一步需要根据src_input和tgt_input来构造相关的mask向量，具体代码如下：
1 def generate_square_subsequent_mask(self, sz, device):
2 mask = (torch.triu(torch.ones((sz, sz),device=device)) ==1).transpose(0, 1)
3 mask = mask.float().masked_fill(mask == 0, float('- inf')).
4 masked_fill(mask == 1, float(0.0))
5 return mask
6 def create_mask(self, src, tgt, device='cpu'):
7 src_seq_len = src.shape[0]
8 tgt_seq_len = tgt.shape[0]
9 tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device) # [tgt_len,tgt_len] Decoder的注意力Mask输入，用于掩盖当前position
10 # 之后的position，所以这里是一个对称矩阵
11 src_mask = torch.zeros((src_seq_len, src_seq_len),device=device).type(torch.bool)
12 #Encoder的注意力Mask输入，这部分其实对于Encoder来说是没有用的，所以全是0
13 src_padding_mask = (src == self.PAD_IDX).transpose(0, 1)
14 # 用于mask掉Encoder的Token序列中的padding部分,[batch_size, src_len]
15 tgt_padding_mask = (tgt == self.PAD_IDX).transpose(0, 1)
16 # 用于mask掉Decoder的Token序列中的padding部分,batch_size, tgt_len
17 return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask在上述代码中，第1~5行是用来生成一个形状为[sz,sz]的注意力掩码矩阵，用于在解码过程中掩盖当前position之后的position；第6~17行用来返回Transformer中各种情况下的mask矩阵，其中src_mask在这里并没有作用。
第6步：构造DataLoade与使用示例
经过前面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在上述代码中，第2~3行便是分别用来将训练集和测试集转换为Token序列；第4~7行则是分别构造2个DataLoader，其中 generate_batch 将作为一个参数传入来对每个batch的样本进行处理。在完成类 LoadCoupletDataset 所有的编码过程后，便可以通过如下形式进行使用：
1 if __name__ == '__main__':
2 config = Config()
3 data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
4 batch_size=config.batch_size,
5 tokenizer=my_tokenizer,
6 min_freq=config.min_freq)
7 train_iter, test_iter = data_loader.load_train_val_test_data(config.test_corpus_file_paths,
8 config.test_corpus_file_paths)
9 print(data_loader.PAD_IDX)
10 for src, tgt in train_iter:
11 tgt_input = tgt[:-1, :]
12 tgt_out = tgt[1:, :]
13 src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = data_loader.create_mask(src, tgt_input)
14 print(src)
15 print(tgt)
16 print("src shape：", src.shape) # [in_tensor_len,batch_size]
17 print("tgt shape:", tgt.shape) # [in_tensor_len,batch_size]
18 print("src input shape:", src.shape)
19 print("src_padding_mask shape (batch_size, src_len): ",src_padding_mask.shape)
20 print("tgt input shape:", tgt_input.shape)
21 print("tgt_padding_mask shape: (batch_size, tgt_len) ",tgt_padding_mask.shape)
22 print("tgt output shape:", tgt_out.shape)
23 print("tgt_mask shape (tgt_len,tgt_len): ", tgt_mask.shape)在介绍完数据集构建的整个过程后，下面就开始正式进入到翻译模型的构建中。如果对于这部分不是特别理解的话，建议先看第5节中的数据处理流程图进行理解。
7.2 对联生成模型#
7.2.1 网络结构#
总体来说，基于Transformer的对联生成模型的网络结构其实就是图7-1所展示的所有部分，只是在前面介绍Transformer网络结构时并没有把Embedding部分的实现给加进去。这是因为对于不同的文本生成模型，其Embedding部分会不一样（例如在本场景中编码器和解码器共用一个TokenEmbedding即可，而在翻译模型中就需要两个），所以将两者进行了拆分。同时，待模型训练完成后，在inference过程中Encoder只需要执行一次，所以在此过程中也需要单独使用Transformer中的Encoder和Decoder。
首先，需要定义一个名为CoupletModel的类，其前向传播过程代码如下所示：
1 class CoupletModel(nn.Module):
2 def __init__(self, vocab_size,
3 d_model=512, nhead=8, num_encoder_layers=6,
4 num_decoder_layers=6, dim_feedforward=2048,
5 dropout=0.1):
6 super(CoupletModel, self).__init__()
7 self.my_transformer = MyTransformer(d_model=d_model,
8 nhead=nhead,
9 num_encoder_layers=num_encoder_layers,
10 num_decoder_layers=num_decoder_layers,
11 dim_feedforward=dim_feedforward,
12 dropout=dropout)
13 self.pos_embedding = PositionalEncoding(d_model=d_model,dropout=dropout)
14 self.token_embedding = TokenEmbedding(vocab_size, d_model)
15 self.classification = nn.Linear(d_model, vocab_size)
16
17 def forward(self, src=None, tgt=None, src_mask=None,
18 tgt_mask=None, memory_mask=None, src_key_padding_mask=None,
19 tgt_key_padding_mask=None, memory_key_padding_mask=None):
20 """
21 :param src: Encoder的输入 [src_len,batch_size]
22 :param tgt: Decoder的输入 [tgt_len,batch_size]
23 :param src_key_padding_mask: 用来Mask掉Encoder中不同序列的padding部 分,[batch_size, src_len]
24 :param tgt_key_padding_mask: 用来Mask掉Decoder中不同序列的padding部分 [batch_size, tgt_len]
25 :memory_key_padding_mask: 用来Mask掉Encoder输出的memory中不同序列的 padding部分 [batch_size, src_len]
26 :return:
27 """
28 src_embed=self.token_embedding(src) #[src_len,batch_size,embed_dim]
29 src_embed=self.pos_embedding(src_embed) #[src_len,batch_size,embed_dim]
30 tgt_embed = self.token_embedding(tgt) # [tgt_len, batch_size, embed_dim]
31 tgt_embed = self.pos_embedding(tgt_embed) #[tgt_len,batch_size,embed_dim]
32 outs = self.my_transformer(src=src_embed, tgt=tgt_embed,src_mask=src_mask,
33 tgt_mask=tgt_mask,memory_mask=memory_mask,
34 src_key_padding_mask=src_key_padding_mask,
35 tgt_key_padding_mask=tgt_key_padding_mask,
36 memory_key_padding_mask=memory_key_padding_mask)
37 # [tgt_len,batch_size,embed_dim]
38 logits = self.classification(outs) # [tgt_len,batch_size,tgt_vocab_size]
39 return logits在上述代码中，第7~12行便是用来定义一个Transformer结构；第13~15分别用来定义Positional Embedding、Token Embedding和最后的分类器（需要注意的是这里是共用同一个Token Embedding）；第28~38行便是用来执行整个前向传播过程，其中Transformer的整个前向传播过程在第4节中已经介绍过，在这里就不再赘述。
在定义完logits的前向传播过后，便可以通过如下形式进行使用：
1 if __name__ == '__main__':
2 src_len = 7
3 batch_size = 2
4 dmodel = 32
5 tgt_len = 8
6 num_head = 4
7 src = torch.tensor([[4, 3, 2, 6, 0, 0, 0], # 转换成 [src_len, batch_size]
8 [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1)
9 src_key_padding_mask = torch.tensor([[True, True, True, True, False, False, False],
10 [True, True, True, True, True, False, False]])
11 tgt = torch.tensor([[1, 3, 3, 5, 4, 3, 0, 0],
12 [1, 6, 8, 2, 9, 1, 0, 0]]).transpose(0, 1)
13 tgt_key_padding_mask = torch.tensor([[True, True, True, True, True, True, False, False],
14 [True, True, True, True, True, True, False, False]])
15 trans_model = CoupletModel(vocab_size=10, d_model=dmodel,
16 nhead=num_head, num_encoder_layers=6,
17 num_decoder_layers=6, dim_feedforward=30,
18 dropout=0.1)
19 tgt_mask = trans_model.my_transformer.generate_square_subsequent_mask(tgt_len)
20 logits = trans_model(src, tgt=tgt, tgt_mask=tgt_mask,
21 src_key_padding_mask=src_key_padding_mask,
22 tgt_key_padding_mask=tgt_key_padding_mask,
23 memory_key_padding_mask=src_key_padding_mask)
24 print(logits.shape)接着，我们需要再定义一个Encoder和Decoder在inference中进行使用，代码如下：
1 def encoder(self, src):
2 src_embed = self.token_embedding(src)#[src_len,batch_size,embed_dim]
3 src_embed = self.pos_embedding(src_embed)#[src_len,batch_size,embed_dim]
4 memory = self.my_transformer.encoder(src_embed)
5 return memory # [src_len, batch_size, embed_dim]
6
7 def decoder(self, tgt, memory, tgt_mask):
8 tgt_embed = self.tgt_token_embedding(tgt) #[tgt_len,batch_size,embed_dim]
9 tgt_embed = self.pos_embedding(tgt_embed) #[tgt_len,batch_size,embed_dim]
10 outs = self.my_transformer.decoder(tgt_embed, memory=memory,
11 tgt_mask=tgt_mask) # [tgt_len,batch_size,embed_dim]
12 return outs在上述代码中，第1~5行用于在inference时对输入序列进行编码并得到memory（只需要执行一次）；第7~11行用于根据memory和当前解码时刻的输入对输出进行预测，需要循环执行多次，这部分内容详见7.2.3模型预测部分。
7.2.2 模型训练#
在定义完成整个对联生成模型的网络结构后下面就可以开始训练模型了。由于这部分代码较长，所以下面依旧以分块的形式进行介绍：
第1步：载入数据集
1 def train_model(config):
2 data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
3 batch_size=config.batch_size,
4 tokenizer=my_tokenizer,
5 min_freq=config.min_freq)
6 train_iter, test_iter = data_loader.load_train_val_test_data(
7 config.train_corpus_file_paths,config.test_corpus_file_paths)首先我们可以根据前面的介绍，通过类LoadCoupletDataset来载入数据集，其中config中定义了模型所涉及到的所有配置参数。
第2步：定义模型并初始化权重
1 couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
2 d_model=config.d_model,
3 nhead=config.num_head,
4 num_encoder_layers=config.num_encoder_layers,
5 num_decoder_layers=config.num_decoder_layers,
6 dim_feedforward=config.dim_feedforward,
7 dropout=config.dropout)
8 for p in couplet_model.parameters():
9 if p.dim() > 1:
10 nn.init.xavier_uniform_(p)在载入数据后，便可以定义模型CoupletModel，并根据相关参数对其进行实例化；同时，可以对整个模型中的所有参数进行初始化操作。
第3步：定义损失学习率与优化器
1 loss_fn = torch.nn.CrossEntropyLoss(ignore_index=data_loader.PAD_IDX)
2 learning_rate = CustomSchedule(config.d_model)
3 optimizer = torch.optim.Adam(couplet_model.parameters(),lr=0.,
4 betas=(config.beta1, config.beta2),
5 eps=config.epsilon)在上述代码中，第1行是定义交叉熵损失函数，并同时指定需要忽略的索引ignore_index部分的损失。因为根据tgt_output可知，有些位置上的标签值其实是Padding后的结果，因此在计算损失的时候需要将这些位置给忽略掉。第2行代码则是论文中所提出来的动态学习率计算过程，具体在第5节中已经介绍过来，这里就不再赘述。
第4步：开始训练
1 for epoch in range(config.epochs):
2 losses = 0
3 start_time = time.time()
4 for idx, (src, tgt) in enumerate(train_iter):
5 src = src.to(config.device) # [src_len, batch_size]
6 tgt = tgt.to(config.device)
7 tgt_input = tgt[:-1, :] # 解码部分的输入, [tgt_len,batch_size]
8 src_mask, tgt_mask, src_padding_mask, tgt_padding_mask \
9 = data_loader.create_mask(src, tgt_input, config.device)
10 logits = couplet_model(
11 src=src, # Encoder的token序列输入，[src_len,batch_size]
12 tgt=tgt_input, # Decoder的token序列输入,[tgt_len,batch_size]
13 src_mask=src_mask, # Encoder的注意力Mask输入，这部分其实对于Encoder来说是没有用的
14 tgt_mask=tgt_mask, # Decoder的注意力Mask输入，
15 # 用于掩盖当前position之后的position[tgt_len,tgt_len]
16 src_key_padding_mask=src_padding_mask, # 用于mask掉Encoder的Token序列中的padding部分
17 tgt_key_padding_mask=tgt_padding_mask, # 用于mask掉Decoder的Token序列中的padding部分
18 memory_key_padding_mask=src_padding_mask)
19 # logits 输出shape为[tgt_len,batch_size,tgt_vocab_size]
20 optimizer.zero_grad()
21 tgt_out = tgt[1:, :] #解码部分的真实值 shape: [tgt_len,batch_size]
22 loss = loss_fn(logits.reshape(-1, logits.shape[-1]),tgt_out.reshape(-1))
23 # [tgt_len*batch_size, tgt_vocab_size] with [tgt_len*batch_size,]
24 loss.backward()在上述代码中，第5~9行是用来得到模型各个部分的输入；第10~18行是计算模型整个前向传播的过程；第20~24行则是执行损失计算与反向传播。
以下便是模型训练过程中的输出：
1 -- INFO: Epoch: 0, Batch[29/3010], Train loss : 8.965, Train acc: 0.094
2 -- INFO: Epoch: 0, Batch[59/3010], Train loss : 8.618, Train acc: 0.098
3 -- INFO: Epoch: 0, Batch[89/3010], Train loss : 8.366, Train acc: 0.099
4 -- INFO: Epoch: 0, Batch[119/3010], Train loss : 8.137, Train acc: 0.109
5 ...7.2.3 模型预测#
在介绍完模型的训练过程后接下来就来看模型的预测部分。生成模型的预测部分不像普通的分类任务只需要将网络最后的输出做argmax操作即可，生成模型在预测过程中往往需要按时刻一步步进行来进行。因此，下面我们这里定义一个couplet函数来执行这一过程，具体代码如下：
1 def couplet(model, src, data_loader, config):
2 vocab = data_loader.vocab
3 tokenizer = data_loader.tokenizer
4 model.eval()
5 tokens = [vocab.stoi[tok] for tok in tokenizer(src)] # 构造一个样本
6 num_tokens = len(tokens)
7 src = (torch.LongTensor(tokens).reshape(num_tokens, 1)) # 将src_len 作为第一个维度
8 tgt_tokens = greedy_decode(model, src, max_len=num_tokens + 5,
9 start_symbol=data_loader.BOS_IDX,config=config,
10 data_loader=data_loader).flatten() # 解码的预测结果
11 return "".join([vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")在上述代码中，第5行是将待翻译的源序列进行序列化操作；第7~10行则是通过函数greedy_decode函数来对输入进行解码；第11行则是将最后解码后的结果由Token序列在转换成实际的目标语言。同时，greedy_decode函数实现如下：
1 def greedy_decode(model, src, max_len, start_symbol, config, data_loader):
2 src = src.to(config.device)
3 memory = model.encoder(src) # 对输入的Token序列进行解码翻译
4 ys = torch.ones(1, 1).fill_(start_symbol). \
5 type(torch.long).to(config.device) # 解码的第一个输入，起始符号
6 for i in range(max_len - 1):
7 memory = memory.to(config.device)
8 out = model.decoder(ys, memory) # [tgt_len,tgt_vocab_size]
9 out = out.transpose(0, 1) # [tgt_vocab_size, tgt_len]
10 prob = model.classification(out[:, -1]) # 只对对预测的下一个词进行分类
11 _, next_word = torch.max(prob, dim=1) # 选择概率最大者
12 next_word = next_word.item()
13 ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
14 # 将当前时刻解码的预测输出结果，同之前所有的结果堆叠作为输入再去预测下一个词。
15 if next_word == data_loader.EOS_IDX: # 如果当前时刻的预测输出为结束标志，则跳出循环结束预测。
16 break
17 return ys在上述代码中，第3行是将源序列输入到Transformer的编码器中进行编码并得到Memory；第4~5行是初始化解码阶段输入的第1个时刻的，在这里也就是
最后，我们只需要调用如下函数便可以完成对原始输入上联的下联生成任务：
1 def do_couplet(srcs, config):
2 data_loader = LoadCoupletDataset(config.train_corpus_file_paths,
3 batch_size=config.batch_size,
4 tokenizer=my_tokenizer,
5 min_freq=config.min_freq)
6 couplet_model = CoupletModel(vocab_size=len(data_loader.vocab),
7 d_model=config.d_model,
8 nhead=config.num_head,
9 num_encoder_layers=config.num_encoder_layers,
10 num_decoder_layers=config.num_decoder_layers,
11 dim_feedforward=config.dim_feedforward,
12 dropout=config.dropout)
13 couplet_model = couplet_model.to(config.device)
14 loaded_paras = torch.load(config.model_save_dir + '/model.pkl')
15 couplet_model.load_state_dict(loaded_paras)
16 results = []
17 for src in srcs:
18 r = couplet(couplet_model, src, data_loader, config)
19 results.append(r)
20 return results
21
22
23 if __name__ == '__main__':
24 srcs = ["晚风摇树树还挺",
25 "忽忽几晨昏，离别间之，疾病间之，不及终年同静好",
26 "风声、雨声、读书声，声声入耳",
27 "上海自来水来自海上"]
28 tgts = ["晨露润花花更红",
29 "茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神",
30 "家事、国事、天下事，事事关心",
31 ""]
32 config = Config()
33 srcs = [" ".join(src) for src in srcs]
34 results = do_couplet(srcs, config)
35 for src, tgt, r in zip(srcs, tgts, results):
36 print(f"上联：{''.join(src.split())}")
37 print(f" AI：{r}")
38 print(f"下联：{tgts}")
39 print("=======")在上述代码中，第6~15行是定义网络结构，以及恢复本地保存的网络权重；第16~20行则是开始执行多个上联的下联生成任务；第34-39行为生成示例，其输出结果为：
1 上联：晚风摇树树还挺
2 A I：朝露沾花花更红
3 下联：晨露润花花更红
4
5 上联：忽忽几晨昏，离别间之，疾病间之，不及终年同静好
6 A I：茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神
7 下联：茕茕小儿女，孱羸若此，娇憨若此，更烦二老费精神
8
9 上联：风声、雨声、读书声，声声入耳
10 A I：山色、水色、烟霞色，色色宜人
11 下联：家事、国事、天下事，事事关心
12
13 上联：上海自来水来自海上
14 A I：中山落叶松叶落山中
15
16 上联：月来客栈，迎五湖好友，谈百态人生。
17 A I：花发老舍，请四海嘉宾，喝一杯清茶。以上完整代码可参见 https://github.com/mlwithme/TransformerCouplet。