基于Transformer实现机器翻译（日译中）

基于Transformer & PyTorch 的中日机器翻译模型

机器翻译是一项重要的自然语言处理任务，而Transformer模型是一种广泛应用于机器翻译任务的强大模型。它在2017年被提出，通过引入自注意力机制（self-attention）来解决了传统循环神经网络在长距离依赖建模上的限制。在本文中，我们将学习如何使用Transformer模型进行机器翻译。

以下是一个使用Jupiter 笔记本、 PyTorch、 Torchtext 和 SentencePiece 的教程

Import required packages

首先，让我们确保在我们的系统中安装了以下包，如果您发现有些包丢失，请务必安装它们。

import math
import torchtext
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from collections import Counter
from torchtext.vocab import Vocab
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io
import time
import pandas as pd
import numpy as np
import pickle
import tqdm
import sentencepiece as spm
torch.manual_seed(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(torch.cuda.get_device_name(0)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

device

device(type='cpu')

Get the parallel dataset

在本教程中，我们将使用从 JParaCrawl 下载的日语-英语并行数据集！[ http://www.kecl.NTT.co.jp/icl/lirg/jparacrawl ]被描述为“ NTT 创建的最大的公开可用的英日平行语料库。它主要是通过网络爬行和自动对齐并行句创建的。”你也可以在这里看到那篇论文。

df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None) # 使用pandas库读取中日双语对照数据集文件
trainen = df[2].values.tolist()#[:10000] # 将数据集中第三列（索引为2）的中文句子转换为列表存储
trainja = df[3].values.tolist()#[:10000] # 将数据集中第四列（索引为3）的日文句子转换为列表存储
# trainen.pop(5972)
# trainja.pop(5972)

在导入所有的日语和英语对应项之后，我删除了数据集中的最后一个数据，因为它缺少一个值。两个网站的句子总数为5,973,071，但是，为了学习的目的，经常建议在一次性使用所有数据之前抽样数据并确保一切正常，以节省时间。

下面是数据集中包含的句子示例：

print(trainen[500]) # 打印列表中第500个位置的中文句子
print(trainja[500]) # 打印列表中第500个位置的日文句子

Chinese HS Code Harmonized Code System < HS编码 2905 无环醇及其卤化、磺化、硝化或亚硝化衍生物 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...
Japanese HS Code Harmonized Code System < HSコード 2905 非環式アルコール並びにそのハロゲン化誘導体、スルホン化誘導体、ニトロ化誘導体及びニトロソ化誘導体 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...

我们也可以使用不同的并行数据集来跟随本文，只是要确保我们可以将数据处理成上面所示的两个字符串列表，其中包含日语和英语句子。

这两行代码分别打印出了trainen列表（存储中文句子的列表）和trainja列表（存储日文句子的列表）中索引为500的元素，也就是第501个句子（因为Python列表索引是从0开始的）。这样做通常是为了检查数据加载是否正确，以及直观地了解数据集中句子的内容。

Prepare the tokenizers

与英语或其他按字母顺序排列的语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的标记器，该标记器是使用日语和英语的句子片段创建的，您可以访问JParaCrawl网站下载它们，或单击此处。

en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')

在加载了标记器之后，您可以测试它们，例如，通过执行以下代码。

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')

['▁All',
 '▁residents',
 '▁aged',
 '▁20',
 '▁to',
 '▁59',
 '▁years',
 '▁who',
 '▁live',
 '▁in',
 '▁Japan',
 '▁must',
 '▁enroll',
 '▁in',
 '▁public',
 '▁pension',
 '▁system',
 '.']

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type='str')

['▁',
 '年金',
 '▁日本',
 'に住んでいる',
 '20',
 '歳',
 '~',
 '60',
 '歳の',
 '全ての',
 '人は',
 '、',
 '公的',
 '年金',
 '制度',
 'に',
 '加入',
 'しなければなりません',
 '。']

Build the TorchText Vocab objects and convert the sentences into Torch tensors

然后，使用标记器和原始句子，构建从 TorchText 导入的词汇对象。这个过程可能需要几秒钟或几分钟，这取决于我们的数据集的大小和计算能力。不同的标记器也会影响构建单词所需的时间，我尝试了其他几个日语标记器，但 SentencePiece 似乎对我来说工作得很好，速度也足够快。

def build_vocab(sentences, tokenizer):
  counter = Counter() # 初始化计数器，用于统计单词出现的频次
# 遍历句子列表
  for sentence in sentences:
    # 使用tokenizer对句子进行编码，获得单词列表，并更新计数器
    counter.update(tokenizer.encode(sentence, out_type=str))
  return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>']) # 根据单词频次创建词汇表，包含特殊标记

ja_vocab = build_vocab(trainja, ja_tokenizer) # 使用自定义的ja_tokenizer处理日语文本数据，构建日语词汇表
en_vocab = build_vocab(trainen, en_tokenizer) # 使用自定义的en_tokenizer处理英语文本数据，构建英语词汇表

这段代码定义了一个build_vocab函数，其作用是根据给定的句子列表和一个分词器（tokenizer），统计句子中所有单词的出现频次，并创建一个词汇表（Vocab）对象。
这个词汇表不仅包括了正常的单词，还特别包含了四个特殊标记：未知词标记、填充标记、句子开始标记和句子结束标记。
之后，该函数被调用来分别构建了日语和英语的词汇表。

在获得词汇表对象之后，我们可以使用词汇表和 tokenizer 对象为我们的训练数据构建张量。

def data_process(ja, en):
  data = [] # 初始化一个空列表，用于存储处理后的数据对
  for (raw_ja, raw_en) in zip(ja, en):
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long) # 使用ja_tokenizer对日语句子进行编码，然后将每个token转换为词汇表中的索引，并构建为张量
    en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
                            dtype=torch.long) # 使用en_tokenizer对英语句子进行编码，然后将每个token转换为词汇表中的索引，并构建为张量
    data.append((ja_tensor_, en_tensor_)) # 将处理后的日语和英语句子张量作为元组添加到data列表中
  return data # 返回处理完毕的数据列表
