seq2seq+attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# pytorch实现seq2seq+attention转换日期\n",
    "***\n",
    "使用keras实现加入注意力机制的seq2seq比较麻烦，所以这里我尝试使用机器翻译的seq2seq+attention模型实现人造日期对标准日期格式的转换。\n",
    "\n",
    "\n",
    "所copy的代码来自[practical-pytorch教程](https://github.com/spro/practical-pytorch/tree/master/seq2seq-translation)，以及[pytorch-seq2seq教程](https://github.com/bentrevett/pytorch-seq2seq)\n",
    "\n",
    "\n",
    "所用的数据来自[注意力机制keras实现](https://github.com/Choco31415/Attention_Network_With_Keras/tree/master/data)。   \n",
    "python3   \n",
    "pytorch版本 0.4.0  \n",
    "可能需要GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "from matplotlib import ticker\n",
    "from numpy import *\n",
    "from collections import Counter\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 预处理 \n",
    "---\n",
    "这里先生成字符和数字相互转换的字典，如果是句子也可以按照词为单位。我在字典的开头添加了4种表示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def build_vocab(texts, n=None):\n",
    "    counter = Counter(''.join(texts))  # char level\n",
    "    char2index = {w: i for i, (w, c) in enumerate(counter.most_common(n), start=4)}\n",
    "    char2index['~'] = 0  # pad  不足长度的文本在后边填充0\n",
    "    char2index['^'] = 1  # sos  表示句子的开头\n",
    "    char2index['$'] = 2  # eos  表示句子的结尾\n",
    "    char2index['#'] = 3  # unk  表示句子中出现的字典中没有的未知词\n",
    "    index2char = {i: w for w, i in char2index.items()}\n",
    "    return char2index, index2char"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "先看一下数据的格式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['six hours and fifty five am', '06:55']]\n"
     ]
    }
   ],
   "source": [
    "pairs = json.load(open('Time Dataset.json', 'rt', encoding='utf-8'))\n",
    "print(pairs[:1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们将目标文本和原文本分开，求出两边句子的最大长度，然后建立两边各自的字典。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data = array(pairs)\n",
    "src_texts = data[:, 0]\n",
    "trg_texts = data[:, 1]\n",
    "src_c2ix, src_ix2c = build_vocab(src_texts)\n",
    "trg_c2ix, trg_ix2c = build_vocab(trg_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里按批量跟新，定义一个随机批量生成的函数，它能够将文本转换成字典中的数字表示，并同时返回batch_size个样本和它们的长度，这些样本按照长度降序排序。pad的长度以batch中最长的为准。这主要是为了适应pack_padded_sequence这个函数，因为输入RNN的序列不需要将pad标志也输入RNN中计算，RNN只需要循环计算到其真实长度即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(43, 7)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def indexes_from_text(text, char2index):\n",
    "    return [1] + [char2index[c] for c in text] + [2]  # 手动添加开始结束标志\n",
    "def pad_seq(seq, max_length):\n",
    "    seq += [0 for _ in range(max_length - len(seq))]\n",
    "    return seq\n",
    "\n",
    "max_src_len = max(list(map(len, src_texts))) + 2\n",
    "max_trg_len = max(list(map(len, trg_texts))) + 2\n",
    "max_src_len, max_trg_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def random_batch(batch_size, pairs, src_c2ix, trg_c2ix):\n",
    "    input_seqs, target_seqs = [], []\n",
    "\n",
    "    for i in random.choice(len(pairs), batch_size):\n",
    "        input_seqs.append(indexes_from_text(pairs[i][0], src_c2ix))\n",
    "        target_seqs.append(indexes_from_text(pairs[i][1], trg_c2ix))\n",
    "\n",
    "    seq_pairs = sorted(zip(input_seqs, target_seqs), key=lambda p: len(p[0]), reverse=True)\n",
    "    input_seqs, target_seqs = zip(*seq_pairs)\n",
    "    input_lengths = [len(s) for s in input_seqs]\n",
    "    input_padded = [pad_seq(s, max(input_lengths)) for s in input_seqs]\n",
    "    target_lengths = [len(s) for s in target_seqs]\n",
    "    target_padded = [pad_seq(s, max(target_lengths)) for s in target_seqs]\n",
    "\n",
    "    input_var = torch.LongTensor(input_padded).transpose(0, 1)  \n",
    "    # seq_len x batch_size\n",
    "    target_var = torch.LongTensor(target_padded).transpose(0, 1)\n",
    "    input_var = input_var.to(device)\n",
    "    target_var = target_var.to(device)\n",
    "\n",
    "    return input_var, input_lengths, target_var, target_lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以先打印一下，batch_size=3时的返回结果。注意这里batch经过了转置。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[  1,   1,   1],\n",
       "         [  6,  18,   6],\n",
       "         [ 31,  18,  15],\n",
       "         [  5,   4,  25],\n",
       "         [ 30,  29,   9],\n",
       "         [ 28,   5,  26],\n",
       "         [  5,  12,  26],\n",
       "         [  4,   7,   2],\n",
       "         [  7,  10,   0],\n",
       "         [ 33,   5,   0],\n",
       "         [ 22,   4,   0],\n",
       "         [ 30,  23,   0],\n",
       "         [  7,   4,   0],\n",
       "         [ 22,  13,   0],\n",
       "         [ 34,   9,   0],\n",
       "         [  4,  11,   0],\n",
       "         [  6,   2,   0],\n",
       "         [ 16,   0,   0],\n",
       "         [ 14,   0,   0],\n",
       "         [ 10,   0,   0],\n",
       "         [  6,   0,   0],\n",
       "         [ 20,   0,   0],\n",
       "         [  4,   0,   0],\n",
       "         [  5,   0,   0],\n",
       "         [ 14,   0,   0],\n",
       "         [ 39,   0,   0],\n",
       "         [ 16,   0,   0],\n",
       "         [  6,   0,   0],\n",
       "         [  4,   0,   0],\n",
       "         [ 13,   0,   0],\n",
       "         [  6,   0,   0],\n",
       "         [  4,   0,   0],\n",
       "         [ 11,   0,   0],\n",
       "         [ 14,   0,   0],\n",
       "         [ 32,   0,   0],\n",
       "         [  8,   0,   0],\n",
       "         [ 14,   0,   0],\n",
       "         [ 39,   0,   0],\n",
       "         [ 16,   0,   0],\n",
       "         [  6,   0,   0],\n",
       "         [  2,   0,   0]], device='cuda:0'),\n",
       " [41, 17, 8],\n",
       " tensor([[  1,   1,   1],\n",
       "         [  5,   5,   6],\n",
       "         [  5,   7,  10],\n",
       "         [  4,   4,   4],\n",
       "         [  8,   8,   5],\n",
       "         [ 11,  11,   5],\n",
       "         [  2,   2,   2]], device='cuda:0'),\n",
       " [7, 7, 7])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_batch(3, data, src_c2ix, trg_c2ix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型定义\n",
    "---\n",
    "模型分为encoder和decoder两个部分，decoder部分比较简单，就是一层Embedding层加上两层GRU。之前处理的batch的格式主要是为了使用pack_padded_sequence和pad_packed_sequence这两个类对GRU输入输出批量处理。一定要注意各个变量的shape。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, embedding_dim, hidden_dim, num_layers=2, dropout=0.2):\n",
    "        super().__init__()\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.dropout = dropout\n",
    "        # input_dim = vocab_size + 1\n",
    "        self.embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "\n",
    "        self.rnn = nn.GRU(embedding_dim, hidden_dim,\n",
    "                          num_layers=num_layers, dropout=dropout)\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, input_seqs, input_lengths, hidden=None):\n",
    "        # src = [sent len, batch size]\n",
    "        embedded = self.dropout(self.embedding(input_seqs))\n",
    "\n",
    "        # embedded = [sent len, batch size, emb dim]\n",
    "        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)\n",
    "\n",
    "        outputs, hidden = self.rnn(packed, hidden)\n",
    "        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs)\n",
    "        # outputs, hidden = self.rnn(embedded, hidden)\n",
    "        # outputs = [sent len, batch size, hid dim * n directions]\n",
    "        # hidden = [n layers, batch size, hid dim]\n",
    "        # outputs are always from the last layer\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先定义一下Attention层，这里主要是对encoder的输出进行attention操作，也可以直接对embedding层的输出进行attention。   \n",
    "论文[Neural Machine Translation by Jointly Learning to Align and Translate](https://arxiv.org/abs/1409.0473)中定义了attention的计算公式。 \n",
    "\n",
    "decoder的输出取决于decoder先前的输出和 $\\mathbf x$, 这里 $\\mathbf x$ 包括当前GRU输出的hidden state（这部分已经考虑了先前的输出） 以及attention（上下文向量，由encoder的输出求得）。 计算公式如下：函数 $g$ 非线性激活的全连接层，输入是 $y_{i-1}$, $s_i$, and $c_i$ 三者的拼接。\n",
    "\n",
    "$$\n",
    "p(y_i \\mid \\{y_1,...,y_{i-1}\\},\\mathbf{x}) = g(y_{i-1}, s_i, c_i)\n",
    "$$\n",
    "\n",
    "所谓的上下文向量就是对encoder的所有输出进行加权求和，$a_{ij}$ 表示输出的第 i 个词对encoder第 j 个输出 $h_j$ 的权重。\n",
    "\n",
    "$$\n",
    "c_i = \\sum_{j=1}^{T_x} a_{ij} h_j\n",
    "$$\n",
    "\n",
    "每个 $a_{ij}$ 通过对所有 $e_{ij}$ 进行softmax，而每个 $e_{ij}$ 是decoder的上一个hidden state $s_{i-1}$ 和指定的encoder的输出 $h_j$ 经过某些线性操作 $a$ 计算得分。\n",
    "\n",
    "$$\n",
    "a_{ij} = \\dfrac{exp(e_{ij})}{\\sum_{k=1}^{T} exp(e_{ik})} \n",
    "\\\\\n",
    "e_{ij} = a(s_{i-1}, h_j)\n",
    "$$\n",
    "\n",
    "此外，论文[Effective Approaches to Attention-based Neural Machine Translation](https://arxiv.org/abs/1508.04025)中提出了计算分值的不同方式。这里用到的是第三种。\n",
    "\n",
    "$$\n",
    "score(h_t, \\bar h_s) =\n",
    "\\begin{cases}\n",
    "h_t ^\\top \\bar h_s & dot \\\\\n",
    "h_t ^\\top \\textbf{W}_a \\bar h_s & general \\\\\n",
    "v_a ^\\top \\textbf{W}_a [ h_t ; \\bar h_s ] & concat\n",
    "\\end{cases}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, hidden_dim):\n",
    "        super(Attention, self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.attn = nn.Linear(self.hidden_dim * 2, hidden_dim)\n",
    "        self.v = nn.Parameter(torch.rand(hidden_dim))\n",
    "        self.v.data.normal_(mean=0, std=1. / np.sqrt(self.v.size(0)))\n",
    "\n",
    "    def forward(self, hidden, encoder_outputs):\n",
    "        #  encoder_outputs:(seq_len, batch_size, hidden_size)\n",
    "        #  hidden:(num_layers * num_directions, batch_size, hidden_size)\n",
    "        max_len = encoder_outputs.size(0)\n",
    "\n",
    "        h = hidden[-1].repeat(max_len, 1, 1)\n",
    "        # (seq_len, batch_size, hidden_size)\n",
    "\n",
    "        attn_energies = self.score(h, encoder_outputs)  # compute attention score\n",
    "        return F.softmax(attn_energies, dim=1)  # normalize with softmax\n",
    "\n",
    "    def score(self, hidden, encoder_outputs):\n",
    "        # (seq_len, batch_size, 2*hidden_size)-> (seq_len, batch_size, hidden_size)\n",
    "        energy = F.tanh(self.attn(torch.cat([hidden, encoder_outputs], 2)))\n",
    "        energy = energy.permute(1, 2, 0)  # (batch_size, hidden_size, seq_len)\n",
    "        v = self.v.repeat(encoder_outputs.size(1), 1).unsqueeze(1)  # (batch_size, 1, hidden_size)\n",
    "        energy = torch.bmm(v, energy)  # (batch_size, 1, seq_len)\n",
    "        return energy.squeeze(1)  # (batch_size, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面是加了attention层的decoder，GRU的输出进过全连接层后，又进行了log_softmax操作计算输出词的概率，主要是为了方便NLLLoss损失函数，如果用CrossEntropyLoss损失函数，可以不加log_softmax操作。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, embedding_dim, hidden_dim, num_layers=2, dropout=0.2):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hid_dim = hidden_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        self.embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.attention = Attention(hidden_dim)\n",
    "        self.rnn = nn.GRU(embedding_dim + hidden_dim, hidden_dim,\n",
    "                          num_layers=num_layers, dropout=dropout)\n",
    "        self.out = nn.Linear(embedding_dim + hidden_dim * 2, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        # input = [bsz]\n",
    "        # hidden = [n layers * n directions, batch size, hid dim]\n",
    "        # encoder_outputs = [sent len, batch size, hid dim * n directions]\n",
    "        input = input.unsqueeze(0)\n",
    "        # input = [1, bsz]\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        # embedded = [1, bsz, emb dim]\n",
    "        attn_weight = self.attention(hidden, encoder_outputs)\n",
    "        # (batch_size, seq_len)\n",
    "        context = attn_weight.unsqueeze(1).bmm(encoder_outputs.transpose(0, 1)).transpose(0, 1)\n",
    "        # (batch_size, 1, hidden_dim * n_directions)\n",
    "        # (1, batch_size, hidden_dim * n_directions)\n",
    "        emb_con = torch.cat((embedded, context), dim=2)\n",
    "        # emb_con = [1, bsz, emb dim + hid dim]\n",
    "        _, hidden = self.rnn(emb_con, hidden)\n",
    "        # outputs = [sent len, batch size, hid dim * n directions]\n",
    "        # hidden = [n layers * n directions, batch size, hid dim]\n",
    "        output = torch.cat((embedded.squeeze(0), hidden[-1], context.squeeze(0)), dim=1)\n",
    "        output = F.log_softmax(self.out(output), 1)\n",
    "        # outputs = [sent len, batch size, vocab_size]\n",
    "        return output, hidden, attn_weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们再定义一个Seq2seq类，将encoder和decoder结合起来，通过一个循环，模型对每一个batch从前往后依次生成序列，训练的时候可以使用teacher_forcing随机使用真实词或是模型输出的词作为target，测试的时候就不需要了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder, device, teacher_forcing_ratio=0.5):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "        self.teacher_forcing_ratio = teacher_forcing_ratio\n",
    "\n",
    "    def forward(self, src_seqs, src_lengths, trg_seqs):\n",
    "        # src_seqs = [sent len, batch size]\n",
    "        # trg_seqs = [sent len, batch size]\n",
    "        batch_size = src_seqs.shape[1]\n",
    "        max_len = trg_seqs.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        # tensor to store decoder outputs\n",
    "        outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)\n",
    "        # hidden used as the initial hidden state of the decoder\n",
    "        # encoder_outputs used to compute context\n",
    "        encoder_outputs, hidden = self.encoder(src_seqs, src_lengths)\n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        output = trg_seqs[0, :]\n",
    "        for t in range(1, max_len): # skip sos\n",
    "            output, hidden, _ = self.decoder(output, hidden, encoder_outputs)\n",
    "            outputs[t] = output\n",
    "            teacher_force = random.random() < self.teacher_forcing_ratio\n",
    "            output = (trg_seqs[t] if teacher_force else output.max(1)[1])\n",
    "        return outputs\n",
    "    \n",
    "    def predict(self, src_seqs, src_lengths, max_trg_len=20, start_ix=1):\n",
    "        max_src_len = src_seqs.shape[0]\n",
    "        batch_size = src_seqs.shape[1]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        outputs = torch.zeros(max_trg_len, batch_size, trg_vocab_size).to(self.device)\n",
    "        encoder_outputs, hidden = self.encoder(src_seqs, src_lengths)\n",
    "        output = torch.LongTensor([start_ix] * batch_size).to(self.device)\n",
    "        attn_weights = torch.zeros((max_trg_len, batch_size, max_src_len))\n",
    "        for t in range(1, max_trg_len):\n",
    "            output, hidden, attn_weight = self.decoder(output, hidden, encoder_outputs)\n",
    "            outputs[t] = output\n",
    "            output = output.max(1)[1]\n",
    "            attn_weights[t] = attn_weight\n",
    "        return outputs, attn_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练\n",
    "---\n",
    "这里直接取1000个batch进行更新。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "current loss: 0.8066\n",
      "current loss: 0.3415\n",
      "current loss: 0.2091\n",
      "current loss: 0.1309\n",
      "current loss: 0.0722\n",
      "current loss: 0.0437\n",
      "current loss: 0.0351\n",
      "current loss: 0.0236\n",
      "current loss: 0.0105\n",
      "current loss: 0.0149\n"
     ]
    }
   ],
   "source": [
    "embedding_dim = 100\n",
    "hidden_dim = 100\n",
    "batch_size = 256\n",
    "clip = 5\n",
    "\n",
    "encoder = Encoder(len(src_c2ix) + 1, embedding_dim, hidden_dim)\n",
    "decoder = Decoder(len(trg_c2ix) + 1, embedding_dim, hidden_dim)\n",
    "model = Seq2Seq(encoder, decoder, device).to(device)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "criterion = nn.NLLLoss(ignore_index=0).to(device)\n",
    "\n",
    "model.train()\n",
    "for batch_id in range(1, 1001):\n",
    "    src_seqs, src_lengths, trg_seqs, _ = random_batch(batch_size, pairs, src_c2ix, trg_c2ix)\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    output = model(src_seqs, src_lengths, trg_seqs)\n",
    "    loss = criterion(output.view(-1, output.shape[2]), trg_seqs.view(-1))\n",
    "    loss.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "    optimizer.step()\n",
    "    \n",
    "    if batch_id % 100 == 0:\n",
    "        print('current loss: {:.4f}'.format(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型测试\n",
    "---\n",
    "在进行测试时，生成的句子不超过最大目标句子的长度，同时要保存生成的每个词对原端每个词的attention权重，以便可视化。生成时不超过最大长度，如果下一个词是终止符，生成结束。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def show_attention(input_words, output_words, attentions):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions, cmap='bone')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    ax.set_xticklabels([''] + input_words)\n",
    "    ax.set_yticklabels([''] + output_words)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator())\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator())\n",
    "\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "def evaluate(model, text, src_c2ix, trg_ix2c):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        seq = torch.LongTensor(indexes_from_text(text, src_c2ix)).view(-1, 1).to(device)\n",
    "        outputs, attn_weights = model.predict(seq, [seq.size(0)], max_trg_len)\n",
    "        outputs = outputs.squeeze(1).cpu().numpy()\n",
    "        attn_weights = attn_weights.squeeze(1).cpu().numpy()\n",
    "        output_words = [trg_ix2c[np.argmax(word_prob)] for word_prob in outputs]\n",
    "        show_attention(list('^' + text + '$'), output_words, attn_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面是我随便写的一个日期，可以看出attention的效果还是有的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAVwAAADxCAYAAACH4w+oAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAFcpJREFUeJzt3Xu0HWV5x/HvL4dbQmIiCYImYFBTlaKgctGibdTSHig2\n4mUZUVm4WlkoFG1pK65ldal1KdKLdYmGCClaqdRWlEiD6MLiHU0CCIQYGhEhaA0HCZdECOfsp3/M\nnHazs/dczpk9e5/h92HN4sx+Z5559+3Ju9955x1FBGZm1n+zBl0BM7MnCidcM7OaOOGamdXECdfM\nrCZOuGZmNXHCNTOriROumVlNnHDNzGrihGtmVpO9Bl0BM7MqjI6OxtjYWKFtN27ceE1EjPa5Sntw\nwjWzRhgbG2PDhg2FtpW0qM/V6coJ18waY9jnhnHCNbNGCGCi1Rp0NTI54ZpZQwSBW7hmZv0X0Bru\nfOuEa2bN4T5cM7MaBNBywjUzq4dbuGZmNYgIj1IwM6uLW7hmZjXxsDAzsxokJ80GXYtsTrhm1hju\nUjAzq4NPmpmZ1SNwC9fMrDa+8MHMrCZu4ZqZ1cKzhZmZ1SI8W5iZWX1aHqVgZtZ/ni3MzKxGPmlm\nZlaHCLdwzczq4haumVkNAphwwjUzq4dbuGZmNXHCNTOrQfikmZlZfdzCNTOriROumVkNklEKvrTX\nzKwWnrzGzKwOEe5SMDOrg2+xY2ZWo2EfFjZr0BUwM6tKpN0KeUsRkkYlbZG0VdJ5XcrnS/qqpB9L\n2iTprXkx3cI1s0aICm+TLmkEuBA4AdgGrJe0NiJua9vsLOC2iHiVpAOBLZIui4jdveK6hWtmjREF\n/yvgWGBrRNyRJtDLgRV7HA7mSRIwF/g1MJ4V1C1cM2uMCoeFLQbublvfBhzXsc0ngbXAL4B5wBsi\nsgcCu4VrZo0wOUqhYB/uIkkb2pYzpnDIPwRuAp4GHAV8UtKTsnZwC9fMGqPEsLCxiDg6o/we4JC2\n9SXpY+3eCnw0koNulfQz4DnAj3oFdQvXzJohPWlWZClgPbBM0mGS9gFWknQftLsLeCWApIOAZwN3\nZAV1C9fMGqHKCx8iYlzS2cA1wAiwJiI2STozLV8FfAi4VNItgIB3R8RYVlwnXDNrjCovfIiIdcC6\njsdWtf39C+APysR0wjWzxig45GtgnHDNrDGG/MpeJ1wza4Zg+OdScMI1s2ao8NLefnHCNbNG8PSM\nZmY1csI1M6uJ+3DNzGpReCawgXHCNbNGiPCwMDOz2niUgplZDTwO18ysRh6lYGZWhxI3iBwUJ1wz\naw4nXDOzerQmnHDNzPouGRbmhGtmVgsnXDOzWvikmZlZbaLlhGtm1nfuwzUzq1EM+aW9s/p9AEl7\nSfpPSWOSjuhSvkDSO3rsu1TSrQWO8f0q6lq0XhXEXiNpe5Hn1mXfQq9JToxzJG2WdNl04lRJ0n6S\nfiTpx5I2SfpAH47xcB9iTvv9yInf873q52e039LX7fSq405OYJO3DErfEy7waeAnwKuBf5O0pKN8\nATCtD01E/E6R7ZQo+pynXa8MlwKjfYpdxDuAEyLiTdMJUvL1zPMo8IqIOBI4ChiV9OKKYs9kWe9V\nZZ/Rit/LvGO9Hbga+JCk6yQdXEngCKJVbBmUvr7Akt4PPBAR50bEd4E/Bb4gaX7bZh8FninpJkkX\ndAkzIukzaavn65JmdzlOz5ZL+i/pFkmfA24FDmkr2z9tff9Y0q2S3lCkXpI+KOldbesflvTO7Ffj\n/0XEt4FfF92+i70kXZa2fP5D0pyO+r05bS3eJOkiSSNtZauAZwBXS/rzzsCS/iJ9LW5tf45t5Vmv\nZ8/j5onE5Pu4d7oU/mZIOk3Szel7+S9F90v3zXvOhWJLeoakGyUdU+LYX5G0Mf18n9FRlvlekfPd\nyXs/er2X6eObs753vV6zvH0lzQM+ALwJ+BvgdGBn0dcrT6SX9+YtA1O0gv1agKXArRll48BR6foX\ngTd32e7hnPgt4MVdyl4LfKZtfX6Jet2Q/j0L+CmwsKrnXWC/AI5P19cAf9lW/lzgq8De6fqngNM6\nYtwJLOoS+0XALcD+wFxgE/CCIq9nkeMWeG4jwE3Aw8D5Jfb7beD2yecEHNBjuz0+J3nPOS/25PsI\nPBu4ETiy5HM+IP3/7DTOwo7yru9Vgc9okc9Br/dyKRnfu6zXrMC++wO/An4fOL3s5z9rOfRZy+Ki\nq64ptAAbqjx20WUmnDT7WUTclP69keQNLevnEXF9l8dvAf5e0vnAVRHxnSLBIuJOSfdJegFwEHBj\nRNw3hXpN1d0R8b30788D5wB/l66/kuQLsV4SJF/k7QXjvhT4ckTsBJB0BfAykkTSrtvrOZ3jAhAR\nE8BRkhYAX5Z0REQU6R99BfDvETGWxinz6yHvOReJfSBwJfCaiLitxLEBzpF0Svr3IcAyoIrPUtH3\no9d3I+t7l/ea9dw3InZKehvwEeBgJed13hcRu0o/wy7SpD60ZkLCfbTt7wmSD05ZXX+yRMTtkl4I\nnAT8raRrI+KDBWNeTPJz6GCSVmadOj9V7esCPhsR7+nj8bu9npUdNyJ2SPovkn7uvp2QqtADwF0k\niahwwpW0nKSl95KI2CXpOmC/iupU9P3o9XN+Ot+7zH0jYq2km4FXAUcD5wIfKhG/uwhi4gk+SqGA\nh4B5gziwpKcBuyLi88AFwAtL1OvLJAnhGOCavlWyu0MlvST9+1Tgu21l1wKvk/QUAEkHSHp6wbjf\nAV4taY6k/YFT0seKmM5xkXRg2rIl7fM7geRkaxHfBF4vaeHksYsel/znXCT27nS/0ySdWuLY84H7\n02T7HKDsScKsz+i03o8cU/6cSJrbVo+HgM1U+P0v0ZUzEANv4UbEfZK+p2RozdUR8Vc1Hv55wAWS\nWsBjwNuL1isidqetsB3pT+HCJH0BWA4skrQNeH9EXFIixBbgLElrSFpUn26r122S3gt8XclZ58eA\ns4Cf5wWNiBskXQr8KH3o4ojo7E7ote+Uj5t6KvDZ9MTOLOCLEXFVwWNvkvRh4FuSJkh+2p5ecN/M\n51w0dvpT+WTgG5Iejoi1BQ7/NeBMSZtJ3tNuP+2z6t7zM1rB+5F13Cl/TkhOhl4ELAQWkfwyKPOP\nVE7dqorUHxr2Po9hlX6IbwBeHxH/Pej6mM0kkpYCyyPi0qpiHvqsZfHuj3280LZnv/bkjRFxdFXH\nLmoYuhRmHEmHA1uBa51szaZkB8mIlOqEuxQaKT0T/YxB18NspoqI6hMuQWvIT5o54ZpZYwx7F6kT\nrpk1QsyA2cJq68PtvGyxyvJhjT2s9XJsxx507Lx9pyyGe/aaOk+a5b3A0ykf1tjDWi/HduxBx+5L\nwo1WsWVQPErBzBqjylEKkkaVTO6zVdJ5PbZZrmRyoE2SvpUXs299uJL2eFbdHquqfFhjD2u9HNux\nBx27o2wsIg7MipUrglZFE5CnF+BcSHLF4zaSOSnWts+VoeTKyE8BoxFx1+RVfVncwjWzYTD9K+Co\ntIV7LLA1Iu6IiN3A5cCKjm1OBa6IiLtIjp07WZMTrpk1Q1DlBOSLgbvb1relj7X7LeDJSiZR3yjp\ntLyguV0KSiY+XkkyQcc/k0xSsQL4XkT8oEjNzcxqUXwEwiJJG9rWV0fE6pJH24tkCsxXksyI9gNJ\n10fE7Vk75DkIOB54JvBe4K9JJhX+YeeG6VCP/gz3MDPLVOqy3bGcuRTuoe1uJsCS9LF224D70nmB\nd0r6NnAkyYT1XeUm3IiYPDu3BXhLzrargdWQ35luZla1VnX3K1sPLJN0GEmiXcmes5pdCXxS0l7A\nPsBxwD9mBfWVZmbWCJH24VYTK8YlnU0y1/UIsCadqvPMtHxVRGyW9DXgZpJbFV2cd4cSJ1wza4wq\nL+2NiHXAuo7HVnWsX0By84JCCifcdLLrk4HtEXFE0f2aZtas3jeibbVKzUNuwP77z+9ZtmBB9rDG\n+fOzy/fdp/ddYR58cCxz32339OyGA+DRRyu5BdeMst9+czPL58zpfeOGvPGxO3b8akp16tSkuRQu\nJbmljJnZECp1t+SBKNzCjYhvp7O0m5kNnxkwW5j7cM2sEQKIiSdQwvU4XDMbpCdUC9fjcM1sYAbc\nP1uEuxTMrDGqGofbL2WGhX0BWE5yDfI24P0RcUm/KjasPPSrWgcdtLRn2cKFnXOFPN6Tn3xwZvn2\n7Xf2LPvV9uzJqcbHd2eWPxE98sjD0yqvw7C3cAsPC4uIN0bEU4H9gHuBU/pWKzOzkiqenrEvptKl\n8E5gM/CkiutiZjZ1EURFE5D3S6n5cCUtAf4IuLg/1TEzm7phv6dZ2Rbux0mmZ+x6DZ+HhZnZIDWm\nD1fS5DwKG3ttExGrI+LonHkmzcyqF83qwz0e+GNJJ5GcOHuSpM9HxJv7UzUzs+ImT5oNszKjFN4T\nEUsiYinJZLzfdLI1s+ERtCZahZZB8YUPJb3q5LN6lv30pzdm7nvb5rxbwGX966xp7Junn7Gzjd27\nrWfZAw/cm7nv7NnZA2Ueeui+nmWPPLIzc9+JCY+37iRlt89mzepdvtde+2TuW8l0l02dvCYirgOu\nq7QmZmbT1cSEa2Y2jIY835YbhztJ0jpJT6u6MmZmU9XUK82IiJO6Pe5xuGY2MBXeRLJfPD2jmTVE\n5N47bdDch2tmjTHsoxTch2tmzRFRbBmQMvPhHgJ8DjiIpH/69cA/9aleQ+urV13Ys+yBXdljCefP\nmTONI0/vQ5I1hnLffbPr1c95Th/MGCub7xeV1cPyRc6sLxMZFxRMTIxXXZ09RMP6cMeBcyPiBknz\ngI2SvhERt/WpbmZmpQx5j0Kp26T/Evhl+vdDkjYDiwEnXDMbAg29p5mkpcALgB9WWRkzsykLmjdK\nQdJc4EvAuyLiwY4yj8M1s4EImtWHi6S9SZLtZRFxRWe5x+Ga2SA1pktBkoBLgM0R8Q/9q5KZ2VQM\ndshXEWXG4R4PvAV4haSb0qXrJb5mZrVr0h0fIuK7kt4FvA0YAT4TEev6VrMZaHrjbGFkpPfbkTeO\nMe9DlPxA6a6f42zN6tSaaEgLV9IRJMn2WOBI4GRJz+pXxczMypgJs4WV6VJ4LvDDiNgVEePAt4DX\n9KdaZmYlzYAuhTIJ91bgZZIWSpoDnAQc0r6BpDMkbZC0ocpKmpnlK5ZsZ0TCjYjNwPnA14GvATcB\nEx3b+DbpZjYwVSZcSaOStkjaKum8jO2OkTQu6XV5MUvNFhYRl0TEiyLid4H7gdvL7G9m1k/RikJL\nHkkjwIXAicDhwBslHd5ju8mGaK5SCVfSU9L/H0rSf/uvZfY3M+uXydnCqki4JIMDtkbEHRGxG7gc\nWNFluz8juRhse5GgZS/t/ZKkhcBjwFkRsaPk/o2WNawL8m+9nVU+e/a8zH333Xd2Zvl0pmd87LFH\nMsuz5F3bPmdO7+eV93pl7QvZz2tOzuuZd9vuVsZUhffee3fmvllD9CB7iF/ez+Hx8ccyy/t5y/ss\n+d+NaqZvrLB/djHQ/kZuA45r30DSYuAU4OXAMUWClkq4EfGydOKa5RFxbZl9zcz6q9QJsUUdJ/dX\np1MTlPFx4N0R0cr7R3RS2bkU3g6cA8yVdDqwMiL+p2QlzcyqV24C8rGck/v38PhRWEvSx9odDVye\nJttFwEmSxiPiK72ClplLYR7wAWAUeD5wHbCz6P5mZv1WYZfCemCZpMNIEu1K4NSOYx02+bekS4Gr\nspItlGvhtkg6gA5ID3Zn5waentHMBmXySrNKYkWMSzobuIZkKoM1EbFJ0plp+aqpxC0zl8JOSW8D\nPgIcnF7q+76I2NW2jadnNLMBCaLCCcjTuWLWdTzWNdFGxOlFYpYdh7uW5OaRHwMOBM4ts7+ZWd8E\nRKvYMihl+nDnAgvT1YeAzaTdC2Zmw6AxE5ADewMXkSTdRcBddHQidxgDft62vih9rJfplA9F7C5j\nCSur129+81DfYneZnrG213PnzgemHHuQ9XbscvsW+G48PSNOYY1JuBFxPzDaNg730pztD2xfl7Qh\naxjGdMqHNfaw1suxHXvQsfP2nYoqT5r1y1Tu2ruDZOIaM7PhEUFromF37U0v53XCNbPh08AW7lTl\nXTY3nfJhjT2s9XJsxx507LKX0RYSA5oroigNe5+HmVkRCxYcFMuXryy07ZVXfmLjIObtrrOFa2bW\nR0EMcpBtAU64ZtYYw/6L3QnXzBojb/7lQXPCNbNGSO5X5oRrZlYPdymYmdVj2IeFOeGaWWP4pJmZ\nWS2CViv7xqOD5oRrZo0Q4RaumVltnHDNzGrihGtmVovwsDAzs7oEvvDBzKzvInxpr5lZTcJ9uGZm\ndfFcCmZmNXEL18ysJk64ZmZ1CA8LMzOrRQCt8FwKZmY18CgFM7PaOOGamdXECdfMrAbJOTOPwzUz\nq0EQvrTXzKwevqeZmVlN3IdrZlaLGPo+3FmDroCZWRUm72lWZClC0qikLZK2SjqvS/mbJN0s6RZJ\n35d0ZF5Mt3DNrDGq6lKQNAJcCJwAbAPWS1obEbe1bfYz4Pci4n5JJwKrgeOy4jrhmlljVDgB+bHA\n1oi4A0DS5cAK4P8SbkR8v23764EleUGdcM2sIQKK9+EukrShbX11RKxuW18M3N22vo3s1uufAFfn\nHdQJ18wao8SwsLGIOLqKY0p6OUnCfWnetk64ZtYIkyfNKnIPcEjb+pL0sceR9HzgYuDEiLgvL6hH\nKZhZY1Q4SmE9sEzSYZL2AVYCa9s3kHQocAXwloi4vUhQt3DNrCGqG4cbEeOSzgauAUaANRGxSdKZ\nafkq4H3AQuBTkgDG87opNOxXZpiZFTF79txYuvR5hbb9yU+u31hVH24ZbuGaWSNU3IfbF064ZtYQ\nvqeZmVltguGeS8EJ18waw10KZma1iCov7e0LJ1wzawTfYsfMrEbuUjAzq4kTrplZLTwszMysNr6J\npJlZDSKg1ZoYdDUyOeGaWUMUv1/ZoDjhmlljOOGamdXECdfMrCa+8MHMrA7hYWFmZrUIoOUWrplZ\nPdylYGZWCw8LMzOrjROumVkNfE8zM7PaBOFLe83M6uHJa8zMauIuBTOzmjjhmpnVICI8DtfMrC5u\n4ZqZ1cS3STczq4tbuGZmdQgCt3DNzPrOV5qZmdXICdfMrCZOuGZmtQjfJt3MrA4zoQ931qArYGZW\nmcn7muUtBUgalbRF0lZJ53Upl6RPpOU3S3phXkwnXDNriCj8Xx5JI8CFwInA4cAbJR3esdmJwLJ0\nOQP4dF5cJ1wza4yIVqGlgGOBrRFxR0TsBi4HVnRsswL4XCSuBxZIempWUPfhmlljVHhp72Lg7rb1\nbcBxBbZZDPyyV1AnXDNrimuARQW33U/Shrb11RGxug91ehwnXDNrhIgYrTDcPcAhbetL0sfKbvM4\n7sM1M9vTemCZpMMk7QOsBNZ2bLMWOC0drfBi4IGI6NmdAG7hmpntISLGJZ1N0k0xAqyJiE2SzkzL\nVwHrgJOArcAu4K15cTXsA4XNzJrCXQpmZjVxwjUzq4kTrplZTZxwzcxq4oRrZlYTJ1wzs5o44ZqZ\n1cQJ18ysJv8LEoOrIojVb1MAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f6b478c5128>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "text = 'thirsty 1 before 3 clock afternon'\n",
    "evaluate(model, text, src_c2ix, trg_ix2c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}