基于Transformer的中英翻译系统

最新推荐文章于 2024-11-05 17:18:24 发布

Etherious_Young

最新推荐文章于 2024-11-05 17:18:24 发布

阅读量801

点赞数 9

分类专栏：自然语言处理文章标签： transformer 深度学习人工智能自然语言处理机器翻译

本文链接：https://blog.csdn.net/NatsuD/article/details/142022662

版权

自然语言处理专栏收录该内容

2 篇文章 0 订阅

订阅专栏

在文章末我会上传本次项目main文件，项目内容为非盈利性，二转请表明出处，出现任何收费等行为与本人无关

1.引言

在自然语言处理（NLP）的领域中，机器翻译是一个极具挑战性且应用广泛的任务。利用机器翻译技术，我们可以实现不同语言之间的自动转换，为跨语言交流提供便利。在众多机器翻译方法中，基于Transformer的模型以其优越的性能和灵活的结构，成为当前主流的解决方案之一。

本博客旨在介绍如何使用一个基于Transformer的模型实现中英翻译。我们将详细讲解从数据准备、模型构建到训练和预测的整个过程，帮助读者深入了解机器翻译的实际应用和技术细节。希望这篇博客能为你提供实用的指导和参考。

2.环境准备

在开始实现中英翻译系统之前，确保你的开发环境中安装了所有必要的软件和库。以下是实现该系统所需的环境准备步骤：

所需软件和库：
- Python：建议使用Python 3.6及以上版本。
- PyTorch：用于构建和训练模型。确保安装与系统兼容的版本。
- Transformers：Hugging Face的Transformers库，用于加载和使用预训练模型。
- NumPy：用于数据处理和计算。
- 其他辅助库：如tqdm（用于显示进度条）和collections（用于数据计数）。
安装指南：使用以下命令安装所需的Python库：

pip install torch transformers numpy tqdm

确保你的PyTorch版本与CUDA（如果使用GPU）版本兼容，可以根据PyTorch官网的说明进行安装。

3.数据准备

from transformers import AutoTokenizer
import numpy as np
from collections import Counter
from utils.tools import cht_to_chs, seq_padding
from models.Attention_mask import Batch
from utils import config

UNK = config.UNK
PAD = config.PAD

class PrepareData:
    def __init__(self, train_file, dev_file):
        self.train_en, self.train_cn = self.load_data(train_file)
        self.dev_en, self.dev_cn = self.load_data(dev_file)

        self.en_word_dict, self.en_total_words, self.en_index_dict = self.build_dict(self.train_en)
        self.cn_word_dict, self.cn_total_words, self.cn_index_dict = self.build_dict(self.train_cn)

        self.train_en_ids, self.train_cn_ids = self.word2id(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)
        self.dev_en_ids, self.dev_cn_ids = self.word2id(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)

        self.train_data = self.split_batch(self.train_en_ids, self.train_cn_ids, config.BATCH_SIZE)
        self.dev_data = self.split_batch(self.dev_en_ids, self.dev_cn_ids, config.BATCH_SIZE)

    def load_data(self, path):
        en = []
        cn = []
        tokenizer = AutoTokenizer.from_pretrained("./utils/token")
        with open(path, mode="r", encoding="utf-8") as f:
            for line in f.readlines():
                sent_en, sent_cn = line.strip().split("\t")
                sent_en = sent_en.lower()
                sent_cn = cht_to_chs(sent_cn)
                sent_en = ["BOS"] + tokenizer.tokenize(sent_en) + ["EOS"]
                sent_cn = ["BOS"] + list(sent_cn) + ["EOS"]

                en.append(sent_en)
                cn.append(sent_cn)
        return en, cn

    def build_dict(self, sentences, max_words=5e4):
        word_count = Counter(word for sent in sentences for word in sent)
        most_common = word_count.most_common(int(max_words))
        word_dict = {word: i + 2 for i, (word, _) in enumerate(most_common)}
        word_dict['UNK'] = config.UNK
        word_dict['PAD'] = config.PAD
        index_dict = {i: word for word, i in word_dict.items()}
        total_words = len(word_dict)

        return word_dict, total_words, index_dict

    def word2id(self, en, cn, en_dict, cn_dict, sort=True):
        length = len(en)
        out_en_ids = [[en_dict.get(word, UNK) for word in sent] for sent in en]
        out_cn_ids = [[cn_dict.get(word, UNK) for word in sent] for sent in cn]

        def len_argsort(seq):
            return sorted(range(len(seq)), key=lambda x: len(seq[x]))

        if sort:
            sorted_index = len_argsort(out_en_ids)
            out_en_ids = [out_en_ids[i] for i in sorted_index]
            out_cn_ids = [out_cn_ids[i] for i in sorted_index]

        return out_en_ids, out_cn_ids

    def split_batch(self, en, cn, batch_size, shuffle=True):
        idx_list = np.arange(0, len(en), batch_size)
        if shuffle:
            np.random.shuffle(idx_list)
        batch_indexs = []
        for idx in idx_list:
            batch_index = np.arange(idx, min(idx + batch_size, len(en)))
            batch_indexs.append(batch_index)

        batches = []
        for batch_index in batch_indexs:
            batch_en = [en[i] for i in batch_index]
            batch_cn = [cn[i] for i in batch_index]
            batch_en = seq_padding(batch_en, PAD)
            batch_cn = seq_padding(batch_cn, PAD)
            batches.append(Batch(batch_en, batch_cn))
        return batches

