Task03_2_机器翻译.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 10.12 机器翻译"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:04.802600Z",
     "start_time": "2020-02-16T15:23:04.187123Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.1.0 cuda\n"
     ]
    }
   ],
   "source": [
    "import collections\n",
    "import os\n",
    "import io\n",
    "import math\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as Data\n",
    "import seq\n",
    "import sys\n",
    "# sys.path.append(\"..\") \n",
    "import utils as d2l\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "print(torch.__version__, device)\n",
    "\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = \"all\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10.12.1 读取和预处理数据"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接着定义两个辅助函数对后面重新的数据进行预先。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:19:55.168827Z",
     "start_time": "2020-02-16T15:19:55.153673Z"
    },
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "outputs": [],
   "source": [
    "import torchtext\n",
    "PAD, BOS, EOS = '<pad>', '<bos>', '<eos>'\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "# 将一个序列中所有的词记录在all_tokens中以便之后构造词典，然后在该序列后面添加PAD直到序列\n",
    "# 长度变为max_seq_len，然后将序列保存在all_seqs中\n",
    "def process_one_seq(seq_tokens, all_tokens, all_seqs, max_seq_len):\n",
    "    all_tokens.extend(seq_tokens)\n",
    "    seq_tokens += [EOS] + [PAD] * (max_seq_len - len(seq_tokens) - 1)\n",
    "    all_seqs.append(seq_tokens)\n",
    "\n",
    "# 使用所有的词来构造词典。并将所有序列中的词变换为词索引后构造Tensor\n",
    "def build_dic(all_tokens, all_seqs):\n",
    "    vocab = torchtext.vocab.Vocab(collections.Counter(all_tokens), specials=[PAD, BOS, EOS])\n",
    "    indices = [[vocab.stoi[w] for w in seq] for seq in all_seqs]\n",
    "    return vocab, torch.tensor(indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了演示方便，我们在这里使用一个很小的法语—英语数据集。在这个数据集里，每一行是一对法语句子和它对应的英语句子，中间使用'\\t'替换。在读取数据时，我们在句末附上“ <eos>”符号，并可能通过添加“ <pad>”符号使每个序列的长度扩展max_seq_len。我们为法语单词和英语单词分别创建字典。法语单词的索引和英语单词的索引相互独立。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:19:55.277824Z",
     "start_time": "2020-02-16T15:19:55.173048Z"
    },
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "outputs": [],
   "source": [
    "def read_data(max_seq_len):\n",
    "    # in和out分别是input和output的缩写\n",
    "    in_tokens, out_tokens, in_seqs, out_seqs = [], [], [], []\n",
    "    with io.open('/home/cc/holdshy/XJQ/Pytorch/Dive_into_DL/eng-fra.txt') as f:\n",
    "        lines = f.readlines()\n",
    "    for line in lines:\n",
    "        in_seq, out_seq = line.rstrip().split('\\t')\n",
    "        in_seq_tokens, out_seq_tokens = in_seq.split(' '), out_seq.split(' ')\n",
    "        if max(len(in_seq_tokens), len(out_seq_tokens)) > max_seq_len - 1:\n",
    "            continue  # 如果加上EOS后长于max_seq_len，则忽略掉此样本\n",
    "        process_one_seq(in_seq_tokens, in_tokens, in_seqs, max_seq_len)\n",
    "        process_one_seq(out_seq_tokens, out_tokens, out_seqs, max_seq_len)\n",
    "    in_vocab, in_data = build_dic(in_tokens, in_seqs)\n",
    "    out_vocab, out_data = build_dic(out_tokens, out_seqs)\n",
    "    return in_vocab, out_vocab, Data.TensorDataset(in_data, out_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将序列的最大长度设置成7，然后查看引用到的第一个样本。该样本分别包含法语单词索引序列和英语单词索引序列。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:19:56.449369Z",
     "start_time": "2020-02-16T15:19:55.282181Z"
    },
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "outputs": [],
   "source": [
    "max_seq_len = 7\n",
    "in_vocab, out_vocab, dataset = read_data(max_seq_len)\n",
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10.12.2 含注意力机制的编码器—解码器\n",
    "### 10.12.2.1 编码器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Image Name](https://cdn.kesci.com/upload/image/q5jcat3c8m.png?imageView2/0/w/640/h/640)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在编码器中，我们将输入语言的词索引通过词嵌入层得到词的表征，然后输入到一个多层门控循环单元中。改为我们在6.5节（循环神经网络的简洁实现）中提到的，PyTorch的nn.GRU实例在前向计算后也会分别返回输出和最终时间步的多层隐藏状态。其中的输出指的是最后一层的隐藏层在各个时间步的隐藏状态，并不涉及输出层计算。关注机制将这些输出作为键项和值项。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sequence to Sequence模型\n",
    "### 模型：\n",
    "训练  \n",
    "![Image Name](https://cdn.kesci.com/upload/image/q5jc7a53pt.png?imageView2/0/w/640/h/640)\n",
    "预测\n",
    "![Image Name](https://cdn.kesci.com/upload/image/q5jcecxcba.png?imageView2/0/w/640/h/640)\n",
    "### 具体结构：\n",
    "![Image Name](https://cdn.kesci.com/upload/image/q5jccjhkii.png?imageView2/0/w/500/h/500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:10.926316Z",
     "start_time": "2020-02-16T15:23:10.908091Z"
    }
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(Encoder, self).__init__(**kwargs)\n",
    "    def forward(self, X, *args):\n",
    "        raise NotImplementedError\n",
    "        \n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(Decoder, self).__init__(**kwargs)\n",
    "    def init_state(self, enc_outputs, *args):\n",
    "        raise NotImplementedError\n",
    "    def forward(self, X, state):\n",
    "        raise NotImplementedError\n",
    "        \n",
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, **kwargs):\n",
    "        super(EncoderDecoder, self).__init__(**kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "    def forward(self, enc_X, dec_X, *args):\n",
    "        enc_outputs = self.encoder(enc_X, *args)\n",
    "        dec_state = self.decoder.init_state(enc_outputs, *args)\n",
    "        return self.decoder(dec_X, dec_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:11.563516Z",
     "start_time": "2020-02-16T15:23:11.558711Z"
    }
   },
   "outputs": [],
   "source": [
    "# 下面我们来创建一个批量大小为4，时间步数为7的小批量序列输入。设门控循环单元的隐藏层个数为2，隐藏单元个数为16。\n",
    "# 编码器该输入执行前向计算后返回的输出形状为（时间步数，批量大小，隐藏单元个数）。\n",
    "# 门控循环单元在最终时间步的多层隐藏状态的形状为（隐藏层个数，批量，隐藏单元个数）。\n",
    "# 对于门控循环单元来说，state就是一个元素，即隐藏状态；如果使用长短期记忆，state是一个元组，包含两个元素即隐藏状态和记忆细胞。\n",
    "\n",
    "# encoder = Encoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)\n",
    "# output, state = encoder(torch.zeros((4, 7)), encoder.begin_state())\n",
    "# output.shape, state.shape # GRU的state是h, 而LSTM的是一个元组(h, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:14.616088Z",
     "start_time": "2020-02-16T15:23:14.599022Z"
    }
   },
   "outputs": [],
   "source": [
    "class Seq2SeqEncoder(Encoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout=0, **kwargs):\n",
    "        super(Seq2SeqEncoder, self).__init__(**kwargs)\n",
    "        self.num_hiddens=num_hiddens\n",
    "        self.num_layers=num_layers\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = nn.LSTM(embed_size,num_hiddens, num_layers, dropout=dropout)\n",
    "\n",
    "    def begin_state(self, batch_size, device):\n",
    "        return [torch.zeros(size=(self.num_layers, batch_size, self.num_hiddens),  device=device),\n",
    "                torch.zeros(size=(self.num_layers, batch_size, self.num_hiddens),  device=device)]\n",
    "    def forward(self, X, *args):\n",
    "        # 输入形状是(批量大小, 时间步数)。将输出互换样本维和时间步维\n",
    "        # batch_size:一个batch几句话, seq_len:一句话几个单词, input_size:embedding后变成input_size维度的词向量\n",
    "        X = self.embedding(X) # X shape: (batch_size, seq_len, embed_size)\n",
    "        # GRU、LSTM:时序为 第 0 个输入\n",
    "        X = X.transpose(0, 1)  # RNN needs first axes to be time\n",
    "        # state = self.begin_state(X.shape[1], device=X.device)\n",
    "        # out:记忆细胞的状态h1、h2...，state:隐层的状态ht\n",
    "        out, state = self.rnn(X)\n",
    "        # The shape of out is (seq_len, batch_size, num_hiddens).\n",
    "        # state contains the hidden state and the memory cell\n",
    "        # of the last time step, the shape is (num_layers, batch_size, num_hiddens)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:15.330836Z",
     "start_time": "2020-02-16T15:23:15.284874Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([7, 4, 16]), 2, torch.Size([2, 4, 16]), torch.Size([2, 4, 16]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8,num_hiddens=16, num_layers=2)\n",
    "X = torch.zeros((4, 7),dtype=torch.long)  # 输入4句话，每句话7个单词\n",
    "output, state = encoder(X)\n",
    "output.shape, len(state), state[0].shape, state[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:16.740802Z",
     "start_time": "2020-02-16T15:23:16.725334Z"
    }
   },
   "outputs": [],
   "source": [
    "class Seq2SeqDecoder(Decoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,\n",
    "                 dropout=0, **kwargs):\n",
    "        super(Seq2SeqDecoder, self).__init__(**kwargs)\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = nn.LSTM(embed_size,num_hiddens, num_layers, dropout=dropout)\n",
    "        self.dense = nn.Linear(num_hiddens,vocab_size)\n",
    "\n",
    "    def init_state(self, enc_outputs, *args):\n",
    "        return enc_outputs[1]\n",
    "\n",
    "    def forward(self, X, state):\n",
    "        X = self.embedding(X).transpose(0, 1)\n",
    "        out, state = self.rnn(X, state)\n",
    "        # Make the batch to be the first dimension to simplify loss computation.\n",
    "        out = self.dense(out).transpose(0, 1)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:17.314848Z",
     "start_time": "2020-02-16T15:23:17.290815Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 7, 10]), 2)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 4, 16]), torch.Size([2, 4, 16]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8,num_hiddens=16, num_layers=2)\n",
    "state = decoder.init_state(encoder(X))\n",
    "out, state = decoder(X, state)\n",
    "\n",
    "out.shape, len(state)\n",
    "# batch_size = 4\n",
    "# state: (num_layers, batch_size, num_hiddens)\n",
    "state[0].shape, state[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:18.555600Z",
     "start_time": "2020-02-16T15:23:18.546757Z"
    }
   },
   "outputs": [],
   "source": [
    "def SequenceMask(X, X_len,value=0):\n",
    "    maxlen = X.size(1)\n",
    "    mask = torch.arange(maxlen)[None, :].to(X_len.device) < X_len[:, None]   \n",
    "    X[~mask]=value\n",
    "    return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:19.199140Z",
     "start_time": "2020-02-16T15:23:19.187807Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 0, 0],\n",
       "        [4, 5, 0]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.tensor([[1,2,3], [4,5,6]])\n",
    "SequenceMask(X,torch.tensor([1,2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:19.992681Z",
     "start_time": "2020-02-16T15:23:19.979793Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.,  1.,  1.,  1.],\n",
       "         [-1., -1., -1., -1.],\n",
       "         [-1., -1., -1., -1.]],\n",
       "\n",
       "        [[ 1.,  1.,  1.,  1.],\n",
       "         [ 1.,  1.,  1.,  1.],\n",
       "         [-1., -1., -1., -1.]]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.ones((2,3, 4))\n",
    "SequenceMask(X, torch.tensor([1,2]),value=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:20.818533Z",
     "start_time": "2020-02-16T15:23:20.807834Z"
    }
   },
   "outputs": [],
   "source": [
    "class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):\n",
    "    # pred shape: (batch_size, seq_len, vocab_size)\n",
    "    # label shape: (batch_size, seq_len)\n",
    "    # valid_length shape: (batch_size, )\n",
    "    def forward(self, pred, label, valid_length):\n",
    "        # the sample weights shape should be (batch_size, seq_len)\n",
    "        weights = torch.ones_like(label)\n",
    "        weights = SequenceMask(weights, valid_length).float()\n",
    "        self.reduction='none'\n",
    "        output=super(MaskedSoftmaxCELoss, self).forward(pred.transpose(1,2), label)\n",
    "        return (output*weights).mean(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:23:21.442134Z",
     "start_time": "2020-02-16T15:23:21.428180Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([2.3026, 1.7269, 0.0000])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = MaskedSoftmaxCELoss()\n",
    "loss(torch.ones((3, 4, 10)), torch.ones((3,4),dtype=torch.long), torch.tensor([4,3,0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练、测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:24:52.960114Z",
     "start_time": "2020-02-16T15:24:52.916065Z"
    },
    "code_folding": [
     7
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import requests\n",
    "from io import BytesIO\n",
    "from torch.utils import data\n",
    "import sys\n",
    "import collections\n",
    "\n",
    "class Vocab(object): # This class is saved in d2l.\n",
    "    def __init__(self, tokens, min_freq=0, use_special_tokens=False):\n",
    "        # sort by frequency and token\n",
    "        counter = collections.Counter(tokens)\n",
    "        token_freqs = sorted(counter.items(), key=lambda x: x[0])\n",
    "        token_freqs.sort(key=lambda x: x[1], reverse=True)\n",
    "        if use_special_tokens:\n",
    "            # padding, begin of sentence, end of sentence, unknown\n",
    "            self.pad, self.bos, self.eos, self.unk = (0, 1, 2, 3)\n",
    "            tokens = ['', '', '', '']\n",
    "        else:\n",
    "            self.unk = 0\n",
    "            tokens = ['']\n",
    "        tokens += [token for token, freq in token_freqs if freq >= min_freq]\n",
    "        self.idx_to_token = []\n",
    "        self.token_to_idx = dict()\n",
    "        for token in tokens:\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, tokens):\n",
    "        if not isinstance(tokens, (list, tuple)):\n",
    "            return self.token_to_idx.get(tokens, self.unk)\n",
    "        else:\n",
    "            return [self.__getitem__(token) for token in tokens]\n",
    "    \n",
    "    def to_tokens(self, indices):\n",
    "        if not isinstance(indices, (list, tuple)):\n",
    "            return self.idx_to_token[indices]\n",
    "        else:\n",
    "            return [self.idx_to_token[index] for index in indices]\n",
    "\n",
    "        \n",
    "def load_data_nmt(batch_size, max_len, num_examples=1000):\n",
    "    \"\"\"Download an NMT dataset, return its vocabulary and data iterator.\"\"\"\n",
    "    # Download and preprocess\n",
    "    def preprocess_raw(text):\n",
    "        text = text.replace('\\u202f', ' ').replace('\\xa0', ' ')\n",
    "        out = ''\n",
    "        for i, char in enumerate(text.lower()):\n",
    "            if char in (',', '!', '.') and text[i-1] != ' ':\n",
    "                out += ' '\n",
    "            out += char\n",
    "        return out \n",
    "    \n",
    "    with open('/home/cc/holdshy/XJQ/Pytorch/Dive_into_DL/eng-fra.txt', 'r') as f:\n",
    "        raw_text = f.read()\n",
    "    text = preprocess_raw(raw_text)\n",
    "    \n",
    "    # Tokenize\n",
    "    source, target = [], []\n",
    "    for i, line in enumerate(text.split('\\n')):\n",
    "        if i >= num_examples:\n",
    "            break\n",
    "        parts = line.split('\\t')\n",
    "        if len(parts) >= 2:\n",
    "            source.append(parts[0].split(' '))\n",
    "            target.append(parts[1].split(' '))\n",
    "\n",
    "    # Build vocab\n",
    "    def build_vocab(tokens):\n",
    "        tokens = [token for line in tokens for token in line]\n",
    "        return Vocab(tokens, min_freq=3, use_special_tokens=True)\n",
    "    \n",
    "    src_vocab, tgt_vocab = build_vocab(source), build_vocab(target)\n",
    "\n",
    "    # Convert to index arrays\n",
    "    def pad(line, max_len, padding_token):\n",
    "        if len(line) > max_len:\n",
    "            return line[:max_len]\n",
    "        return line + [padding_token] * (max_len - len(line))\n",
    "\n",
    "    def build_array(lines, vocab, max_len, is_source):\n",
    "        lines = [vocab[line] for line in lines]\n",
    "        if not is_source:\n",
    "            lines = [[vocab.bos] + line + [vocab.eos] for line in lines]\n",
    "        array = torch.tensor([pad(line, max_len, vocab.pad) for line in lines])\n",
    "        valid_len = (array != vocab.pad).sum(1)\n",
    "        return array, valid_len\n",
    "\n",
    "    src_vocab, tgt_vocab = build_vocab(source), build_vocab(target)\n",
    "    src_array, src_valid_len = build_array(source, src_vocab, max_len, True)\n",
    "    tgt_array, tgt_valid_len = build_array(target, tgt_vocab, max_len, False)\n",
    "    train_data = data.TensorDataset(src_array, src_valid_len, tgt_array, tgt_valid_len)\n",
    "    train_iter = data.DataLoader(train_data, batch_size, shuffle=True)\n",
    "    \n",
    "    return src_vocab, tgt_vocab, train_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:24:55.387494Z",
     "start_time": "2020-02-16T15:24:53.642303Z"
    }
   },
   "outputs": [],
   "source": [
    "# from seq import load_data_nmt\n",
    "embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.0\n",
    "batch_size, max_len, num_examples = 64, 10, 1000\n",
    "lr, num_epochs, ctx = 0.005, 300, seq.try_gpu()\n",
    "\n",
    "src_vocab, tgt_vocab, train_iter = load_data_nmt(batch_size, max_len)\n",
    "encoder = Seq2SeqEncoder(len(src_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "decoder = Seq2SeqDecoder(len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "model = EncoderDecoder(encoder, decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:28:43.362827Z",
     "start_time": "2020-02-16T15:28:43.343781Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from torch import optim\n",
    "def train_ch7(model, data_iter, lr, num_epochs, device): \n",
    "    \"\"\"Train an encoder-decoder model\"\"\"\n",
    "    model.to(device)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "    loss = MaskedSoftmaxCELoss()\n",
    "    tic = time.time()\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "        l_sum, num_tokens_sum = 0.0, 0.0\n",
    "        for batch in data_iter:\n",
    "            optimizer.zero_grad()\n",
    "            X, X_vlen, Y, Y_vlen = [x.to(device) for x in batch]\n",
    "            Y_input, Y_label, Y_vlen = Y[:,:-1], Y[:,1:], Y_vlen-1\n",
    "            Y_hat, _ = model(X, Y_input, X_vlen, Y_vlen)\n",
    "            l = loss(Y_hat, Y_label, Y_vlen).sum()\n",
    "            l.backward()\n",
    "            \n",
    "            with torch.no_grad():\n",
    "                seq.grad_clipping_nn(model, 5, device)\n",
    "                \n",
    "            num_tokens = Y_vlen.sum().item()\n",
    "            optimizer.step()\n",
    "            l_sum += l.sum().item()\n",
    "            num_tokens_sum += num_tokens\n",
    "        if epoch % 50 == 0:\n",
    "            print(\"epoch {0:4d},loss {1:.3f}, time {2:.1f} sec\".format( \n",
    "                  epoch, (l_sum/num_tokens_sum), time.time()-tic))\n",
    "            tic = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:30:15.937534Z",
     "start_time": "2020-02-16T15:28:44.707396Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch   50,loss 0.098, time 15.2 sec\n",
      "epoch  100,loss 0.048, time 15.2 sec\n",
      "epoch  150,loss 0.033, time 15.2 sec\n",
      "epoch  200,loss 0.027, time 15.1 sec\n",
      "epoch  250,loss 0.024, time 15.1 sec\n",
      "epoch  300,loss 0.022, time 15.3 sec\n"
     ]
    }
   ],
   "source": [
    "train_ch7(model, train_iter, lr, num_epochs, ctx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:31:16.725075Z",
     "start_time": "2020-02-16T15:31:16.706511Z"
    }
   },
   "outputs": [],
   "source": [
    "def translate_ch7(model, src_sentence, src_vocab, tgt_vocab, max_len, device):\n",
    "    src_tokens = src_vocab[src_sentence.lower().split(' ')]\n",
    "    src_len = len(src_tokens)\n",
    "    if src_len < max_len:\n",
    "        src_tokens += [src_vocab.pad] * (max_len - src_len)\n",
    "    enc_X = torch.tensor(src_tokens, device=device)\n",
    "    enc_valid_length = torch.tensor([src_len], device=device)\n",
    "    # use expand_dim to add the batch_size dimension.\n",
    "    enc_outputs = model.encoder(enc_X.unsqueeze(dim=0), enc_valid_length)\n",
    "    dec_state = model.decoder.init_state(enc_outputs, enc_valid_length)\n",
    "    dec_X = torch.tensor([tgt_vocab.bos], device=device).unsqueeze(dim=0)\n",
    "    predict_tokens = []\n",
    "    for _ in range(max_len):\n",
    "        Y, dec_state = model.decoder(dec_X, dec_state)\n",
    "        # The token with highest score is used as the next time step input.\n",
    "        dec_X = Y.argmax(dim=2)\n",
    "        py = dec_X.squeeze(dim=0).int().item()\n",
    "        if py == tgt_vocab.eos:\n",
    "            break\n",
    "        predict_tokens.append(py)\n",
    "    return ' '.join(tgt_vocab.to_tokens(predict_tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-16T15:31:17.581819Z",
     "start_time": "2020-02-16T15:31:17.546092Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Go . => va !\n",
      "Wow ! =>  !\n",
      "I'm OK . => je vais bien .\n",
      "I won ! => j'ai gagné !\n"
     ]
    }
   ],
   "source": [
    "for sentence in ['Go .', 'Wow !', \"I'm OK .\", 'I won !']:\n",
    "    print(sentence + ' => ' + translate_ch7(\n",
    "        model, sentence, src_vocab, tgt_vocab, max_len, ctx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Go . => va !\n",
    "# Good Night ! =>   !\n",
    "# I'm OK . => je vais bien .\n",
    "# I won ! => je l'ai emporté !"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Beam Search\n",
    "简单greedy search：\n",
    "\n",
    "![Image Name](https://cdn.kesci.com/upload/image/q5jchqoppn.png?imageView2/0/w/440/h/440)\n",
    "\n",
    "维特比算法：选择整体分数最高的句子（搜索空间太大）\n",
    "集束搜索：\n",
    "\n",
    "![Image Name](https://cdn.kesci.com/upload/image/q5jcia86z1.png?imageView2/0/w/640/h/640)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.6.9"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "592.5px",
    "left": "115px",
    "top": "253.13px",
    "width": "184.154px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}