train_data = data_process(trainja, trainen) # 调用data_process函数处理训练数据

这段代码定义了一个data_process函数，其目的是将原始的日语和英语句子转换为模型可以直接处理的张量格式。函数首先遍历给定的日语和英语句子列表，对每一对句子进行以下处理步骤：去除行尾换行符，使用对应的分词器进行编码，然后将每个token（单词或子词）转换为词汇表中的索引，并将这些索引构建成Long类型的PyTorch张量。处理完成后，将每一对日语和英语的张量作为一个元素添加到结果列表中。最后，调用这个函数处理训练数据，生成train_data，为模型训练做好数据准备。

Create the DataLoader object to be iterated during training

在这里，我将BATCH_SIZE设置为16，以防止“cuda内存不足”，但这取决于各种因素，如机器内存容量、数据大小等，所以可以根据您的需要随意更改批量大小（注意：PyTorch的教程使用Multi30k德语-英语数据集将批量大小设置为128。）

BATCH_SIZE = 8 # 定义批次大小
PAD_IDX = ja_vocab['<pad>'] # 填充符号的索引
BOS_IDX = ja_vocab['<bos>'] # 句子开始符号的索引
EOS_IDX = ja_vocab['<eos>'] # 句子结束符号的索引
def generate_batch(data_batch):
  ja_batch, en_batch = [], [] # 初始化空列表用于存放处理后的日语和英语批次数据
# 遍历批次中的每一对数据
  for (ja_item, en_item) in data_batch:
    # 对于每个句子，添加开始标志<BOS>在句首，结束标志<EOS>在句尾，并使用torch.cat拼接
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
    # 使用pad_sequence函数对所有句子进行填充，保证批内序列长度一致，填充值为PAD_IDX
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
    # 返回处理后的批次数据
  return ja_batch, en_batch
# 使用DataLoader创建训练数据迭代器，指定批次大小、数据混洗以及自定义的生成批次函数
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

定义了批次处理的一些关键变量，如批次大小（BATCH_SIZE）及特殊符号在词汇表中的索引（PAD_IDX, BOS_IDX, EOS_IDX）。之后定义了generate_batch函数，该函数接收一个数据批次，为其中的每对日英句子添加起始和结束标记，并通过填充操作使所有序列达到相同长度，以便于模型批量处理。最后，利用DataLoader创建了一个训练数据迭代器，它会在每次迭代时返回由generate_batch处理过的、大小为BATCH_SIZE的训练批次，并且数据会在每个epoch开始时被随机打乱，以增加模型训练时的泛化能力。

Sequence-to-sequence Transformer

接下来的几个代码和文本解释（用斜体书写）取自PyTorch的原始教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]. 我没有做任何更改，除了BATCH_SIZE和单词de_vocab被改为ja_vocab。

