2 - Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 - 使用RNN编码器-解码器学习短语表示进行统计机器翻译\n",
    "\n",
    "在这个关于使用PyTorch和TorchText的序列到序列模型的第二个笔记本中，我们将实现 [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://arxiv.org/abs/1406.1078) 中的模型。该模型将在编码器和解码器中仅使用单层RNN的情况下实现改进的测试困惑度。\n",
    "\n",
    "## 介绍\n",
    "\n",
    "让我们先回顾一下一般的编码器-解码器模型。\n",
    "\n",
    "![](assets/seq2seq1.png)\n",
    "\n",
    "我们使用编码器（绿色）在嵌入的源序列（黄色）上创建上下文向量（红色）。然后，我们使用该上下文向量与解码器（蓝色）和线性层（紫色）生成目标句子。\n",
    "\n",
    "在先前的模型中，我们使用了多层LSTM作为编码器和解码器。\n",
    "\n",
    "![](assets/seq2seq4.png)\n",
    "\n",
    "先前模型的一个缺点是解码器试图将大量信息压缩到隐藏状态中。在解码过程中，隐藏状态将需要包含关于整个源序列的信息，以及到目前为止已解码的所有标记的信息。通过减轻一些这种信息压缩，我们可以创建一个更好的模型！\n",
    "\n",
    "我们还将使用GRU（门控循环单元）而不是LSTM（长短时记忆）。为什么？主要因为这就是他们在论文中所做的（该论文还引入了GRU），并且因为我们上次使用了LSTM。要了解GRU（和LSTM）与标准RNN有何不同，请查看[此链接](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)。GRU是否比LSTM更好？[研究](https://arxiv.org/abs/1412.3555)表明它们几乎相同，两者都比标准RNN更好。\n",
    "\n",
    "## 准备数据\n",
    "\n",
    "所有数据准备步骤都（几乎）与上次相同，因此我们将非常简要地说明每个代码块的作用。有关简要回顾，请参见上一个笔记本。\n",
    "\n",
    "我们将导入Mindspore、spaCy和一些标准模块。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore as ms\n",
    "import mindspore.context as context\n",
    "import mindspore.nn as nn\n",
    "import mindspore.ops as ops\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import spacy\n",
    "import datasets\n",
    "from tqdm.notebook import tqdm\n",
    "import evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们将设置所有可能的随机种子以确保结果的可重复性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 1234\n",
    "\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "ms.set_seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "实例化我们的德语和英语spaCy模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/root/.cache/huggingface/datasets/bentrevett___json/bentrevett--multi30k-8cca6da32304eb4d/0.0.0/8bb11242116d547c741b2e8a1f18598ffdd40a1d4f2a2872c7a28b697434bc96)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "edbb1f00b2e04f0c82ba0020f8193088",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = datasets.load_dataset(\"bentrevett/multi30k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, valid_data, test_data = dataset[\"train\"], dataset[\"validation\"], dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_nlp = spacy.load(\"en_core_web_sm\")\n",
    "de_nlp = spacy.load(\"de_core_news_sm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "之前我们颠倒了源（德语）句子，但在我们要实现的论文中，他们没有这样做，因此我们也不会这样做。\n",
    "现在编写一个用于将分词器应用于每个数据拆分中所有示例的函数，并进行一些其他处理。\n",
    "\n",
    "此函数接受来自 `Dataset` 对象的示例，应用英语和德语 spaCy 模型的分词器，将标记列表修剪到最大长度，可选地将每个标记转换为小写，然后将序列开始和序列结束标记附加到标记列表的开头和结尾。\n",
    "\n",
    "这个函数将与 `Dataset` 的 `map` 方法一起使用，该方法需要返回一个包含每个示例中存储输出的特征的名称的字典。由于输出特征名称 \"en_tokens\" 和 \"de_tokens\" 在示例中尚未存在（我们只有 \"en\" 和 \"de\" 特征），因此这将在每个示例中创建两个新特征。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_example(\n",
    "    example,\n",
    "    en_nlp,\n",
    "    de_nlp,\n",
    "    max_length,\n",
    "    lower,\n",
    "    sos_token,\n",
    "    eos_token\n",
    "):\n",
    "    en_tokens = [token.text for token in en_nlp.tokenizer(example[\"en\"])][:max_length]\n",
    "    de_tokens = [token.text for token in de_nlp.tokenizer(example[\"de\"])][:max_length]\n",
    "    if lower:\n",
    "        en_tokens = [token.lower() for token in en_tokens]\n",
    "        de_tokens = [token.lower() for token in de_tokens]\n",
    "    en_tokens = [sos_token] + en_tokens + [eos_token]\n",
    "    de_tokens = [sos_token] + de_tokens + [eos_token]\n",
    "    return {\"en_tokens\": en_tokens, \"de_tokens\": de_tokens}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用以下方式使用 `map` 方法应用 `tokenize_example` 函数。\n",
    "\n",
    "在这里，我们将所有序列修剪为最大长度为 1000 个标记，将每个标记转换为小写，并分别使用 `<sos>` 和 `<eos>` 作为序列的开始和结束标记。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "222f2532710a40ba9a564d2fc147c959",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/29000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3f2081134d144d16b32396a98d5d7a08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1014 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa766ddbdf244ce3821ae5008376ec29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "max_length = 1_000\n",
    "lower = True\n",
    "sos_token = \"<sos>\"\n",
    "eos_token = \"<eos>\"\n",
    "\n",
    "fn_kwargs = {\n",
    "    \"en_nlp\": en_nlp, \n",
    "    \"de_nlp\": de_nlp, \n",
    "    \"max_length\": max_length,\n",
    "    \"lower\": lower,\n",
    "    \"sos_token\": sos_token,\n",
    "    \"eos_token\": eos_token,\n",
    "}\n",
    "\n",
    "\n",
    "train_data = train_data.map(tokenize_example, fn_kwargs=fn_kwargs)\n",
    "valid_data = valid_data.map(tokenize_example, fn_kwargs=fn_kwargs)\n",
    "test_data = test_data.map(tokenize_example, fn_kwargs=fn_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们可以查看一个示例，确认已添加了两个新特征；这两个特征都是附加了序列开始/结束标记的小写字符串列表。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'en': 'Two young, White males are outside near many bushes.',\n",
       " 'de': 'Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.',\n",
       " 'en_tokens': ['<sos>',\n",
       "  'two',\n",
       "  'young',\n",
       "  ',',\n",
       "  'white',\n",
       "  'males',\n",
       "  'are',\n",
       "  'outside',\n",
       "  'near',\n",
       "  'many',\n",
       "  'bushes',\n",
       "  '.',\n",
       "  '<eos>'],\n",
       " 'de_tokens': ['<sos>',\n",
       "  'zwei',\n",
       "  'junge',\n",
       "  'weiße',\n",
       "  'männer',\n",
       "  'sind',\n",
       "  'im',\n",
       "  'freien',\n",
       "  'in',\n",
       "  'der',\n",
       "  'nähe',\n",
       "  'vieler',\n",
       "  'büsche',\n",
       "  '.',\n",
       "  '<eos>']}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们将为源语言和目标语言构建*词汇表*。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "class Vocab:\n",
    "    \"\"\"一个词汇表的实现\"\"\"\n",
    "    def __init__(self, tokens:list, min_freq=0, reserved_tokens:list=None) -> None:\n",
    "        self.default_index = None\n",
    "        if tokens is not None:\n",
    "            # 当第一个条件满足时，就不会跳到第二个判断，避免了空列表报错的情况。\n",
    "            if len(tokens)!=0 and isinstance(tokens[0], list):\n",
    "                tokens = [i for line in tokens for i in line]\n",
    "        else:\n",
    "            tokens = []\n",
    "        if reserved_tokens is None:\n",
    "            reserved_tokens = []\n",
    "        counter=collections.Counter(tokens)\n",
    "        # 按出现词频从高到低排序\n",
    "        self._token_freqs = sorted(counter.items(), key=lambda x:x[1], reverse=True)\n",
    "        # 通过列表,利用序号访问词元。\n",
    "        self.idx_to_token = [] + reserved_tokens # 未知词元<unk>的索引为0, 保留词元排在最前\n",
    "        self.token_to_idx = {\n",
    "            i: k\n",
    "            for k, i in enumerate(self.idx_to_token) \n",
    "        }\n",
    "        \n",
    "        for token, freq in self._token_freqs:\n",
    "            if freq < min_freq:  # 过滤掉出现频率低于要求的词\n",
    "                break\n",
    "            if token not in self.token_to_idx:  \n",
    "                self.idx_to_token.append(token)\n",
    "                self.token_to_idx[token] = len(self.idx_to_token) - 1\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.idx_to_token)\n",
    "    \n",
    "    def __getitem__(self, input_tokens):\n",
    "        \"\"\"输入单字串或序列, 将其全部转化为序号编码\"\"\"\n",
    "        if isinstance(input_tokens, str):\n",
    "            out =  self.token_to_idx.get(input_tokens, self.default_index)\n",
    "            if out is None:\n",
    "                raise Exception('Please call \"set_default_index\" before getting unknown index')\n",
    "            return out\n",
    "        return [self.__getitem__(token) for token in input_tokens]\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        show_items = 5 if len(self) > 5 else len(self)\n",
    "        out = f\"<Vocab with {len(self)} tokens: \"\n",
    "        for i in range(show_items):\n",
    "            out += f'\"{self.idx_to_token[i]}\", '\n",
    "        out += \"...>\"\n",
    "        return out\n",
    "\n",
    "    def __contains__(self, token:str) -> bool:\n",
    "        return token in self.idx_to_token\n",
    "\n",
    "    def to_tokens(self, input_keys):\n",
    "        \"\"\"输入单s索引或序列, 将其全部转化为词元\"\"\"\n",
    "        if isinstance(input_keys, int):\n",
    "            return self.idx_to_token[input_keys] if input_keys < len(self) else self.idx_to_token[0]\n",
    "        elif isinstance(input_keys, (list, tuple)):\n",
    "            return [self.to_tokens(keys) for keys in input_keys]\n",
    "        else:\n",
    "            return self.idx_to_token[0]\n",
    "    \n",
    "    def get_itos(self):\n",
    "        return self.idx_to_token\n",
    "    \n",
    "    def get_stoi(self):\n",
    "        return self.token_to_idx\n",
    "    \n",
    "    def set_default_index(self, idx):\n",
    "        if isinstance(idx, int):\n",
    "            self.default_index = idx\n",
    "        else:\n",
    "            raise Exception(f\"Only type int allowed, got {type(idx)}\")\n",
    "    \n",
    "    def lookup_indices(self, input_tokens):\n",
    "        return self.__getitem__(input_tokens)\n",
    "    \n",
    "    def lookup_tokens(self, idx):\n",
    "        return self.to_tokens(idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "创建我们的词汇表，将所有出现不到两次的标记转换为 `<unk>` 标记。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_freq = 2\n",
    "unk_token = \"<unk>\"\n",
    "pad_token = \"<pad>\"\n",
    "\n",
    "special_tokens = [\n",
    "    unk_token,\n",
    "    pad_token,\n",
    "    sos_token,\n",
    "    eos_token,\n",
    "]\n",
    "\n",
    "def build_vocab_from_iterator(tokens, min_freq, specials):\n",
    "    return Vocab(tokens, min_freq, specials)\n",
    "\n",
    "en_vocab = build_vocab_from_iterator(\n",
    "    train_data[\"en_tokens\"],\n",
    "    min_freq=min_freq,\n",
    "    specials=special_tokens,\n",
    ")\n",
    "\n",
    "de_vocab = build_vocab_from_iterator(\n",
    "    train_data[\"de_tokens\"],\n",
    "    min_freq=min_freq,\n",
    "    specials=special_tokens,  \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "以程序化的方式获取 `<unk>` 标记的索引，并检查我们的两个词汇表是否对未知和填充标记具有相同的索引，因为这将简化后面的一些代码。\n",
    "\n",
    "我们还获取 `<pad>` 标记的索引，因为我们稍后将使用它。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert en_vocab[unk_token] == de_vocab[unk_token]\n",
    "assert en_vocab[pad_token] == de_vocab[pad_token]\n",
    "\n",
    "unk_index = en_vocab[unk_token]\n",
    "pad_index = en_vocab[pad_token]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用 `set_default_index` 方法，我们可以设置在尝试获取词汇表范围之外的标记的索引时返回的值。在这种情况下，即未知标记 `<unk>` 的索引。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_vocab.set_default_index(unk_index)\n",
    "de_vocab.set_default_index(unk_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "创建一个 `numericalize_example` 函数，将在数据集的 `map` 方法中使用。这将使用词汇表“数字化”（一个说法是将标记转换为索引）每个示例中的标记，并将结果返回到新的 \"en_ids\" 和 \"de_ids\" 特征中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def numericalize_example(example, en_vocab:Vocab, de_vocab:Vocab):\n",
    "    en_ids = en_vocab.lookup_indices(example[\"en_tokens\"])\n",
    "    de_ids = de_vocab.lookup_indices(example[\"de_tokens\"])\n",
    "    return {\"en_ids\": en_ids, \"de_ids\": de_ids}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "应用 `numericalize_example` 函数，将我们的词汇表传递给 `fn_kwargs` 参数中的字典。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63efe272fecf4194b0bc01f94760fda4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/29000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0583cbf640d4bc3b8ef50b1882f7b15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1014 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c6c30da223f4f9793727078006b1f03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fn_kwargs = {\n",
    "    \"en_vocab\": en_vocab, \n",
    "    \"de_vocab\": de_vocab\n",
    "}\n",
    "\n",
    "train_data = train_data.map(numericalize_example, fn_kwargs=fn_kwargs)\n",
    "valid_data = valid_data.map(numericalize_example, fn_kwargs=fn_kwargs)\n",
    "test_data = test_data.map(numericalize_example, fn_kwargs=fn_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`datasets` 库使用 `Dataset` 类为我们处理的另一件事是将特征转换为正确的类型。每个示例中的索引当前是基本的 Python 整数。但是，为了在 mindspore 中使用它们，它们需要转换为 mindspore 张量。我们可以在将它们传递到模型之前将它们转换，但是现在这样做更加方便。\n",
    "\n",
    "`with_format` 方法将由 `columns` 参数指示的特征转换为给定的 `type`。在这里，我们将类型指定为 \"numpy\"（用于 mindspore），并将列指定为 \"en_ids\" 和 \"de_ids\"（我们要转换为 mindspore 张量的特征）。默认情况下，`with_format` 将删除未在传递给 `columns` 的特征列表中的特征。我们希望保留这些特征，可以使用 `output_all_columns=True`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_type = \"numpy\"\n",
    "format_columns = [\"en_ids\", \"de_ids\"]\n",
    "\n",
    "train_data = train_data.with_format(\n",
    "    type=data_type, \n",
    "    columns=format_columns, \n",
    "    output_all_columns=True\n",
    ")\n",
    "\n",
    "valid_data = valid_data.with_format(\n",
    "    type=data_type, \n",
    "    columns=format_columns, \n",
    "    output_all_columns=True,\n",
    ")\n",
    "\n",
    "test_data = test_data.with_format(\n",
    "    type=data_type, \n",
    "    columns=format_columns, \n",
    "    output_all_columns=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "是创建数据加载器。这些加载器可以进行迭代，以返回一批数据，每个批次都是一个包含数值化的英语和德语句子的字典（它们也已被填充）作为 Mindspore 张量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sequence(sequences:list, padding_value:int):\n",
    "    '''将序列填充到等长并返回mindspore张量'''\n",
    "    # Find the length of the longest sequence in the batch\n",
    "    max_length = max(len(seq) for seq in sequences)\n",
    "    padded_sequences:ms.Tensor = ops.full((len(sequences), max_length), padding_value, dtype=ms.int64)\n",
    "    # Copy the sequences into the padded array\n",
    "    for i, seq in enumerate(sequences):\n",
    "        padded_sequences[i, :len(seq)] = ms.tensor(seq).astype(np.int64)\n",
    "    # 换轴，保证输出为时序优先\n",
    "    padded_sequences = padded_sequences.swapaxes(0, 1)\n",
    "    return padded_sequences  \n",
    "\n",
    "def get_collate_fn(pad_index):\n",
    "    \n",
    "    def collate_fn(batch):\n",
    "        batch_en_ids = [example[\"en_ids\"] for example in batch]\n",
    "        batch_de_ids = [example[\"de_ids\"] for example in batch]\n",
    "        batch_en_ids = pad_sequence(batch_en_ids, padding_value=pad_index)\n",
    "        batch_de_ids = pad_sequence(batch_de_ids, padding_value=pad_index)\n",
    "        batch = {\n",
    "            \"en_ids\": batch_en_ids,\n",
    "            \"de_ids\": batch_de_ids,\n",
    "        }\n",
    "        return batch\n",
    "    \n",
    "    return collate_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "编写使用  `DataLoader` 类创建的数据加载器的函数。\n",
    "\n",
    "`get_data_loader` 使用 `Dataset`、批次大小、填充标记索引（用于在 `collate_fn` 中创建批次）以及一个布尔值（决定在迭代数据加载器时是否应该对示例进行洗牌）创建。\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader:\n",
    "    def __init__(self, source, batch_size, shuffle=False, per_batch_map=None):\n",
    "        self.source = source\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.per_batch_map = per_batch_map\n",
    "        self.indices = np.arange(len(source))\n",
    "        self.current_index = 0\n",
    "\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(self.indices)\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        if self.current_index >= len(self.source):\n",
    "            self.current_index = 0\n",
    "            raise StopIteration\n",
    "\n",
    "        batch_indices = self.indices[self.current_index:self.current_index + self.batch_size]\n",
    "        batch_data = [self.source[int(i)] for i in batch_indices]\n",
    "\n",
    "        if self.per_batch_map:\n",
    "            batch_data = self.per_batch_map(batch_data)\n",
    "\n",
    "        self.current_index += self.batch_size\n",
    "        return batch_data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.source) // self.batch_size\n",
    "\n",
    "\n",
    "def get_data_loader(dataset, batch_size, pad_index, shuffle=False):\n",
    "    collate_fn = get_collate_fn(pad_index)\n",
    "    dataloader = DataLoader(dataset, batch_size, shuffle=shuffle, per_batch_map=collate_fn)\n",
    "    \n",
    "    return dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "\n",
    "train_data_loader = get_data_loader(train_data, batch_size, pad_index, shuffle=True)\n",
    "valid_data_loader = get_data_loader(valid_data, batch_size, pad_index)\n",
    "test_data_loader = get_data_loader(test_data, batch_size, pad_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRU(nn.Cell):\n",
    "    def __init__(self, embedding_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.para  = self.get_gru_params(embedding_dim, hidden_dim)\n",
    "        self.state = self.init_gru_state(2, hidden_dim)\n",
    "    \n",
    "    def construct(self, src):\n",
    "        # src = [src length, batch size]\n",
    "        y, state = self.gru_forward(src, self.state, self.para)\n",
    "        return hidden\n",
    "    \n",
    "    def get_gru_params(self, embed_size:int, num_hiddens:int):\n",
    "        \"\"\"获取GRU中的所有参数:\n",
    "        `W_xz, W_hz, b_z, W_xr, W_hr, b_r, W_xh, W_hh, b_h, W_hq, b_q`\n",
    "        \"\"\"\n",
    "        num_inputs = num_outputs = embed_size\n",
    "        def wwb():\n",
    "            return (ops.randn((num_inputs, num_hiddens))/100,\n",
    "                    ops.randn((num_hiddens, num_hiddens))/100,\n",
    "                    ops.zeros((num_hiddens)),\n",
    "            )\n",
    "        W_xz, W_hz, b_z = wwb()\n",
    "        W_xr, W_hr, b_r = wwb()\n",
    "        W_xh, W_hh, b_h = wwb()\n",
    "        W_hq = ops.randn((num_hiddens, num_outputs))/100\n",
    "        b_q  = ops.zeros((num_outputs))\n",
    "        params = [W_xz, W_hz, b_z, W_xr, W_hr, b_r, W_xh, W_hh, b_h, W_hq, b_q]\n",
    "        return params\n",
    "    \n",
    "    def init_gru_state(self, batch_size, num_hiddens):\n",
    "        \"\"\" 初始化隐变量 \"\"\"\n",
    "        return ops.zeros((batch_size, num_hiddens))\n",
    "    \n",
    "    def gru_forward(self, inputs, state, params):\n",
    "        W_xz, W_hz, b_z, W_xr, W_hr, b_r, W_xh, W_hh, b_h, W_hq, b_q = params\n",
    "        H = state\n",
    "        outputs = []\n",
    "        for x in inputs:\n",
    "            Z = ops.sigmoid((x @ W_xz) + (H @ W_hz) + b_z)\n",
    "            R = ops.sigmoid((x @ W_xr) + (H @ W_hr) + b_r)\n",
    "            H_tilda = ops.tanh((x @ W_xh) + ((R * H) @ W_hh) + b_h )\n",
    "            H = Z * H + (1 - Z) * H_tilda\n",
    "            Y = H @ W_hq + b_q\n",
    "            outputs.append(Y)\n",
    "        return ops.cat(outputs, axis=0), H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建 Seq2Seq 模型\n",
    "\n",
    "### 编码器（Encoder）\n",
    "\n",
    "编码器与先前的模型类似，多层 LSTM 被单层 GRU 替换。我们也不将 dropout 作为参数传递给 GRU，因为 dropout 是在多层 RNN 的每一层之间使用的。因为我们只有一层，如果我们尝试将 dropout 值传递给它，PyTorch 将显示警告。\n",
    "\n",
    "关于 GRU 还要注意的一点是，它只需要并返回一个隐藏状态，没有像 LSTM 中那样的 cell 状态。\n",
    "\n",
    "$$\\begin{align*}\n",
    "h_t &= \\text{GRU}(e(x_t), h_{t-1})\\\\\n",
    "(h_t, c_t) &= \\text{LSTM}(e(x_t), h_{t-1}, c_{t-1})\\\\\n",
    "h_t &= \\text{RNN}(e(x_t), h_{t-1})\n",
    "\\end{align*}$$\n",
    "\n",
    "从上面的方程式中，看起来 RNN 和 GRU 是相同的。然而，在 GRU 内部，有一些控制信息流入和流出隐藏状态的 *门控机制*（类似于 LSTM）。再次说明，欲了解更多信息，请查看[这篇](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)博文。\n",
    "\n",
    "编码器的其余部分应该与上个教程非常相似，它接收一个序列 $X = \\{x_1, x_2, ... , x_T\\}$，通过嵌入层传递它，循环计算隐藏状态 $H = \\{h_1, h_2, ..., h_T\\}$，并返回一个上下文向量（最终隐藏状态）$z=h_T$。\n",
    "\n",
    "$$h_t = \\text{EncoderGRU}(e(x_t), h_{t-1})$$\n",
    "\n",
    "这与通用 seq2seq 模型的编码器相同，所有的“魔法”都发生在 GRU 内部（绿色）。\n",
    "\n",
    "![](assets/seq2seq5.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Cell):\n",
    "    def __init__(self, input_dim, embedding_dim, hidden_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "        self.rnn = nn.GRU(embedding_dim, hidden_dim)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "    def construct(self, src):\n",
    "        # src = [src length, batch size]\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        # embedded = [src length, batch size, embedding dim]\n",
    "        outputs, hidden = self.rnn(embedded) # no cell state in GRU!\n",
    "        # outputs = [src length, batch size, hidden dim * n directions]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # outputs are always from the top hidden layer\n",
    "        return hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 解码器 (Decoder)\n",
    "\n",
    "解码器是实现与先前模型显著不同之处，我们通过一些方法减轻了信息压缩。\n",
    "\n",
    "解码器中的 GRU 不再仅仅将嵌入目标token $d(y_t)$ 和先前隐藏状态 $s_{t-1}$ 作为输入，它还接受上下文向量 $z$。\n",
    "\n",
    "$$s_t = \\text{DecoderGRU}(d(y_t), s_{t-1}, z)$$\n",
    "\n",
    "请注意，这个上下文向量 $z$ 没有 $t$ 下标，这意味着我们在解码器的每个时间步骤中都重用编码器返回的相同上下文向量。\n",
    "\n",
    "在之前的模型中，我们使用线性层 $f$ 仅使用该时间步长的顶层解码器隐藏状态 $s_t$ 预测下一个token $\\hat{y}_{t+1}$，即 $\\hat{y}_{t+1}=f(s_t^L)$。现在，我们还将当前token的嵌入 $d(y_t)$ 和上下文向量 $z$ 传递到线性层。\n",
    "\n",
    "$$\\hat{y}_{t+1} = f(d(y_t), s_t, z)$$\n",
    "\n",
    "因此，我们的解码器现在看起来像这样：\n",
    "\n",
    "![](assets/seq2seq6.png)\n",
    "\n",
    "请注意，初始隐藏状态 $s_0$ 仍然是上下文向量 $z$，因此在生成第一个token时，我们实际上将两个相同的上下文向量输入到 GRU 中。\n",
    "\n",
    "这两个更改如何减少信息压缩呢？嗯，假设解码器隐藏状态 $s_t$ 不再需要包含关于源序列的信息，因为它始终可用作输入。因此，它只需要包含有关到目前为止生成了哪些token的信息。将 $y_t$ 添加到线性层中还意味着该层可以直接看到token是什么，而不必从隐藏状态获取此信息。\n",
    "\n",
    "然而，这个假设只是一个假设，不可能确定模型实际上是如何使用提供给它的信息的（不要听任何说法与此不同的人）。尽管如此，这是一个很好的直觉，并且结果似乎表明这些修改是一个好主意！\n",
    "\n",
    "在实现中，我们将 $d(y_t)$ 和 $z$ 传递给 GRU，方法是将它们连接在一起，因此 GRU 的输入维度现在为 `emb_dim + hid_dim`（因为上下文向量的大小将为 `hid_dim`）。线性层将 $d(y_t), s_t$ 和 $z$ 通过将它们连接在一起传递，因此输入维度现在为 `emb_dim + hid_dim*2`。由于 GRU 只使用一个层，我们也不传递 dropout 的值给它。\n",
    "\n",
    "在 `forward` 中，现在需要一个 `context` 参数。在 `forward` 内部，我们将 $y_t$ 和 $z$ 连接在一起，形成 `emb_con`，然后将其馈送到 GRU 中，我们将 $d(y_t)$，$s_t$ 和 $z$ 连接在一起，形成 `output`，然后将其馈送到线性层，以接收我们的预测 $\\hat{y}_{t+1}$。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Cell):\n",
    "    def __init__(self, output_dim, embedding_dim, hidden_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.rnn = nn.GRU(embedding_dim + hidden_dim, hidden_dim)\n",
    "        self.fc_out = nn.Dense(embedding_dim + hidden_dim * 2, output_dim)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "    def construct(self, input:ms.Tensor, hidden:ms.Tensor, context:ms.Tensor):\n",
    "        # input = [batch size]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # context = [n layers * n directions, batch size, hidden dim]\n",
    "        # n layers and n directions in the decoder will both always be 1, therefore:\n",
    "        # hidden = [1, batch size, hidden dim]\n",
    "        # context = [1, batch size, hidden dim]\n",
    "        input = input.unsqueeze(0)\n",
    "        #input = [1, batch size]\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        #embedded = [1, batch size, embedding dim]\n",
    "        emb_con = ops.cat((embedded, context), axis = 2)\n",
    "        #emb_con = [1, batch size, embedding dim + hidden dim]\n",
    "        output, hidden = self.rnn(emb_con, hidden)\n",
    "        # output = [seq len, batch size, hidden dim * n directions]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # seq len, n layers and n directions will always be 1 in this decoder, therefore:\n",
    "        # output = [1, batch size, hidden dim]\n",
    "        # hidden = [1, batch size, hidden dim]\n",
    "        output = ops.cat((\n",
    "            embedded.squeeze(0), \n",
    "            hidden.squeeze(0), \n",
    "            context.squeeze(0)\n",
    "        ),\n",
    "            axis=1)\n",
    "        # output = [batch size, embedding dim + hidden dim * 2]\n",
    "        prediction = self.fc_out(output)\n",
    "        # prediction = [batch size, output dim]\n",
    "        return prediction, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Seq2Seq 模型\n",
    "\n",
    "将编码器和解码器放在一起，我们得到：\n",
    "\n",
    "![](assets/seq2seq7.png)\n",
    "\n",
    "同样，在此实现中，我们需要确保编码器和解码器中的隐藏维度相同。\n",
    "\n",
    "简要地概述所有步骤：\n",
    "- 创建 `outputs` 张量以保存所有预测 $\\hat{Y}$\n",
    "- 将源序列 $X$ 输入编码器以接收 `context` 向量\n",
    "- 将初始解码器隐藏状态设置为 `context` 向量，$s_0 = z = h_T$\n",
    "- 我们使用一批 `<sos>` token作为第一个 `input`，$y_1$\n",
    "- 然后我们在循环内解码：\n",
    "  - 将输入token $y_t$，先前的隐藏状态 $s_{t-1}$ 和上下文向量 $z$ 插入解码器\n",
    "  - 接收预测 $\\hat{y}_{t+1}$ 和新的隐藏状态 $s_t$\n",
    "  - 然后，我们决定是否要使用 teacher force，根据情况设置下一个输入（目标序列中的下一个真实token或最高预测的下一个token）\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Cell):\n",
    "    def __init__(self, encoder:Encoder, decoder:Decoder):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        assert encoder.hidden_dim == decoder.hidden_dim, \\\n",
    "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
    "        \n",
    "    def construct(self, src:ms.Tensor, trg:ms.Tensor, teacher_forcing_ratio):\n",
    "        # src = [src length, batch size]\n",
    "        # trg = [trg length, batch size]\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        # e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
    "        batch_size = trg.shape[1]\n",
    "        trg_length = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        # tensor to store decoder outputs\n",
    "        # outputs = torch.zeros(trg_length, batch_size, trg_vocab_size)\n",
    "        outputs = []\n",
    "        # last hidden state of the encoder is the context\n",
    "        context = self.encoder(src)\n",
    "        # context = [n layers * n directions, batch size, hidden dim]\n",
    "        # context also used as the initial hidden state of the decoder\n",
    "        hidden = context\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        input = trg[0,:]\n",
    "        for t in range(1, trg_length):\n",
    "            # insert input token embedding, previous hidden state and the context state\n",
    "            # receive output tensor (predictions) and new hidden state\n",
    "            output, hidden = self.decoder(input, hidden, context)\n",
    "            if len(outputs) == 0:\n",
    "                outputs.append(ops.zeros(output.shape, dtype=output.dtype))\n",
    "            # output = [batch size, output dim]\n",
    "            # hidden = [1, batch size, hidden dim]\n",
    "            # place predictions in a tensor holding predictions for each token\n",
    "            # outputs[t] = output\n",
    "            outputs.append(output)\n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1) \n",
    "            if self.training:\n",
    "                teacher_force = random.random() < teacher_forcing_ratio\n",
    "                # if teacher forcing, use actual next token as next input\n",
    "                # if not, use predicted token\n",
    "                input = trg[t] if teacher_force else top1\n",
    "                # input = [batch size]\n",
    "            else:\n",
    "                input = top1\n",
    "        return ops.stack(outputs, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练 Seq2Seq 模型\n",
    "\n",
    "这个教程的其余部分与前一个非常相似。\n",
    "\n",
    "我们初始化编码器、解码器和 Seq2Seq 模型。与之前一样，嵌入维度和使用的 dropout 量可以在编码器和解码器之间不同，但隐藏维度必须保持相同。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dim = len(de_vocab)\n",
    "output_dim = len(en_vocab)\n",
    "encoder_embedding_dim = 256\n",
    "decoder_embedding_dim = 256\n",
    "hidden_dim = 512\n",
    "encoder_dropout = 0.5\n",
    "decoder_dropout = 0.5\n",
    "\n",
    "encoder = Encoder(\n",
    "    input_dim,\n",
    "    encoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    encoder_dropout,\n",
    ")\n",
    "\n",
    "decoder = Decoder(\n",
    "    output_dim,\n",
    "    decoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    decoder_dropout,\n",
    ")\n",
    "\n",
    "\n",
    "model = Seq2Seq(encoder, decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们初始化参数。该论文指出，参数应该从均值为 0、标准差为 0.01 的正态分布中初始化，即 $\\mathcal{N}(0, 0.01)$。\n",
    "\n",
    "它还指出我们应该将循环参数初始化为特殊的初始化值，但为了保持简单，我们也将它们初始化为 $\\mathcal{N}(0, 0.01)$。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Seq2Seq<\n",
       "  (encoder): Encoder<\n",
       "    (embedding): Embedding<vocab_size=7853, embedding_size=256, use_one_hot=False, embedding_table=Parameter (name=encoder.embedding.embedding_table, shape=(7853, 256), dtype=Float32, requires_grad=True), dtype=Float32, padding_idx=None>\n",
       "    (rnn): GRU<\n",
       "      (rnn): _DynamicGRUCPUGPU<>\n",
       "      (dropout_op): Dropout<p=0.0>\n",
       "      >\n",
       "    (dropout): Dropout<p=0.5>\n",
       "    >\n",
       "  (decoder): Decoder<\n",
       "    (embedding): Embedding<vocab_size=5893, embedding_size=256, use_one_hot=False, embedding_table=Parameter (name=decoder.embedding.embedding_table, shape=(5893, 256), dtype=Float32, requires_grad=True), dtype=Float32, padding_idx=None>\n",
       "    (rnn): GRU<\n",
       "      (rnn): _DynamicGRUCPUGPU<>\n",
       "      (dropout_op): Dropout<p=0.0>\n",
       "      >\n",
       "    (fc_out): Dense<input_channels=1280, output_channels=5893, has_bias=True>\n",
       "    (dropout): Dropout<p=0.5>\n",
       "    >\n",
       "  >"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def init_weights(m:nn.Cell):\n",
    "    for name, param in m.parameters_and_names():\n",
    "        param.set_data(ops.uniform(param.shape, ms.Tensor([-0.08]), ms.Tensor([0.08]), dtype=param.dtype))\n",
    "\n",
    "model.apply(init_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们打印出参数的数量。\n",
    "\n",
    "尽管我们的编码器和解码器只有一个单层 RNN，但实际上我们拥有比上一个模型**更多**的参数。这是由于 GRU 和线性层的输入大小增加所致。然而，这不是一个显著的参数量，而且会导致训练时间的极小增加（每个时代额外增加约 3 秒）。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 14,219,781 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model:nn.Cell):\n",
    "    return sum(p.numel() for p in model.get_parameters() if p.requires_grad)\n",
    "\n",
    "print(f\"The model has {count_parameters(model):,} trainable parameters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们初始化优化器。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = nn.Adam(model.trainable_params())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们还初始化损失函数，确保忽略 `<pad>` token上的损失。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，我们创建训练循环...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_fn(src, trg, teacher_forcing_ratio):\n",
    "    # src = [src length, batch size]\n",
    "    # trg = [trg length, batch size]\n",
    "    output = model(src, trg, teacher_forcing_ratio)\n",
    "    # output = [trg length, batch size, trg vocab size]\n",
    "    output_dim = output.shape[-1]\n",
    "    output = output[1:].view(-1, output_dim)\n",
    "    # output = [(trg length - 1) * batch size, trg vocab size]\n",
    "    trg = trg[1:].view(-1)\n",
    "    # trg = [(trg length - 1) * batch size]\n",
    "    loss = criterion(output, trg.astype(ms.int32))\n",
    "    return loss\n",
    "\n",
    "grad_fn = ms.value_and_grad(forward_fn, grad_position=None, weights=model.trainable_params())\n",
    "\n",
    "\n",
    "def train_fn(data_loader, optimizer, clip, teacher_forcing_ratio):\n",
    "    epoch_loss = 0.\n",
    "    model.set_train(True)\n",
    "    for i, batch in enumerate(tqdm(data_loader)):\n",
    "        src = batch[\"de_ids\"]\n",
    "        trg = batch[\"en_ids\"]\n",
    "        loss, grads = grad_fn(src, trg, teacher_forcing_ratio)\n",
    "        # grads = ops.clip_by_norm(grads, max_norm=clip)\n",
    "        optimizer(grads)\n",
    "        epoch_loss += float(loss)\n",
    "\n",
    "    return epoch_loss / len(data_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...和评估循环，记得关闭teaching forcing。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_fn(model:Seq2Seq, data_loader, criterion):\n",
    "    epoch_loss = 0\n",
    "    model.set_train(False)\n",
    "    for i, batch in enumerate(data_loader):\n",
    "        src = batch[\"de_ids\"]\n",
    "        trg = batch[\"en_ids\"]\n",
    "        # src = [src length, batch size]\n",
    "        # trg = [trg length, batch size]\n",
    "        output = model(src, trg, 0) # turn off teacher forcing\n",
    "        # output = [trg length, batch size, trg vocab size]\n",
    "        output_dim = output.shape[-1]\n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        # output = [(trg length - 1) * batch size, trg vocab size]\n",
    "        trg = trg[1:].view(-1)\n",
    "        # trg = [(trg length - 1) * batch size]\n",
    "        loss = criterion(output, trg.astype(ms.int32))\n",
    "        epoch_loss += float(loss)\n",
    "    return epoch_loss / len(data_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，我们训练我们的模型，保存给出最佳(best)验证损失的参数。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1fb921396de543c98da43d3e1af9ca4a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c76d4762b07144709ebc4b6de1284f7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   4.760 | Train PPL: 116.738\n",
      "\tValid Loss:   5.648 | Valid PPL: 283.596\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d65cf1694524a978999a5577e0172ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   4.263 | Train PPL:  71.009\n",
      "\tValid Loss:   5.302 | Valid PPL: 200.643\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "414e0b9414874b74bfde88618cc39e76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   3.728 | Train PPL:  41.596\n",
      "\tValid Loss:   4.764 | Valid PPL: 117.161\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32c3285dee1544879c3a60e5cb2ef44d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   3.285 | Train PPL:  26.705\n",
      "\tValid Loss:   4.515 | Valid PPL:  91.397\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2307da5d50124a4c9b36a762d1b6b641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   2.943 | Train PPL:  18.967\n",
      "\tValid Loss:   4.338 | Valid PPL:  76.528\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2e23c86993044c0ab0f4d1a3a52134e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   2.699 | Train PPL:  14.859\n",
      "\tValid Loss:   4.196 | Valid PPL:  66.407\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af4c5af14d8e4b648832ee0d961bb1ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   2.462 | Train PPL:  11.725\n",
      "\tValid Loss:   4.266 | Valid PPL:  71.207\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e17e74a80ded495dbfe09497178acf6f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   2.269 | Train PPL:   9.666\n",
      "\tValid Loss:   4.308 | Valid PPL:  74.269\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d8e25a6d8cf4816854bc40371995c24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   2.099 | Train PPL:   8.159\n",
      "\tValid Loss:   4.253 | Valid PPL:  70.292\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94b1820e371d4a2898318dba4d5e29ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/226 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTrain Loss:   1.944 | Train PPL:   6.987\n",
      "\tValid Loss:   4.274 | Valid PPL:  71.792\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 10\n",
    "clip = 1.0\n",
    "teacher_forcing_ratio = 0.5\n",
    "\n",
    "best_valid_loss = float(\"inf\")\n",
    "\n",
    "for epoch in tqdm(range(n_epochs)):\n",
    "        \n",
    "    train_loss = train_fn(\n",
    "        train_data_loader, \n",
    "        optimizer, \n",
    "        clip, \n",
    "        teacher_forcing_ratio, \n",
    "    )\n",
    "    \n",
    "    valid_loss = evaluate_fn(\n",
    "        model, \n",
    "        valid_data_loader, \n",
    "        criterion, \n",
    "    )\n",
    "\n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        ms.save_checkpoint(model, f\"./tut2-model.ckpt\")\n",
    "    \n",
    "    print(f\"\\tTrain Loss: {train_loss:7.3f} | Train PPL: {np.exp(train_loss):7.3f}\")\n",
    "    print(f\"\\tValid Loss: {valid_loss:7.3f} | Valid PPL: {np.exp(valid_loss):7.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，我们使用这些“best”参数在测试集上测试模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Test Loss: 4.142 | Test PPL:  62.906 |\n"
     ]
    }
   ],
   "source": [
    "ms.load_checkpoint(f\"tut2-model.ckpt\", model)\n",
    "\n",
    "test_loss = evaluate_fn(model, test_data_loader, criterion)\n",
    "\n",
    "print(f\"| Test Loss: {test_loss:.3f} | Test PPL: {np.exp(test_loss):7.3f} |\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "仅看测试损失，我们获得了比先前模型更好的性能。这是这种模型架构做得不错的一个很好的迹象！缓解信息压缩似乎是前进的道路，在下一个教程中，我们将进一步扩展这一点，引入“注意力”机制。\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}