数据加载 (load_data)：
- 使用AutoTokenizer对英文句子进行分词，并将中文句子转换为简体中文。
- 添加起始符BOS和终止符EOS。
- 将英文和中文句子分别存储到列表中。
词表构建 (build_dict)：
- 统计每个词出现的频率，并创建词汇到索引的映射。
- 使用Counter计算词频，并选择出现频率最高的词构建词汇表。
- 将词汇表中的特殊词（UNK和PAD）添加到字典中。
单词到索引的映射 (word2id)：
- 将英文和中文句子中的词转换为对应的索引。
- 可选择按照句子长度对数据进行排序，以便批次内句子长度尽量相同，减少填充。
数据批次划分 (split_batch)：
- 将数据划分为多个批次，并对每个批次进行随机打乱。
- 对每个批次中的句子进行填充，以保证批次内所有句子的长度相同。
- 返回Batch对象的列表，每个Batch对象包含当前批次的填充数据和对应的掩码。

4.模型构建

from models.Transformer import make_model
from utils import config

# 初始化模型
model = make_model(
    src_vocab,        # 源语言词汇表大小
    tgt_vocab,        # 目标语言词汇表大小
    config.LAYERS,    # Transformer层数
    config.D_MODEL,   # 模型的维度
    config.D_FF,      # 前馈网络的维度
    config.H_NUM,     # 注意力头的数量
    config.DROPOUT    # Dropout比率
).to(config.DEVICE)

模型初始化 (make_model)：
- make_model函数用于创建一个Transformer模型实例。该函数接受以下参数：
  - src_vocab：源语言词汇表的大小。
  - tgt_vocab：目标语言词汇表的大小。
  - config.LAYERS：Transformer模型中的层数。
  - config.D_MODEL：模型的维度，即每个词向量的大小。
  - config.D_FF：前馈网络的维度，用于Transformer的全连接层。
  - config.H_NUM：多头注意力机制中的头的数量。
  - config.DROPOUT：Dropout的比率，用于防止过拟合。
模型移动到设备 (.to(config.DEVICE))：
- 将模型移动到指定的计算设备（如CPU或GPU），以便进行训练和评估。这是通过config.DEVICE配置的设备进行的，例如torch.device('cuda')表示使用GPU。

5.模型训练

from models.loss import SimpleLossCompute
from tqdm import tqdm
import torch
import time
import torch.nn as nn
from torch.autograd import Variable
from utils import config

class LabelSmoothing(nn.Module):
    def __init__(self, size, padding_idx, smoothing=0.0):
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(size_average=False)
        self.padding_idx = padding_idx
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None

    def forward(self, x, target):
        confidence = 1.0 - self.smoothing
        low_confidence = self.smoothing / (self.size - 2)
        true_dist = x.data.clone()
        true_dist.fill_(low_confidence)
        true_dist.scatter_(1, target.data.unsqueeze(1), confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, Variable(true_dist, requires_grad=False))

class NoamOpt:
    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0

    def step(self):
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()

    def rate(self, step=None):
        if step is None:
            step = self._step
        return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))