Transformer是“注意力就是你所需要的”论文中介绍的一个Seq2Seq模型，用于解决机器翻译任务。转换器模型由编码器和解码器块组成，每个块包含固定数量的层。

编码器通过一系列多头注意和前馈网络层传播输入序列来处理输入序列。编码器的输出（称为存储器）与目标张量一起被馈送到解码器。编码器和解码器使用教师强制技术以端到端的方式进行训练。

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)


class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward:int = 512, dropout:float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        # 初始化编码器层
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 创建编码器，由多个编码器层堆叠而成
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        # 初始化解码器层
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 创建解码器，同样由多个解码器层堆叠
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

        # 用于从模型输出到目标词汇表的线性变换
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        # 初始化源语言和目标语言的词嵌入层
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        # 位置编码层，为输入添加位置信息
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)

    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        # 对源语言和目标语言输入添加位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        # 编码阶段，生成记忆向量
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        # 解码阶段，基于记忆向量生成输出序列
        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
                                        tgt_padding_mask, memory_key_padding_mask)
        # 线性变换得到最终输出
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

文本标记通过使用标记嵌入来表示。位置编码被添加到标记嵌入中，以引入单词顺序的概念。

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size) # 计算位置编码中的sin和cos使用的系数
        pos = torch.arange(0, maxlen).reshape(maxlen, 1) # 生成位置索引张量
        # 初始化位置嵌入矩阵，并分别计算偶数位置（sin）和奇数位置（cos）的值
        pos_embedding = torch.zeros((maxlen, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 为后续操作方便，增加一个维度
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 定义Dropout层
        self.dropout = nn.Dropout(dropout)
        # 将位置嵌入注册为模型的缓冲区，不会被优化器更新
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0),:]) # 将位置编码与词嵌入相加，并应用Dropout

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size) # 初始化词嵌入层
        self.emb_size = emb_size # 词嵌入维度大小，用于后续计算
    def forward(self, tokens: Tensor):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size) # 获取词嵌入并乘以sqrt(emb_size)来缩放词嵌入的初始化范围

该代码定义了两个类，PositionalEncoding和TokenEmbedding，这两个类是自然语言处理任务中使用的Transformer模型的核心组件，旨在增强模型对序列中单词顺序的理解。
这两个类在NLP模型的输入层协同工作，其中TokenEmbedding负责将单词转换为单词嵌入，而PositionalEncoding将位置信息添加到这个基础上。两者的结合使模型能够理解文本中单词的语义及其在句子中的相对位置，这是Transformer模型理解序列数据的基础。

我们创建一个后续单词掩码来阻止目标单词关注其后续单词。我们还创建掩码，用于掩码源和目标填充令牌

def generate_square_subsequent_mask(sz):
    # 生成一个上三角矩阵，对角线及上方为1，下方为0
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将mask转换为float类型，并将原为0的位置替换为负无穷，1的位置替换为0
    # 这样在后续softmax操作中，位置i之后的位置j的得分将被抑制（因为softmax后接近0）
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

def create_mask(src, tgt):
  src_seq_len = src.shape[0] # 获取源序列的长度 
  tgt_seq_len = tgt.shape[0] # 获取目标序列的长度 

  tgt_mask = generate_square_subsequent_mask(tgt_seq_len) # 为目标序列生成后续子序列掩码，用于屏蔽未来信息
  src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool) # 创建源序列的自注意力掩码，这里全为False，意味着源序列中没有需要屏蔽的位置

  src_padding_mask = (src == PAD_IDX).transpose(0, 1) # 创建源序列和目标序列的填充掩码
  tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
  return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask # 返回所有生成的掩码

此代码主要为Transformer模型中的注意力机制实现两种类型的掩码生成功能：
使用Generate.square_subsequence_mask生成一个上三角矩阵，以屏蔽解码器的自关注层中的未来信息，确保在预测第i个单词时只能看到时间i之前的单词。
createMask函数综合生成源序列的自注意掩码、目标序列的后续子序列掩码以及源序列和目标序列的填充掩码。填充掩码用于忽略输入中的PAD符号，避免它们对注意力得分的影响。

定义模型参数并实例化模型。 这里我们服务器实在是计算能力有限，按照以下配置可以训练但是效果应该是不行的。如果想要看到训练的效果请使用你自己的带GPU的电脑运行这一套代码。

