日-中机器翻译模型：基于Transformer和PyTorch_日中翻译专用大模型-CSDN博客

本文链接：https://blog.csdn.net/qq_62734810/article/details/140054445

日-中机器翻译模型：基于Transformer和PyTorch

使用Jupyter Notebook、PyTorch、Torchtext和SentencePiece的教程

导入所需的包

首先，确保我们的系统已安装以下包，如果发现缺少某些包，请确保安装它们。且请保证版本较高的python。

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')

获取平行数据集

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

df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()#[:10000]
trainja = df[3].values.tolist()#[:10000]
# trainen.pop(5972)
# trainja.pop(5972)

在导入了所有的日语和它们的英语对应文本之后，我删除了数据集中的最后一条数据，因为它存在缺失值。总共，在训练集中，包括英文和日文，共有5,973,071个句子。然而，为了学习目的，通常建议对数据进行抽样，并确保一切按预期工作，然后再使用全部数据，以节省时间。

以下是数据集中包含的一条示例句子。

print(trainen[500])
print(trainja[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 ...

我们也可以使用不同的平行数据集来跟随本文，只需确保我们能够将数据处理成上述展示的两个字符串列表，分别包含日语和英语句子即可。

准备分词器

与英语或其他字母语言不同，日语句子中没有空格来分隔单词。我们可以使用 JParaCrawl 提供的分词器，这些分词器使用了 SentencePiece 分别针对日语和英语进行了创建。您可以访问 JParaCrawl 网站下载它们，或者点击这里。

en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='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',
 '歳の',
 '全ての',
 '人は',
 '、',
 '公的',
 '年',
 '金',
 '制度',
 'に',
 '加入',
 'しなければなりません',
 '。']

构建 TorchText 的词汇对象并将句子转换为 Torch 张量

利用分词器和原始句子，我们接着构建从 TorchText 导入的词汇对象。这个过程根据数据集的大小和计算能力可能需要几秒钟或几分钟。不同的分词器也会影响构建词汇所需的时间。我尝试了几种日语分词器，但是 SentencePiece 对我来说效果良好且速度足够快。

def build_vocab(sentences, tokenizer):
    # 使用 Counter 对词汇进行计数
    counter = Counter()
    
    # 遍历所有句子，更新词汇计数器
    for sentence in sentences:
        # 使用指定的 tokenizer 对句子进行编码，并以字符串形式输出
        counter.update(tokenizer.encode(sentence, out_type=str))
    
    # 创建一个 Vocab 对象，使用计数器中的词汇，同时指定特殊标记（unk、pad、bos、eos）
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用 build_vocab 函数构建日语和英语的词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)
en_vocab = build_vocab(trainen, en_tokenizer)

在我们拥有词汇对象之后，我们可以利用词汇和分词器对象来为我们的训练数据构建张量。

def data_process(ja, en):
    data = []
    
    # 遍历日语和英语句子对
    for (raw_ja, raw_en) in zip(ja, en):
        # 使用日语 tokenizer 对原始日语句子进行编码，并转换为 Tensor
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                                  dtype=torch.long)
        
        # 使用英语 tokenizer 对原始英语句子进行编码，并转换为 Tensor
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
                                  dtype=torch.long)
        
        # 将编码后的日语和英语 Tensor 组成元组，加入到数据列表中
        data.append((ja_tensor_, en_tensor_))
    
    # 返回处理后的数据列表
    return data

# 使用 data_process 函数处理训练集，得到日语和英语句子的 Tensor 数据对
train_data = data_process(trainja, trainen)

创建用于训练迭代的DataLoader对象

在这里，我将BATCH_SIZE设置为16，以防止“cuda out of memory”，但这取决于诸如您的机器内存容量、数据大小等各种因素，因此根据需要随时更改批处理大小（注：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 标记，然后连接成 Tensor
        ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
        
        # 在英语句子的开头和结尾添加 BOS 和 EOS 标记，然后连接成 Tensor
        en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
    
    # 对日语和英语句子进行填充，使它们的长度相同
    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)

Sequence-to-sequence Transformer

