基于transformer的日译中机器翻译

斯巴拉西！这里是教你怎么实现日译中（机翻）博客！

其实，“斯巴拉西”其实是日语“素晴らしい”的音译，意为“太棒了“。那么，我们怎么实现更精准的日译中工作呢？我们可以用transformer来训练翻译模型。

1 获取数据集

首先，我们可以去http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl来下载常用的日语-英文平行的数据集。

下载完成后，导入数据集：

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，请在你自己的电脑上尝试运行这一套代码
print(f"Torch version: {torch.__version__}, Device: {device}")


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

 代码从一个名为’zh-ja.bicleaner05.txt’的文件中读取数据，并将第2列（英文句子）和第3列（日文句子）分别提取为Python列表 trainen 和 trainja。导入所有的日语和其英语对应文本后，需要删除数据集中的最后一条数据，因为它存在缺失值。

# 打印第 501 行（索引从 0 开始）的英文句子
print(trainen[500])
 
# 打印第 501 行（索引从 0 开始）的日文句子
print(trainja[500])

打印出的结果如下：

 

2 准备分词器

不像英文或其他语言的字母顺序、日本句子并不含有空格分开的语法。 我们可以使用该系统提供的JParaCrawl，它使用了SentencePiece来处理日语和英语，可以访问JParaCrawl网站下载。

# 使用 SentencePieceProcessor 加载英文模型文件 'spm.en.nopretok.model'，创建英文分词器
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')
 
# 使用 SentencePieceProcessor 加载日文模型文件 'spm.ja.nopretok.model'，创建日文分词器
ja_tokenizer = spm.SentencePieceProcessor(model_file='spm.ja.nopretok.model')

 使用分词器对英文句子“All residents aged 20 to 59 years who live in Japan must enroll in public pension system.”进行编码：

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

结果：

英文句子中每个以空格为分隔符的单词都被划分开；

对日文句子"年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。"进行同样的操作： 

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

 与英文单词不同的是，日文是以词或短语进行划分，不含空格。

3 构建从 TorchText 导入的 Vocab 对象 

使用分词器和原始句子，我们接着构建从 TorchText 导入的 Vocab 对象。取决于我们的数据集和计算能力，这个过程可能需要一段时间。不同的分词器也会影响构建词汇表的时间。

# 定义一个函数 build_vocab，用于构建词汇表
def build_vocab(sentences, tokenizer):
    counter = Counter()  # 创建一个计数器对象，用于统计词频
    for sentence in sentences:
        # 对每个句子使用指定的分词器进行编码，并更新计数器
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 使用统计结果创建词汇表对象 Vocab，同时指定特殊标记
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
 
# 使用 build_vocab 函数构建日文和英文的词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)  # 构建日文词汇表
en_vocab = build_vocab(trainen, en_tokenizer)  # 构建英文词汇表

我们可以使用词汇表和分词器对象来为训练数据构建张量。 

# 定义一个函数 data_process，用于处理输入的日文和英文句子数据
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)  # 使用日文分词器和词汇表编码并转换为张量
        # 处理英文句子
        en_tensor_ = torch.tensor(
            [en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
            dtype=torch.long)  # 使用英文分词器和词汇表编码并转换为张量
        # 将处理后的张量对加入到数据列表中
        data.append((ja_tensor_, en_tensor_))
    return data
 
# 使用 data_process 函数处理训练数据集 trainja 和 trainen
train_data = data_process(trainja, trainen)

4 创建 DataLoader

DataLoader 对象可以用于在训练期间进行迭代。这里，我们将 BATCH_SIZE 设置为 16，以防止 “cuda out of memory” 错误发生。这和电脑GPU也有关系。

BATCH_SIZE = 8  # 定义批量大小为 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:
        # 处理日文句子
        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))
 
    # 对日文和英文句子批次进行填充以保证相同长度
    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
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

5 序列到序列变换器 

接下来的几个代码和注释采取自原来的PyTorch教程中https://pytorch.org/tutorials/beginner/translation_transformer.html. 除了将BATCH_SIZE和词de_vocabwhich改变ja_vocab之外，没有其他变化。

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)
        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)
        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):
        # 将位置编码和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):
        super(TokenEmbedding, self).__init__()
        
        # 创建一个词嵌入层，将词汇表大小和嵌入维度作为输入
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        # 将输入的tokens转换成长整型，并且乘以sqrt(emb_size)
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

下面，创建一个后续词掩码，以阻止目标词注意到其后续词，并且我们还需创建用于屏蔽源和目标填充标记的掩码。 

def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    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)

  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

6 设定超参数，实例化模型

用自己的显卡跑的话，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)  # 将模型移动到指定的设备
 
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和目标序列mask以及padding 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(val_iter):
        src = src.to(device)
        tgt = tgt.to(device)
 
        tgt_input = tgt[:-1, :]  # 获取目标序列输入（不包括最后一个词）
 
        # 创建源序列mask和目标序列mask以及padding 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)  # 返回平均损失

7 训练模型

用GPU跑一段时间可以出结果，但是如果用CPU的话时间会变得非常长~

# 训练模型多个epochs
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()  # 记录结束时间
    # 打印当前epoch的训练损失和训练时间
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
           f"Epoch time = {(end_time - start_time):.3f}s"))

 

最后耗时一个半小时，损失1.76

8 查看训练结果

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

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码法进行序列翻译

    Args:
        model (nn.Module): 训练好的翻译模型
        src (Tensor): 源语言序列张量
        src_mask (Tensor): 源语言序列掩码张量
        max_len (int): 生成目标语言序列的最大长度
        start_symbol (int): 目标语言序列的起始符号索引

    Returns:
        Tensor: 解码后的目标语言序列张量
    """
    # 将源序列和掩码移动到设备（GPU或CPU）
    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 = memory.to(device)
        
        # 创建 memory 掩码，初始值全为 False（即无掩码）
        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)
        
        # 解码器将当前的目标序列 ys 和编码后的 memory 作为输入
        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
        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 (nn.Module): 训练好的翻译模型
        src (str): 源语言序列（句子）
        src_vocab (Vocab): 源语言词汇表
        tgt_vocab (Vocab): 目标语言词汇表
        src_tokenizer (Tokenizer): 源语言分词器

    Returns:
        str: 翻译后的目标语言序列（句子）
    """
    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)
    
    # 创建源序列掩码张量，初始值全为 False（即无掩码）
    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>", "")

然后，翻译句子"HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。"

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

 

从trainen、trainja的列表中移除索引为 5 的元素，再返回这个被移除的元素值：

结果一致。这样就完成了日、英、中的机翻工作。