当你使用自己的GPU的时候，NUM_ENCODER_LAYERS 和 NUM_DECODER_LAYERS 设置为3或者更高，NHEAD设置8，EMB_SIZE设置为512。

SRC_VOCAB_SIZE = len(ja_vocab) # 日语文本词汇表的大小
TGT_VOCAB_SIZE = len(en_vocab) # 英语文本词汇表的大小
EMB_SIZE = 512 # 嵌入层的维度大小
NHEAD = 8 # 多头注意力中的头数
FFN_HID_DIM = 512 # 前馈网络隐藏层的维度
BATCH_SIZE = 16 # 批次大小
NUM_ENCODER_LAYERS = 3 # 编码器的层数
NUM_DECODER_LAYERS = 3 # 解码器的层数
NUM_EPOCHS = 16 # 训练的轮数
# 初始化Seq2SeqTransformer模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)

# 使用Xavier初始化方法初始化模型参数
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

transformer = transformer.to(device) # 将模型移动到预定义的设备上（如GPU）

loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX) # 定义损失函数，忽略PAD_IDX位置的损失
# 设置Adam优化器，用于更新模型参数
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
# 定义训练一个epoch的函数
def train_epoch(model, train_iter, optimizer):
  model.train() # 将模型设置为训练模式
  losses = 0 # 初始化损失总和
  for idx, (src, tgt) in  enumerate(train_iter): # 遍历训练数据迭代器
      src = src.to(device)
      tgt = tgt.to(device) # 将数据移到设备上

      tgt_input = tgt[:-1, :] # 截取目标序列，去除最后一个作为输入

      src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input) # 为当前批次创建掩码

      logits = model(src, tgt_input, src_mask, tgt_mask,
                                src_padding_mask, tgt_padding_mask, src_padding_mask) # 模型前向传播得到logits

      optimizer.zero_grad() # 清零梯度

      tgt_out = tgt[1:,:] # 计算损失，只针对除了PAD之外的token
      loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
      loss.backward() # 反向传播并优化

      optimizer.step()
      losses += loss.item() # 累加损失
  return losses / len(train_iter) # 返回平均损失

# 定义验证函数
def evaluate(model, val_iter):
  model.eval() # 将模型设置为评估模式
  losses = 0 # 初始化损失总和
  for idx, (src, tgt) in (enumerate(valid_iter)):
    src = src.to(device)
    tgt = tgt.to(device)

    tgt_input = tgt[:-1, :]

    src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

    logits = model(src, tgt_input, src_mask, tgt_mask,
                              src_padding_mask, tgt_padding_mask, src_padding_mask)
    tgt_out = tgt[1:,:]
    loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
    losses += loss.item()
  return losses / len(val_iter)

上面的代码片段首先设置模型的基本参数，然后初始化Seq2SeqTransformer模型，并使用Xavier初始化其参数。接下来，定义了优化器和损失函数，并实现了用于训练历元和评估模型性能的函数。在训练过程中，通过计算损失和反向传播来更新模型参数；在评估阶段，不进行反向传播，仅用于评估模型在验证集上的性能。

Start training

最后，在准备好必要的类和函数之后，我们就可以训练我们的模型了。这是不言而喻的，但完成训练所需的时间可能会因计算能力、参数和数据集大小等因素的不同而有很大差异。

当我使用JParaCrawl的完整句子列表训练模型时，每种语言大约有590万个句子，使用一个NVIDIA GeForce RTX 3070 GPU每个历元大约需要5个小时。

这是代码：

# 使用tqdm库显示进度条，遍历每个训练轮次（epoch）
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time() # 记录当前轮次开始的时间
  train_loss = train_epoch(transformer, train_iter, optimizer) # 训练一个epoch并计算训练损失
  end_time = time.time() # 记录当前轮次结束的时间
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, " # 显示当前epoch和训练损失（保留三位小数）
          f"Epoch time = {(end_time - start_time):.3f}s")) # 显示该轮次训练耗时（秒，保留三位小数）

  0%|          | 0/16 [00:00<?, ?it/s]

此代码执行模型的训练循环，并使用tqdm.tqdm库动态显示训练进度。对于每个epoch（训练回合），它记录开始时间，调用train_epoch函数进行训练，计算训练损失，然后记录结束时间以计算回合的持续时间。最后，打印出当前轮次数、训练损失值和完成轮次所需的总时间，以帮助监控训练过程。

注：由于条件不足，只有CPU无法支持成功的训练。

Try translating a Japanese sentence using the trained model