接下来的几行代码和文本解释（以斜体显示）摘自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]。我没有进行任何修改，除了BATCH_SIZE和词汇de_vocab更改为ja_vocab之外。

Transformer是一种Seq2Seq模型，最初在“Attention is all you need”论文中提出，用于解决机器翻译任务。Transformer模型包括编码器和解码器模块，每个模块包含固定数量的层。

编码器通过一系列的多头注意力和前馈网络层处理输入序列。编码器的输出称为记忆，与目标张量一起馈送到解码器中。编码器和解码器通过强制教学技术进行端到端训练。

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__()
        
        # Transformer Encoder 定义
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # Transformer Decoder 定义
        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)
        
        # 源语言和目标语言的token嵌入
        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))
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        
        # 对目标语言序列进行解码
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        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)

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

import torch
import torch.nn as nn
import math
from torch import Tensor

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        
        # 计算位置编码的数值
        # den 是一个指数衰减函数，用于计算正弦和余弦函数的周期
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        
        # pos 是位置编码的索引，从0到maxlen
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        
        # 初始化位置编码矩阵
        pos_embedding = torch.zeros((maxlen, emb_size))
        
        # 根据位置编码公式，分别计算偶数和奇数位置的正弦和余弦值
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        
        # 增加一个维度，以便能够与token_embedding相加
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 使用dropout层，防止过拟合
        self.dropout = nn.Dropout(dropout)
        
        # 将位置编码矩阵注册为buffer，这样可以确保在保存和加载模型时不会更新该参数
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 返回带有位置编码的token_embedding
        # token_embedding 的维度是 [seq_len, batch_size, emb_size]
        # pos_embedding 的维度是 [maxlen, 1, emb_size]
        # 只选取与token_embedding相同长度的位置编码，然后进行dropout和加法操作
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0), :, :])

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size: int):
        super(TokenEmbedding, self).__init__()
        
        # 定义词嵌入层，用于将token转换为其对应的嵌入向量
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        # 返回token对应的嵌入向量，并乘以sqrt(emb_size)，以改善模型的训练效果
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续词mask，用于阻止目标词关注其后续词。我们还创建了用于遮蔽源语言和目标语言填充标记的mask。

def generate_square_subsequent_mask(sz):
    # 生成一个上三角矩阵，表示每个位置只能看到其之前的位置信息
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    
    # 将mask转换为浮点型，并用-infinity填充不可见部分
    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]  # 获取目标序列的长度

    # 生成目标序列的mask，用于遮蔽未来信息
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    
    # 生成源序列的mask，这里用全零矩阵表示没有任何遮蔽
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

    # 生成源序列和目标序列的填充mask，用于遮蔽pad部分
    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

Define model parameters and instantiate model. 这里我们服务器实在是计算能力有限，按照以下配置可以训练但是效果应该是不行的。如果想要看到训练的效果请使用你自己的带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  # Feedforward网络隐藏层维度
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)  # 定义损失函数

optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)  # 定义优化器

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, :]  # 获取目标序列的输入部分（不包括结束标记）

        # 创建各种mask
        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)

        optimizer.zero_grad()  # 梯度清零

        tgt_out = tgt[1:, :]  # 获取目标序列的输出部分（不包括起始标记）
        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, :]  # 获取目标序列的输入部分（不包括结束标记）

        # 创建各种mask
        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)  # 返回平均损失

开始训练

最终，在准备好必要的类和函数之后，我们准备好训练我们的模型了。毫无疑问，完成训练所需的时间会因计算能力、参数设置以及数据集大小等诸多因素而有很大差异。

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

以下是代码：

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()
  train_loss = train_epoch(transformer, train_iter, optimizer)
  end_time = time.time()
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))