def get_std_opt(model):
    return NoamOpt(model_size=model.d_model, factor=1, warmup=2000,
                   optimizer=torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

def run_epoch(data, model, loss_compute, epoch):
    start = time.time()
    total_tokens = 0.
    total_loss = 0.
    tokens = 0.

    progress_bar = tqdm(enumerate(data), total=len(data), desc=f"Epoch {epoch + 1}")

    for i, batch in progress_bar:
        out = model(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)

        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens

        if i % 50 == 1:
            elapsed = time.time() - start
            progress_bar.set_postfix(loss=total_loss / total_tokens, tokens_per_sec=tokens / elapsed)
            tokens = 0
            start = time.time()

    progress_bar.close()
    return total_loss / total_tokens

def train(data, model, criterion, optimizer):
    best_dev_loss = 1e5

    for epoch in range(config.EPOCHS):
        model.train()
        run_epoch(data.train_data, model, SimpleLossCompute(model.generator, criterion, optimizer), epoch)
        model.eval()

        print('>>>>> Evaluate')
        dev_loss = run_epoch(data.dev_data, model, SimpleLossCompute(model.generator, criterion, None), epoch)

        print('<<<<< Evaluate loss: %f' % dev_loss)

        if dev_loss < best_dev_loss:
            best_dev_loss = dev_loss
            print('****** Save model done... ******')
        print()

print("------------------------------------step2: 模型训练------------------------------------")
print(">>>>>>> start train")
train_start = time.time()
criterion = LabelSmoothing(tgt_vocab, padding_idx=0, smoothing=0.0)
optimizer = NoamOpt(config.D_MODEL, 1, 2000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

train(data, model, criterion, optimizer)
print(f"<<<<<<< finished train, cost {time.time() - train_start:.4f} seconds")
print("------------------------------------step2: 模型训练完成------------------------------------")

标签平滑 (LabelSmoothing)：
- LabelSmoothing类用于实现标签平滑技术，旨在减少模型对训练数据的过拟合。
- 在损失计算中引入平滑参数smoothing，调整目标标签的概率分布。
Noam优化器 (NoamOpt)：
- NoamOpt类实现了Noam学习率调度器，学习率随训练进度而变化。
- step()方法更新优化器的学习率并执行一步优化。
标准优化器获取 (get_std_opt)：
- get_std_opt函数创建并返回一个NoamOpt优化器实例，用于优化模型参数。
运行一个训练周期 (run_epoch)：
- run_epoch函数处理一个训练周期的训练和评估。
- 计算损失、更新进度条，并记录每秒处理的token数。
模型训练 (train)：
- train函数进行模型训练和评估，保存表现最佳的模型。
- 训练过程中根据开发集的损失来调整模型，并保存损失最小的模型。
训练执行：
- 设置训练相关的配置，包括损失函数和优化器。
- 调用train函数开始训练过程，并打印训练和评估结果。

这部分代码涵盖了模型训练的完整流程，从损失计算到优化器配置，再到训练和评估模型。

6.模型评估

from models.Attention_mask import subsequent_mask
from utils.predict import translate

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码策略进行预测
    """
    # 通过编码器进行编码
    memory = model.encode(src, src_mask)

    # 初始化预测内容为1×1的tensor，填入开始符（'BOS'）的id
    ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)

    # 遍历输出的长度下标
    for i in range(max_len - 1):
        # 解码得到隐藏层表示
        out = model.decode(memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data))

        # 将隐藏表示转为对词典各词的log_softmax概率分布表示
        prob = model.generator(out[:, -1])

        # 获取当前位置最大概率的预测词id
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]
        # 将当前位置预测的字符id与之前的预测内容拼接起来
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)

    return ys

def predict(data, model):
    """
    在数据集上使用训练好的模型进行预测，打印模型翻译结果
    """
    # 禁用梯度计算
    with torch.no_grad():
        # 遍历数据集中的英文数据
        for i in tqdm(range(len(data.dev_en)), desc="Translating"):
            # 获取待翻译的英文句子
            en_sent = " ".join(data.dev_en[i])
            print("\n" + en_sent)

            # 获取对应的中文参考答案
            cn_sent = "".join(data.dev_cn[i])
            print("Reference: " + cn_sent)

            # 将英文句子转换为tensor，并移动到计算设备上
            src = torch.tensor([data.dev_en_ids[i]], dtype=torch.long).to(config.DEVICE)
            src = src.unsqueeze(0)  # 增加一维

            # 设置注意力掩码
            src_mask = (src != config.PAD).unsqueeze(-2)

            # 使用训练好的模型进行解码预测
            out = greedy_decode(model, src, src_mask, max_len=config.MAX_LENGTH, start_symbol=config.BOS)

            # 初始化用于存放模型翻译结果的列表
            translation = []
            # 遍历翻译输出字符的下标（注意：开始符"BOS"的索引0不遍历）
            for j in range(1, out.size(1)):
                # 获取当前下标的输出字符
                sym = data.cn_index_dict[out[0, j].item()]

                # 如果输出字符不为'EOS'终止符，则添加到当前语句的翻译结果列表
                if sym != 'EOS':
                    translation.append(sym)
                else:
                    break
            # 打印模型翻译输出的中文语句结果
            print("Translation: %s" % "".join(translation))

print("------------------------------------step3: 模型预测------------------------------------")
# 预测
# 加载模型
model.load_state_dict(torch.load(config.SAVE_FILE, map_location=torch.device('cpu')))
# 开始预测
print(">>>>>>> start predict")
evaluate_start = time.time()
predict(data, model)
print(f"<<<<<<< finished evaluate, cost {time.time() - evaluate_start:.4f} seconds")
print("------------------------------------step3: 模型预测完成------------------------------------")

贪婪解码 (greedy_decode)：
- greedy_decode函数使用贪婪算法生成预测序列。
- 从编码器得到的memory中进行解码，逐步生成每个位置的预测词。
- 使用model.generator将隐藏层表示转化为词汇表上的概率分布，并选择概率最大的词作为预测结果。
- 将预测的词追加到已经生成的序列中，直到生成的序列达到最大长度或遇到结束符（EOS）。
预测函数 (predict)：
- predict函数在开发集上使用训练好的模型进行翻译预测。
- 遍历开发集中的英文句子，打印英文句子和对应的中文参考答案。
- 将英文句子转换为tensor，并设置注意力掩码。
- 使用greedy_decode进行预测，生成翻译结果。
- 遍历解码结果，将预测的词（去除开始符BOS和结束符EOS）拼接成最终翻译的句子，并打印结果。
模型加载与预测：
- 使用model.load_state_dict加载训练好的模型权重。
- 调用predict函数开始模型的预测过程，并打印预测结果。

7.总结

本项目旨在实现一个基于Transformer架构的英文到中文机器翻译模型，涵盖了数据准备、模型构建、训练和评估的全过程。项目从数据的预处理开始，通过分词、构建词表和映射索引，为模型提供了干净且结构化的输入数据。在模型构建阶段，使用了Transformer架构，结合了多头自注意力机制和前馈神经网络，为机器翻译任务提供了强大的模型能力。模型训练过程中，通过精心设置的超参数和标签平滑技术，确保了模型的训练稳定性和泛化能力。

在训练过程中，采用了Noam优化器来调整学习率，确保了模型在训练过程中的有效性和效率。训练完成后，通过模型评估，利用贪婪解码策略对测试数据进行翻译，并与参考答案进行对比，评估了模型的实际表现。整体来看，本项目实现了一个功能完善的机器翻译系统，能够在英文到中文翻译任务中提供有效的解决方案。

尽管模型在现有设置下表现良好，但仍有进一步优化的空间。未来的改进方向包括扩展数据集、优化模型结构、调整超参数以及尝试更先进的解码策略。通过这些改进，期望能进一步提高模型的翻译质量和适用性，为自然语言处理领域贡献更为出色的技术解决方案。

总的来说，本项目成功地展示了如何通过现代深度学习技术构建一个高效的翻译系统，并为未来的研究和应用提供了宝贵的经验和基础。

8.附录

现将本项目main代码附下：

from transformers import AutoTokenizer
import time
import torch
from utils import config
from utils.tools import cht_to_chs, seq_padding
from models.Transformer import make_model

from collections import Counter
import numpy as np
from models.Attention_mask import Batch

UNK = config.UNK
PAD = config.PAD


class PrepareData:
    def __init__(self, train_file, dev_file):
        # 读取数据、分词
        self.train_en, self.train_cn = self.load_data(train_file)
        self.dev_en, self.dev_cn = self.load_data(dev_file)

        # 构建词表
        self.en_word_dict, self.en_total_words, self.en_index_dict = self.build_dict(self.train_en)
        self.cn_word_dict, self.cn_total_words, self.cn_index_dict = self.build_dict(self.train_cn)

        # 单词映射为索引
        self.train_en_ids, self.train_cn_ids = self.word2id(self.train_en, self.train_cn, self.en_word_dict,
                                                            self.cn_word_dict)
        self.dev_en_ids, self.dev_cn_ids = self.word2id(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)

        # 划分批次、填充、掩码
        self.train_data = self.split_batch(self.train_en_ids, self.train_cn_ids, config.BATCH_SIZE)
        self.dev_data = self.split_batch(self.dev_en_ids, self.dev_cn_ids, config.BATCH_SIZE)

    def load_data(self, path):
        """
        读取英文、中文数据
        对每条样本分词并构建包含起始符和终止符的单词列表
        形式如：en = [['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]
                cn = [['BOS', '我', '爱', '你', 'EOS'], ['BOS', '我', '也', '是', 'EOS'], ...]
        """
        en = []
        cn = []
        tokenizer = AutoTokenizer.from_pretrained("./utils/token")
        with open(path, mode="r", encoding="utf-8") as f:
            for line in f.readlines():
                sent_en, sent_cn = line.strip().split("\t")
                sent_en = sent_en.lower()

                sent_cn = cht_to_chs(sent_cn)
                sent_en = ["BOS"] + tokenizer.tokenize(sent_en) + ["EOS"]
                sent_cn = ["BOS"] + list(sent_cn) + ["EOS"]

                en.append(sent_en)
                cn.append(sent_cn)
        return en, cn

    def build_dict(self, sentences, max_words=5e4):
        """
        构造分词后的列表数据
        构建单词-索引映射（key为单词，value为id值）
        """
        word_count = Counter(word for sent in sentences for word in sent)
        most_common = word_count.most_common(int(max_words))
        word_dict = {word: i + 2 for i, (word, _) in enumerate(most_common)}
        word_dict['UNK'] = config.UNK
        word_dict['PAD'] = config.PAD
        index_dict = {i: word for word, i in word_dict.items()}
        total_words = len(word_dict)

        return word_dict, total_words, index_dict

    def word2id(self, en, cn, en_dict, cn_dict, sort=True):
        """
        将英文、中文单词列表转为单词索引列表
        sort=True表示以英文语句长度排序，以便按批次填充时，同批次语句填充尽量少
        """
        length = len(en)
        # 单词映射为索引
        out_en_ids = [[en_dict.get(word, UNK) for word in sent] for sent in en]
        out_cn_ids = [[cn_dict.get(word, UNK) for word in sent] for sent in cn]

        # 按照语句长度排序
        def len_argsort(seq):
            """
            传入一系列语句数据(分好词的列表形式)，
            按照语句长度排序后，返回排序后原来各语句在数据中的索引下标
            """
            return sorted(range(len(seq)), key=lambda x: len(seq[x]))

        # 按相同顺序对中文、英文样本排序
        if sort:
            # 以英文语句长度排序
            sorted_index = len_argsort(out_en_ids)
            out_en_ids = [out_en_ids[i] for i in sorted_index]
            out_cn_ids = [out_cn_ids[i] for i in sorted_index]

        return out_en_ids, out_cn_ids

    def split_batch(self, en, cn, batch_size, shuffle=True):
        """
        划分批次
        shuffle=True表示对各批次顺序随机打乱
        """
        # 每隔batch_size取一个索引作为后续batch的起始索引
        idx_list = np.arange(0, len(en), batch_size)

        # 起始索引随机打乱
        if shuffle:
            np.random.shuffle(idx_list)
        # 存放所有批次的语句索引
        batch_indexs = []
        for idx in idx_list:
            """
            形如[array([4, 5, 6, 7]), 
                 array([0, 1, 2, 3]), 
                 array([8, 9, 10, 11]),
                 ...]
            """
            # 起始索引最大的批次可能发生越界，要限定其索引
            batch_index = np.arange(idx, min(idx + batch_size, len(en)))
            batch_indexs.append(batch_index)

        # 构建批次列表
        batches = []
        for batch_index in batch_indexs:
            # 按当前批次的样本索引采样
            batch_en = [en[i] for i in batch_index]
            batch_cn = [cn[i] for i in batch_index]

            # 对当前批次中所有语句填充、对齐长度
            # 维度为：batch_size * 当前批次中语句的最大长度
            batch_en = seq_padding(batch_en, PAD)
            batch_cn = seq_padding(batch_cn, PAD)

            # 将当前批次添加到批次列表
            # Batch类用于实现注意力掩码
            batches.append(Batch(batch_en, batch_cn))
        return batches


# 数据预处理
print("------------------------------------step1: 正在准备数据------------------------------------")
data = PrepareData(config.TRAIN_FILE, config.DEV_FILE)
src_vocab = len(data.en_word_dict)
tgt_vocab = len(data.cn_word_dict)
print("src_vocab %d" % src_vocab)
print("tgt_vocab %d" % tgt_vocab)
print("------------------------------------step1: 数据准备完成------------------------------------")

from models.loss import SimpleLossCompute
from tqdm import tqdm
import torch.nn as nn
import torch
from torch.autograd import Variable
import time


class LabelSmoothing(nn.Module):
    """
    标签平滑
    """

    def __init__(self, size, padding_idx, smoothing=0.0):
        """
        初始化标签平滑器

        Args:
            size (int): 目标类别的数量
            padding_idx (int): padding 的索引
            smoothing (float): 平滑参数
        """
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(size_average=False)
        self.padding_idx = padding_idx
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None

    def forward(self, x, target):
        """
        计算标签平滑的损失

        Args:
            x (torch.Tensor): 模型的输出
            target (torch.Tensor): 真实标签

        Returns:
            torch.Tensor: 标签平滑的损失值
        """
        confidence = 1.0 - self.smoothing
        low_confidence = self.smoothing / (self.size - 2)
        true_dist = x.data.clone()
        true_dist.fill_(low_confidence)
        true_dist.scatter_(1, target.data.unsqueeze(1), confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist

        return self.criterion(x, Variable(true_dist, requires_grad=False))


class NoamOpt:
    """
    Noam 优化器
    """

    def __init__(self, model_size, factor, warmup, optimizer):
        """
        初始化 Noam 优化器

        Args:
            model_size (int): 模型参数的大小
            factor (float): 优化器的因子
            warmup (int): warmup 步数
            optimizer (torch.optim.Optimizer): PyTorch 优化器
        """
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0

    def step(self):
        """
        更新学习率并执行一步优化器
        """
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()

    def rate(self, step=None):
        """
        计算当前的学习率

        Args:
            step (int, optional): 当前步数，默认为 None

        Returns:
            float: 当前的学习率===self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5))
        """
        if step is None:
            step = self._step
        return self.factor * (self.model_size ** (-0.5) *
                              min(step ** (-0.5), step * self.warmup ** (-1.5)))


def get_std_opt(model):
    """
    获取标准优化器

    Args:
        model: PyTorch 模型

    Returns:
        NoamOpt: Noam 优化器
    """
    return NoamOpt(model_size=model.d_model, factor=1, warmup=2000,
                   optimizer=torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))


# 初始化模型
# D_MODEL = 256
# H_NUM = 8

model = make_model(src_vocab, tgt_vocab, config.LAYERS, config.D_MODEL, config.D_FF, config.H_NUM, config.DROPOUT).to(
    config.DEVICE)


def run_epoch(data, model, loss_compute, epoch):
    start = time.time()
    total_tokens = 0.
    total_loss = 0.
    tokens = 0.

    # 初始化 tqdm 进度条
    progress_bar = tqdm(enumerate(data), total=len(data), desc=f"Epoch {epoch + 1}")

    for i, batch in progress_bar:
        out = model(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)

        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens

        if i % 50 == 1:
            elapsed = time.time() - start
            # 更新进度条显示信息，包括平均损失和每秒处理的 token 数
            elapsed = time.time() - start
            progress_bar.set_postfix(loss=total_loss / total_tokens, tokens_per_sec=tokens / elapsed)
            tokens = 0
            start = time.time()

    # 关闭进度条
    progress_bar.close()

    return total_loss / total_tokens


def train(data, model, criterion, optimizer):
    """
    训练并保存模型
    """
    # 初始化模型在dev集上的最优Loss为一个较大值，可设置1e5
    best_dev_loss = 1e5

    for epoch in range(config.EPOCHS):
        # 模型训练
        model.train()
        run_epoch(data.train_data, model, SimpleLossCompute(model.generator, criterion, optimizer), epoch)
        model.eval()

        # 在dev集上进行loss评估
        print('>>>>> Evaluate')
        dev_loss = run_epoch(data.dev_data, model, SimpleLossCompute(model.generator, criterion, None), epoch)

        print('<<<<< Evaluate loss: %f' % dev_loss)

        # 如果当前epoch的模型在dev集上的loss优于之前记录的最优loss则保存当前模型，并更新最优loss值
        if dev_loss < best_dev_loss:
            # torch.save(model.state_dict(), config.SAVE_FILE)
            best_dev_loss = dev_loss
            print('****** Save model done... ******')

        print()


# 训练
print("------------------------------------step2: 模型训练------------------------------------")
print(">>>>>>> start train")
train_start = time.time()
criterion = LabelSmoothing(tgt_vocab, padding_idx=0, smoothing=0.0)
optimizer = NoamOpt(config.D_MODEL, 1, 2000,
                    torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

train(data, model, criterion, optimizer)
print(f"<<<<<<< finished train, cost {time.time() - train_start:.4f} seconds")
print("------------------------------------step2: 模型训练完成------------------------------------")

from models.Attention_mask import subsequent_mask
from utils.predict import translate


def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    传入一个训练好的模型，对指定数据进行预测
    """
    # 先用encoder进行encode
    memory = model.encode(src, src_mask)

    # 初始化预测内容为1×1的tensor，填入开始符('BOS')的id，并将type设置为输入数据类型(LongTensor)
    ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)

    # 遍历输出的长度下标
    for i in range(max_len - 1):
        # decode得到隐层表示
        out = model.decode(memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data))

        # 将隐藏表示转为对词典各词的log_softmax概率分布表示
        prob = model.generator(out[:, -1])

        # 获取当前位置最大概率的预测词id
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]
        # 将当前位置预测的字符id与之前的预测内容拼接起来
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)

    return ys