首先，我们创建翻译新句子的功能，包括获取日语句子、标记化、转换为张量、推理，然后将结果解码回句子等步骤，但这次是用英语。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device) # 源序列数据移到GPU
    src_mask = src_mask.to(device) # 源序列遮罩移到GPU
    memory = model.encode(src, src_mask) # 编码源序列得到记忆向量
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device) # 初始化解码器的输入序列，以开始符号开始
    for i in range(max_len-1): # 迭代直到达到最大长度或遇到结束符号
        memory = memory.to(device)  # 确保记忆向量在GPU上
        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool) # 创建记忆遮罩
        tgt_mask = (generate_square_subsequent_mask(ys.size(0)) 
                                    .type(torch.bool)).to(device) # 为当前的解码序列生成自注意力遮罩
        out = model.decode(ys, memory, tgt_mask) # 解码一步得到输出
        out = out.transpose(0, 1) # 调整输出形状以便访问最后一行
        prob = model.generator(out[:, -1]) # 使用生成器模型得到词的概率分布
        _, next_word = torch.max(prob, dim = 1) # 选择概率最高的词作为下一个词
        next_word = next_word.item()
        ys = torch.cat([ys,
                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0) # 将选择的词添加到序列中
        # 如果选择了结束符号，则停止生成
        if next_word == EOS_IDX:
          break
    return ys # 返回最终生成的词序列
def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    model.eval() # 确保模型处于评估模式
    # 预处理源文本：添加开始和结束符号，然后转为索引表示
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
    num_tokens = len(tokens)
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1) ) # 转为张量并移到GPU
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool) # 创建源序列遮罩
    tgt_tokens = greedy_decode(model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten() # 使用贪心解码策略生成目标序列
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "") # 将目标序列的索引转换回单词并移除开始与结束符号

这两段代码定义了贪婪解码和翻译的过程。greedy_decode函数实现贪婪解码策略，从给定的源序列逐渐生成目标序列。它首先对源序列进行编码，然后根据起始符号逐渐生成下一个最有可能的单词，直到达到最大长度或生成结束符号。translate函数封装了整个翻译过程，包括准备输入数据，将模型设置为评估模式，调用贪婪解码函数生成目标序列，最后将目标序列的索引转换回实际单词序列，去除特殊标记，并返回翻译结果。

然后，我们可以直接调用translate函数并传递所需的参数。

translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)

' Type Type Type 叁氯 叁氯 叁氯 叁氯 ด ▁Different ▁Different ▁Different ▁Different 叁氯 ▁日本 麓 麓 ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different ▁Different'

trainen.pop(5) # 从训练的英文数据集中移除第5个元素

'Chinese HS Code Harmonized Code System < HS编码 8515 : 电气(包括电热气体)、激光、其他光、光子束、超声波、电子束、磁脉冲或等离子弧焊接机器及装置,不论是否 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

trainja.pop(5) # 从训练的日文数据集中移除第5个元素

'Japanese HS Code Harmonized Code System < HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)、レーザーその他の光子ビーム式、超音波式、電子ビーム式、 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

Save the Vocab objects and trained model

最后，在训练完成后，我们将首先使用Pickle保存Vocab对象（en_Vocab和ja_Vocab）。

import pickle
# open a file, where you want to store the data
file = open('en_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(en_vocab, file)
file.close()
file = open('ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()

最后，我们还可以使用PyTorch保存和加载函数保存模型以供以后使用。通常，有两种方法可以保存模型，具体取决于我们以后要使用它们。第一个仅用于推理，我们可以稍后加载模型，并使用它将日语翻译为英语。

# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

第二个也用于推理，但也用于稍后加载模型并恢复训练时。

# save model + checkpoint to resume training later
torch.save({
  'epoch': NUM_EPOCHS,
  'model_state_dict': transformer.state_dict(),
  'optimizer_state_dict': optimizer.state_dict(),
  'loss': train_loss,
  }, 'model_checkpoint.tar')

---------------------------------------------------------------------------

NameError                                 Traceback (most recent call last)

<ipython-input-23-2e0ecdf61e01> in <module>
      4   'model_state_dict': transformer.state_dict(),
      5   'optimizer_state_dict': optimizer.state_dict(),
----> 6   'loss': train_loss,
      7   }, 'model_checkpoint.tar')


NameError: name 'train_loss' is not defined

注：由于模型训练不成功，此处出现运行错误，读者不必担心

Conclusion

That’s it! Thank you for reading.