0%| | 0/16 [00:00<?,
?it/s]/opt/conda/lib/python3.11/site-packages/torch/nn/functional.py:5076:
UserWarning: Support for mismatched key_padding_mask and attn_mask is
deprecated. Use same type for both instead. warnings.warn( 6%|▋
| 1/16 [03:19<49:53, 199.57s/it] Epoch: 1, Train loss: 4.470, Epoch
time = 199.575s 12%|█▎ | 2/16 [06:35<46:03, 197.36s/it] Epoch:
2, Train loss: 3.483, Epoch time = 195.811s 19%|█▉ | 3/16
[09:51<42:39, 196.91s/it] Epoch: 3, Train loss: 3.080, Epoch time =
196.363s 25%|██▌ | 4/16 [13:09<39:28, 197.39s/it] Epoch: 4, Train loss: 2.783, Epoch time = 198.121s 31%|███▏ | 5/16
[16:27<36:12, 197.48s/it] Epoch: 5, Train loss: 2.563, Epoch time =
197.636s 38%|███▊ | 6/16 [19:43<32:50, 197.06s/it] Epoch: 6, Train loss: 2.396, Epoch time = 196.259s 44%|████▍ | 7/16
[23:00<29:31, 196.89s/it] Epoch: 7, Train loss: 2.287, Epoch time =
196.519s 50%|█████ | 8/16 [26:16<26:12, 196.61s/it] Epoch: 8, Train loss: 2.192, Epoch time = 196.030s 56%|█████▋ | 9/16
[29:31<22:53, 196.29s/it] Epoch: 9, Train loss: 2.110, Epoch time =
195.564s 62%|██████▎ | 10/16 [32:53<19:46, 197.83s/it] Epoch: 10, Train loss: 2.040, Epoch time = 201.276s 69%|██████▉ | 11/16
[36:06<16:22, 196.57s/it] Epoch: 11, Train loss: 1.980, Epoch time =
193.724s 75%|███████▌ | 12/16 [39:24<13:07, 196.82s/it] Epoch: 12, Train loss: 1.927, Epoch time = 197.384s 81%|████████▏ | 13/16
[42:40<09:50, 196.77s/it] Epoch: 13, Train loss: 1.881, Epoch time =
196.644s 88%|████████▊ | 14/16 [45:56<06:33, 196.51s/it] Epoch: 14, Train loss: 1.841, Epoch time = 195.926s 94%|█████████▍| 15/16
[49:13<03:16, 196.70s/it] Epoch: 15, Train loss: 1.803, Epoch time =
197.120s 100%|██████████| 16/16 [52:27<00:00, 196.71s/it] Epoch: 16, Train loss: 1.768, Epoch time = 193.372s

尝试使用训练好的模型翻译一句日语句子

首先，我们创建函数来翻译一个新的句子，包括获取日语句子、进行标记化、转换为张量、推理，然后将结果解码回一个英语句子的步骤。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码生成目标序列。

    Args:
    - model: Seq2SeqTransformer模型实例
    - src: 源语言序列张量
    - src_mask: 源语言序列的mask
    - max_len: 生成序列的最大长度
    - start_symbol: 目标语言序列的起始符号索引

    Returns:
    - ys: 生成的目标语言序列张量
    """
    src = src.to(device)
    src_mask = src_mask.to(device)
    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)
        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()  # 转换为Python整数
        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):
    """
    翻译函数，将源语言文本翻译为目标语言文本。

    Args:
    - model: 训练好的Seq2SeqTransformer模型
    - src: 源语言文本字符串
    - src_vocab: 源语言词汇表
    - tgt_vocab: 目标语言词汇表
    - src_tokenizer: 源语言文本分词器

    Returns:
    - translation: 翻译后的目标语言文本字符串
    """
    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))  # 将分词后的索引转换为张量
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)  # 创建源语言序列的mask
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()  # 使用贪婪解码生成目标语言序列
    translation = " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")  # 将生成的目标语言序列转换为文本字符串
    return translation

然后，我们只需调用翻译函数并传入必要的参数即可。

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

' UE 歷 อ ▁Instant ▁CAGR 歷 歷 歷 歷 อ อ อ อ อ อ อ 歷 อ 歷 肾盂 歷 肾盂 歷 歷 อ ▁web hour hour 椅 蒸 芝 椅 ▁silence อ อ ▁wedding'

trainen.pop(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)

'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 ...'

保存词汇对象和训练好的模型

最后，在训练完成后，我们将首先使用Pickle保存词汇对象（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')