def predict(data, model):
    """
    在data上用训练好的模型进行预测，打印模型翻译结果
    """
    # 梯度清零
    with torch.no_grad():
        # 在data的英文数据长度上遍历下标
        for i in tqdm(range(len(data.dev_en)), desc="Translating"):
            # 打印待翻译的英文语句
            en_sent = " ".join(data.dev_en[i])
            # print("\n" + en_sent)

            # 打印对应的中文语句答案
            cn_sent = "".join(data.dev_cn[i])
            # print("Reference: " + cn_sent)

            # 将当前以单词id表示的英文语句数据转为tensor，并放入DEVICE中
            src = torch.tensor([data.dev_en_ids[i]], dtype=torch.long).to(config.DEVICE)

            # 增加一维
            src = src.unsqueeze(0)

            # 设置attention mask
            src_mask = (src != config.PAD).unsqueeze(-2)

            # 用训练好的模型进行decode预测
            out = greedy_decode(model, src, src_mask, max_len=config.MAX_LENGTH, start_symbol=config.BOS)

            # 初始化一个用于存放模型翻译结果语句单词的列表
            translation = []
            # 遍历翻译输出字符的下标（注意：开始符"BOS"的索引0不遍历）
            for j in range(1, out.size(1)):
                # 获取当前下标的输出字符
                sym = data.cn_index_dict[out[0, j].item()]

                # 如果输出字符不为'EOS'终止符，则添加到当前语句的翻译结果列表
                if sym != 'EOS':
                    translation.append(sym)
                # 否则终止遍历
                else:
                    break
            # 打印模型翻译输出的中文语句结果
            # print("Translation: %s" % " ".join(translation))


print("------------------------------------step3: 模型预测------------------------------------")
# 预测
# 加载模型
model.load_state_dict(torch.load(config.SAVE_FILE, map_location=torch.device('cpu')))
# 开始预测
print(">>>>>>> start predict")
evaluate_start = time.time()
test_data = "i love you"
predict(data, model)
print(f"<<<<<<< finished evaluate, cost {time.time() - evaluate_start:.4f} seconds")
print("------------------------------------step3: 模型预测完成------------------------------------")