Tensorflow seq2seq template based on g2p.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll write a simple template for seq2seq using Tensorflow. For demonstration, we attack the g2p task. G2p is a task of converting graphemes (spelling) to phonemes (pronunciation). It's a very good source for this purpose as it's simple enough for you to up and run. If you want to know more about g2p, see my [repo](https://github.com/kyubyong/g2p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"kyubyong\"\n",
    "__address__ = \"https://github.com/kyubyong/nlp_made_easy\"\n",
    "__email__ = \"kbpark.linguist@gmail.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TEXqpZ_U738q"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from distance import levenshtein\n",
    "import os\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "_7vuctbU7381",
    "outputId": "f8ee2cbf-1f04-432f-ba42-d25fec61669b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.12.0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "B6te4HKk738_"
   },
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CWS2hkce739C"
   },
   "outputs": [],
   "source": [
    "class Hparams:\n",
    "    batch_size = 128\n",
    "    enc_maxlen = 20\n",
    "    dec_maxlen = 20\n",
    "    num_epochs = 10\n",
    "    hidden_units = 128\n",
    "    graphemes = [\"<pad>\", \"<unk>\", \"</s>\"] + list(\"abcdefghijklmnopqrstuvwxyz\")\n",
    "    phonemes = [\"<pad>\", \"<unk>\", \"<s>\", \"</s>\"] + ['AA0', 'AA1', 'AA2', 'AE0', 'AE1', 'AE2', 'AH0', 'AH1', 'AH2', 'AO0',\n",
    "                    'AO1', 'AO2', 'AW0', 'AW1', 'AW2', 'AY0', 'AY1', 'AY2', 'B', 'CH', 'D', 'DH',\n",
    "                    'EH0', 'EH1', 'EH2', 'ER0', 'ER1', 'ER2', 'EY0', 'EY1', 'EY2', 'F', 'G', 'HH',\n",
    "                    'IH0', 'IH1', 'IH2', 'IY0', 'IY1', 'IY2', 'JH', 'K', 'L', 'M', 'N', 'NG', 'OW0', 'OW1',\n",
    "                    'OW2', 'OY0', 'OY1', 'OY2', 'P', 'R', 'S', 'SH', 'T', 'TH', 'UH0', 'UH1', 'UH2', 'UW',\n",
    "                    'UW0', 'UW1', 'UW2', 'V', 'W', 'Y', 'Z', 'ZH']\n",
    "    lr = 0.001\n",
    "    eval_steps = 500\n",
    "    logdir = \"log/04\"\n",
    "hp = Hparams()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nz-hD6dn739L"
   },
   "source": [
    "# Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-as4PHs-739N"
   },
   "outputs": [],
   "source": [
    "import nltk\n",
    "# nltk.download('cmudict')# <- if you haven't downloaded, do this.\n",
    "from nltk.corpus import cmudict\n",
    "cmu = cmudict.dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "39gQ3vOi739S"
   },
   "outputs": [],
   "source": [
    "def load_vocab():\n",
    "    g2idx = {g: idx for idx, g in enumerate(hp.graphemes)}\n",
    "    idx2g = {idx: g for idx, g in enumerate(hp.graphemes)}\n",
    "\n",
    "    p2idx = {p: idx for idx, p in enumerate(hp.phonemes)}\n",
    "    idx2p = {idx: p for idx, p in enumerate(hp.phonemes)}\n",
    "\n",
    "    return g2idx, idx2g, p2idx, idx2p # note that g and p mean grapheme and phoneme, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Zslytxn6739Z"
   },
   "outputs": [],
   "source": [
    "def prepare_data():\n",
    "    words = [\" \".join(list(word)) for word, prons in cmu.items()]\n",
    "    prons = [\" \".join(prons[0]) for word, prons in cmu.items()]\n",
    "    indices = list(range(len(words)))\n",
    "    from random import shuffle\n",
    "    shuffle(indices)\n",
    "    words = [words[idx] for idx in indices]\n",
    "    prons = [prons[idx] for idx in indices]\n",
    "    num_train, num_test = int(len(words)*.8), int(len(words)*.1)\n",
    "    train_words, eval_words, test_words = words[:num_train], \\\n",
    "                                          words[num_train:-num_test],\\\n",
    "                                          words[-num_test:]\n",
    "    train_prons, eval_prons, test_prons = prons[:num_train], \\\n",
    "                                          prons[num_train:-num_test],\\\n",
    "                                          prons[-num_test:]    \n",
    "    return train_words, eval_words, test_words, train_prons, eval_prons, test_prons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WHBXkAPG739j"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q u a l i t a t i v e\n",
      "K W AA1 L AH0 T EY2 T IH0 V\n"
     ]
    }
   ],
   "source": [
    "train_words, eval_words, test_words, train_prons, eval_prons, test_prons = prepare_data()\n",
    "print(train_words[0])\n",
    "print(train_prons[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "def drop_lengthy_samples(words, prons, enc_maxlen, dec_maxlen):\n",
    "    \"\"\"We only include such samples less than maxlen.\"\"\"\n",
    "    _words, _prons = [], []\n",
    "    for w, p in zip(words, prons):\n",
    "        if len(w.split()) + 1 > enc_maxlen: continue\n",
    "        if len(p.split()) + 1 > dec_maxlen: continue # 1: <EOS>\n",
    "        _words.append(w)\n",
    "        _prons.append(p)\n",
    "    return _words, _prons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_words, train_prons = drop_lengthy_samples(train_words, train_prons, hp.enc_maxlen, hp.dec_maxlen)\n",
    "# We do NOT apply this constraint to eval and test datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KfHMTzeH7394"
   },
   "source": [
    "# Data Loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(inp, type, dict):\n",
    "    '''type: \"x\" or \"y\"'''\n",
    "    inp_str = inp.decode(\"utf-8\")\n",
    "    if type==\"x\": tokens = inp_str.split() + [\"</s>\"]\n",
    "    else: tokens = [\"<s>\"] + inp_str.split() + [\"</s>\"]\n",
    "\n",
    "    x = [dict.get(t, dict[\"<unk>\"]) for t in tokens]\n",
    "    return x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "G6jSAgus7399"
   },
   "outputs": [],
   "source": [
    "def generator_fn(words, prons):\n",
    "    '''\n",
    "    words: 1d byte array. e.g., [b\"w o r d\", ]\n",
    "    prons: 1d byte array. e.g., [b'W ER1 D', ]\n",
    "    \n",
    "    yields\n",
    "    xs: tuple of\n",
    "        x: list of encoded x. encoder input\n",
    "        x_seqlen: scalar.\n",
    "        word: string\n",
    "        \n",
    "    ys: tuple of\n",
    "        decoder_input: list of decoder inputs\n",
    "        y: list of encoded y. label.\n",
    "        y_seqlen: scalar.\n",
    "        pron: string\n",
    "    '''\n",
    "    g2idx, idx2g, p2idx, idx2p = load_vocab()\n",
    "    for word, pron in zip(words, prons):\n",
    "        x = encode(word, \"x\", g2idx)\n",
    "        y = encode(pron, \"y\", p2idx)\n",
    "        decoder_input, y = y[:-1], y[1:]\n",
    "\n",
    "        x_seqlen, y_seqlen = len(x), len(y)\n",
    "        yield (x, x_seqlen, word), (decoder_input, y, y_seqlen, pron)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QEvz4YTR73-I"
   },
   "outputs": [],
   "source": [
    "def input_fn(words, prons, batch_size, shuffle=False):\n",
    "    '''Batchify data\n",
    "    words: list of words. e.g., [\"word\", ]\n",
    "    prons: list of prons. e.g., ['W ER1 D',]\n",
    "    batch_size: scalar.\n",
    "    shuffle: boolean\n",
    "    '''\n",
    "    shapes = ( ([None], (), ()),\n",
    "               ([None], [None], (), ())  )\n",
    "    types = (  (tf.int32, tf.int32, tf.string),\n",
    "               (tf.int32, tf.int32, tf.int32, tf.string)  )\n",
    "    paddings = (  (0, 0, ''),\n",
    "                  (0, 0, 0, '')  )\n",
    "\n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        generator_fn,\n",
    "        output_shapes=shapes,\n",
    "        output_types=types,\n",
    "        args=(words, prons)) # <- converted to np string arrays\n",
    "        \n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(128*batch_size)    \n",
    "    dataset = dataset.repeat() # iterate forever\n",
    "    dataset = dataset.padded_batch(batch_size, shapes, paddings).prefetch(1)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(words, prons, batch_size, shuffle=False):\n",
    "    '''Gets training / evaluation mini-batches\n",
    "    fpath1: source file path. string.\n",
    "    fpath2: target file path. string.\n",
    "    maxlen1: source sent maximum length. scalar.\n",
    "    maxlen2: target sent maximum length. scalar.\n",
    "    vocab_fpath: string. vocabulary file path.\n",
    "    batch_size: scalar\n",
    "    shuffle: boolean\n",
    "\n",
    "    Returns\n",
    "    batches\n",
    "    num_batches: number of mini-batches\n",
    "    num_samples\n",
    "    '''\n",
    "    batches = input_fn(words, prons, batch_size, shuffle=shuffle)\n",
    "    num_batches = calc_num_batches(len(words), batch_size)\n",
    "    return batches, num_batches, len(words)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "22mif4xf73-M"
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_idx_to_token_tensor(inputs, idx2token):\n",
    "    '''Converts int32 tensor to string tensor.\n",
    "    inputs: 1d int32 tensor. indices.\n",
    "    idx2token: dictionary\n",
    "\n",
    "    Returns\n",
    "    1d string tensor.\n",
    "    '''\n",
    "    def my_func(inputs):\n",
    "        return \" \".join(idx2token[elem] for elem in inputs)\n",
    "\n",
    "    return tf.py_func(my_func, [inputs], tf.string)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HA39FU4-73-O"
   },
   "outputs": [],
   "source": [
    "class Net:\n",
    "    def __init__(self, hp):\n",
    "        self.g2idx, self.idx2g, self.p2idx, self.idx2p = load_vocab()\n",
    "        self.hp = hp\n",
    "    \n",
    "    def encode(self, xs):\n",
    "        '''\n",
    "        xs: tupple of \n",
    "            x: (N, T). int32\n",
    "            seqlens: (N,). int32\n",
    "            words: (N,). string\n",
    "            \n",
    "        returns\n",
    "        last hidden: (N, hidden_units). float32\n",
    "        words: (N,). string\n",
    "        '''\n",
    "        with tf.variable_scope(\"encode\", reuse=tf.AUTO_REUSE):\n",
    "            x, seqlens, words = xs\n",
    "            x = tf.one_hot(x, len(self.g2idx))\n",
    "            cell = tf.contrib.rnn.GRUCell(self.hp.hidden_units)\n",
    "            _, last_hidden = tf.nn.dynamic_rnn(cell, x, seqlens, dtype=tf.float32)\n",
    "            \n",
    "        return last_hidden, words\n",
    "        \n",
    "    \n",
    "    def decode(self, ys, h0=None):\n",
    "        '''\n",
    "        ys: tupple of \n",
    "            decoder_inputs: (N, T). int32\n",
    "            y: (N, T). int32\n",
    "            seqlens: (N,). int32\n",
    "            prons: (N,). string.\n",
    "        h0: initial hidden state. (N, hidden_units)\n",
    "        \n",
    "        returns\n",
    "        logits: (N, T, len(p2idx)). float32. before softmax\n",
    "        y_hat: (N, T). int32.\n",
    "        y: (N, T). int32. label.\n",
    "        prons: (N,). string. ground truth phonemes \n",
    "        last_hidden: (N, hidden_units). This is for autoregressive inference\n",
    "        '''\n",
    "        decoder_inputs, y, seqlens, prons = ys\n",
    "       \n",
    "        with tf.variable_scope(\"decode\", reuse=tf.AUTO_REUSE):\n",
    "            inputs = tf.one_hot(decoder_inputs, len(self.p2idx))\n",
    "            \n",
    "            cell = tf.contrib.rnn.GRUCell(self.hp.hidden_units)\n",
    "            outputs, last_hidden = tf.nn.dynamic_rnn(cell, inputs, initial_state=h0, dtype=tf.float32)\n",
    "\n",
    "            # projection\n",
    "            logits = tf.layers.dense(outputs, len(self.p2idx))\n",
    "            y_hat = tf.to_int32(tf.argmax(logits, axis=-1))\n",
    "        \n",
    "        return logits, y_hat, y, prons, last_hidden\n",
    "            \n",
    "    def train(self, xs, ys):\n",
    "        # forward\n",
    "        last_hidden, words = self.encode(xs)\n",
    "        logits, y_hat, y, prons, last_hidden = self.decode(ys, h0=last_hidden)\n",
    "        \n",
    "        # train scheme\n",
    "        ce = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=y)\n",
    "        nonpadding = tf.to_float(tf.not_equal(y, self.p2idx[\"<pad>\"])) # 0: <pad>\n",
    "        loss = tf.reduce_sum(ce*nonpadding) / (tf.reduce_sum(nonpadding)+1e-7)\n",
    "\n",
    "        global_step = tf.train.get_or_create_global_step()\n",
    "        train_op = tf.train.AdamOptimizer(hp.lr).minimize(loss, global_step=global_step)\n",
    "        \n",
    "        return loss, train_op, global_step\n",
    "\n",
    "    \n",
    "    def eval(self, xs, ys):\n",
    "        '''Predicts autoregressively\n",
    "        At inference input ys is ignored.\n",
    "        Returns\n",
    "        y_hat: (N, T2)\n",
    "        '''\n",
    "        decoder_inputs, y, seqlens, prons = ys\n",
    "        decoder_inputs = tf.ones((tf.shape(xs[0])[0], 1), tf.int32) * self.p2idx[\"<s>\"]\n",
    "        ys = (decoder_inputs, y, seqlens, prons)\n",
    "\n",
    "        last_hidden, words = self.encode(xs)\n",
    "\n",
    "        \n",
    "        h0 = last_hidden\n",
    "        y_hats = []\n",
    "        print(\"Inference graph is being built. Please be patient.\")\n",
    "        for t in tqdm(range(self.hp.dec_maxlen)):\n",
    "            _, y_hat, _, _, h0 = self.decode(ys, h0)\n",
    "            if tf.reduce_sum(y_hat, 1)==0: break\n",
    "           \n",
    "            ys = (y_hat, y, seqlens, prons)\n",
    "            y_hats.append(tf.squeeze(y_hat))\n",
    "        y_hats = tf.stack(y_hats, 1)\n",
    "        \n",
    "        # monitor a random sample\n",
    "        n = tf.random_uniform((), 0, tf.shape(y_hats)[0]-1, tf.int32)\n",
    "        word = words[n]\n",
    "        pred = convert_idx_to_token_tensor(y_hats[n], self.idx2p)\n",
    "        pron = prons[n]\n",
    "        \n",
    "        return y_hats, word, pred, pron\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PKllLnfp73-V"
   },
   "source": [
    "# Train & Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_num_batches(total_num, batch_size):\n",
    "    return total_num // batch_size + int(total_num % batch_size != 0) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation metric\n",
    "def per(ref, hyp):\n",
    "    '''Calc phoneme error rate\n",
    "    hyp: list of predicted phoneme sequences. e.g., [[\"B\", \"L\", \"AA1\", \"K\", \"HH\", \"AW2\", \"S\"], ...]\n",
    "    ref: list of ground truth phoneme sequences. e.g., [[\"B\", \"L\", \"AA1\", \"K\", \"HH\", \"AW2\", \"S\"], ...]\n",
    "    '''\n",
    "    num_phonemes, num_erros = 0, 0\n",
    "    g2idx, idx2g, p2idx, idx2p = load_vocab()\n",
    "    for r, h in zip(ref, hyp):\n",
    "        r = r.split()\n",
    "        h = \" \".join(idx2p[idx] for idx in h)\n",
    "        h = h.split(\"</s>\")[0].strip().split()\n",
    "        \n",
    "        num_phonemes += len(r)\n",
    "        num_erros += levenshtein(h, r)\n",
    "#         print(h, r)\n",
    "    per = round(num_erros / num_phonemes, 2)\n",
    "    return per"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "64f3a-fb73-Y",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "# prepare batches\n",
    "train_batches, num_train_batches, num_train_samples = get_batch(train_words, train_prons,\n",
    "                         hp.batch_size, shuffle=True)\n",
    "eval_batches, num_eval_batches, num_eval_samples = get_batch(eval_words, eval_prons,\n",
    "                         hp.batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iL3CK4NW73-g"
   },
   "outputs": [],
   "source": [
    "# create a iterator of the correct shape and type\n",
    "iter = tf.data.Iterator.from_structure(train_batches.output_types, train_batches.output_shapes)\n",
    "\n",
    "# create the initialisation operations\n",
    "train_init_op = iter.make_initializer(train_batches)\n",
    "eval_init_op = iter.make_initializer(eval_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "# variable specs\n",
    "def print_variable_specs(fpath):\n",
    "    def get_size(shp):\n",
    "        size = 1\n",
    "        for d in range(len(shp)):\n",
    "            size *=shp[d]\n",
    "        return size\n",
    "\n",
    "    params, num_params = [], 0\n",
    "    for v in tf.global_variables():\n",
    "        params.append(\"{}==={}\\n\".format(v.name, v.shape))\n",
    "        num_params += get_size(v.shape)\n",
    "    print(\"num_params:\", num_params)\n",
    "#     with open(fpath, 'w') as fout:\n",
    "#         fout.write(\"num_params: {}\\n\".format(num_params))\n",
    "#         fout.write(\"\\n\".join(params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inference graph is being built. Please be patient.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b3f85d3acbc43dc8ebf4eec2ea9335c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=20), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load model\n",
    "net = Net(hp)\n",
    "xs, ys = iter.get_next()\n",
    "loss, train_op, global_step = net.train(xs, ys)\n",
    "y_hat, word, pred, pron = net.eval(xs, ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 14303
    },
    "colab_type": "code",
    "id": "frKAWTc873-q",
    "outputId": "2d464429-3e88-4f3f-9a9d-d64094d995f9",
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variables initialized\n",
      "num_params: 444513\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc6815075e624866b77282ec647338ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=7721), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch= 1 is done!\n",
      "wrd: m a p e l\n",
      "exp: M AE1 P AH0 L\n",
      "got: M AE1 P L </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s> </s>\n",
      "per=0.56\n",
      "\n",
      "epoch= 2 is done!\n",
      "wrd: s t e a r i c\n",
      "exp: S T IY1 R IH0 K\n",
      "got: S T EH1 R IY0 </s> </s> </s> </s> </s> </s> Z </s> </s> N </s> Z </s> N </s>\n",
      "per=0.41\n",
      "\n",
      "epoch= 3 is done!\n",
      "wrd: s c o l d e d\n",
      "exp: S K OW1 L D AH0 D\n",
      "got: S K OW1 D L D </s> </s> </s> </s> </s> L T V UW0 T </s> UW0 V UW0\n",
      "per=0.34\n",
      "\n",
      "epoch= 4 is done!\n",
      "wrd: s c o l d e d\n",
      "exp: S K OW1 L D AH0 D\n",
      "got: S K OW1 L D </s> </s> D </s> </s> L </s> T L T V EY0 T W EH1\n",
      "per=0.30\n",
      "\n",
      "epoch= 5 is done!\n",
      "wrd: n o a\n",
      "exp: N OW1 AH0\n",
      "got: N OW1 </s> </s> </s> </s> </s> EH1 L EH1 N T </s> EH1 T </s> </s> EH1 T </s>\n",
      "per=0.26\n",
      "\n",
      "epoch= 6 is done!\n",
      "wrd: c o n f i d e n t i a l l y\n",
      "exp: K AA2 N F AH0 D EH1 N SH AH0 L IY0\n",
      "got: K AH0 N F EH1 D AH0 N T AH0 L IY0 </s> </s> </s> </s> </s> </s> AW1 </s>\n",
      "per=0.24\n",
      "\n",
      "epoch= 7 is done!\n",
      "wrd: d o l i n g\n",
      "exp: D OW1 L IH0 NG\n",
      "got: D OW1 L IH0 NG </s> </s> </s> </s> G </s> </s> </s> </s> </s> T AH1 V AH1 T\n",
      "per=0.22\n",
      "\n",
      "epoch= 8 is done!\n",
      "wrd: e q u a t o r\n",
      "exp: IH0 K W EY1 T ER0\n",
      "got: EH1 K W AA2 T ER0 </s> </s> </s> </s> TH </s> </s> AO1 </s> TH AO1 AH1 T AH2\n",
      "per=0.22\n",
      "\n",
      "epoch= 9 is done!\n",
      "wrd: c a r t e r s v i l l e\n",
      "exp: K AA1 R T ER0 Z V IH2 L\n",
      "got: K AA1 R T ER0 S IH0 V L IY0 </s> </s> </s> G TH </s> G M AH1 TH\n",
      "per=0.20\n",
      "\n",
      "epoch= 10 is done!\n",
      "wrd: s e e f e l d t\n",
      "exp: S IY1 F IH0 L T\n",
      "got: S IY1 F EH2 L T </s> </s> D </s> </s> </s> </s> </s> </s> </s> </s> </s> T </s>\n",
      "per=0.19\n",
      "\n",
      "Training Done!\n"
     ]
    }
   ],
   "source": [
    "# Session\n",
    "saver = tf.train.Saver()\n",
    "with tf.Session() as sess:\n",
    "    ckpt = tf.train.latest_checkpoint(hp.logdir)\n",
    "    if ckpt is None:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        print(\"Variables initialized\")\n",
    "    else:\n",
    "        saver.restore(sess, ckpt)\n",
    "        print(\"Restored from file: \", ckpt)\n",
    "\n",
    "    print_variable_specs('specs')\n",
    "\n",
    "    sess.run(train_init_op)\n",
    "    total_steps = hp.num_epochs*num_train_batches\n",
    "    _gs = sess.run(global_step)\n",
    "    for _ in tqdm(range(_gs, total_steps+1)):\n",
    "        # training\n",
    "        _, _gs, _loss = sess.run([train_op, global_step,loss]) \n",
    "\n",
    "        epoch = math.ceil(_gs / num_train_batches)\n",
    "            \n",
    "        if _gs and _gs % num_train_batches == 0: # Be careful that you should evaluate at every epoch due to train_init_op\n",
    "            print(\"epoch=\", epoch, \"is done!\")\n",
    "            sess.run(eval_init_op)\n",
    "            _y_hats = []\n",
    "            for _ in range(num_eval_batches):\n",
    "                _y_hat, _word, _pred, _pron = sess.run([y_hat, word, pred, pron])\n",
    "                _y_hats.extend(_y_hat.tolist())\n",
    "                \n",
    "            # sample monitor\n",
    "            print(\"wrd:\", _word.decode(\"utf-8\"))\n",
    "            print(\"exp:\", _pron.decode(\"utf-8\"))\n",
    "            print(\"got:\", _pred.decode(\"utf-8\"))\n",
    "                \n",
    "            \n",
    "            _per = per(eval_prons, _y_hats)\n",
    "            print(\"per=%.2f\"%_per)\n",
    "            print()\n",
    "                  \n",
    "            sess.run(train_init_op)\n",
    "            \n",
    "            # save\n",
    "            if not os.path.exists(hp.logdir): os.makedirs(hp.logdir)\n",
    "            fname = os.path.join(hp.logdir, \"my_model_loss_%.2f_per_%.2f\" % (_loss, _per))\n",
    "            saver.save(sess, fname, global_step=_gs)\n",
    "   \n",
    "    print(\"Training Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "82t4Dmwp73--"
   },
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "test_batches, num_test_batches, num_test_samples  = get_batch(test_words, test_prons,\n",
    "                                                              hp.batch_size,\n",
    "                                                              shuffle=False)\n",
    "iter = tf.data.Iterator.from_structure(test_batches.output_types, test_batches.output_shapes)\n",
    "\n",
    "# create the initialisation operations\n",
    "test_init_op = iter.make_initializer(test_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inference graph is being built. Please be patient.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b109bf9c46024b1e81b4033a4c5ced26",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=20), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load model\n",
    "xs, ys = iter.get_next()\n",
    "net = Net(hp)\n",
    "y_hat, _, _, _ = net.eval(xs, ys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "log/04/my_model_loss_0.40_per_0.19-7720\n",
      "INFO:tensorflow:Restoring parameters from log/04/my_model_loss_0.40_per_0.19-7720\n",
      "checkpoint restored\n",
      "per=0.20\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# saver for restoration\n",
    "ckpt = tf.train.latest_checkpoint(hp.logdir)\n",
    "print(ckpt)\n",
    "# saver = tf.train.import_meta_graph(ckpt + \".meta\")# <- Do NOT use this as we'll use a distinct graph.\n",
    "saver = tf.train.Saver()\n",
    "   \n",
    "with tf.Session() as sess:\n",
    "    \n",
    "    saver.restore(sess, ckpt); print(\"checkpoint restored\") \n",
    "    sess.run(test_init_op)\n",
    "\n",
    "    _y_hats = []\n",
    "    for _ in range(num_test_batches):\n",
    "        _y_hat = sess.run(y_hat)\n",
    "        _y_hats.extend(_y_hat.tolist())\n",
    "            \n",
    "    _per = per(test_prons, _y_hats)\n",
    "            \n",
    "    print(\"per=%.2f\"%_per)\n",
    "    \n",
    "    # save\n",
    "    g2idx, idx2g, p2idx, idx2p = load_vocab()\n",
    "    \n",
    "    with open(\"result\", 'w') as fout:\n",
    "        fout.write(\"per: %.2f\\n\" % _per)\n",
    "        for w, r, h in zip(test_words, test_prons, _y_hats):\n",
    "            w = w.replace(\" \", \"\")\n",
    "            h = \" \".join(idx2p[idx] for idx in h)\n",
    "            h = h.split(\"</s>\")[0].strip()\n",
    "            fout.write(\"wrd: {}\\nexp: {}\\ngot: {}\\n\\n\".format(w, r, h))\n",
    "            \n",
    "    print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see some results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['wrd: campau',\n",
       " 'exp: K AA1 M P AW0',\n",
       " 'got: K AE1 M P OW2',\n",
       " '',\n",
       " 'wrd: tension',\n",
       " 'exp: T EH1 N SH AH0 N',\n",
       " 'got: T EH1 N S IY0 AH0 N',\n",
       " '',\n",
       " 'wrd: pithy',\n",
       " 'exp: P IH1 TH IY0',\n",
       " 'got: P IH1 TH IY0',\n",
       " '',\n",
       " 'wrd: blaisdell',\n",
       " 'exp: B L EY1 S D AH0 L',\n",
       " 'got: B L EY1 S D AH0 L',\n",
       " '',\n",
       " 'wrd: reflectone',\n",
       " 'exp: R IY0 F L EH1 K T OW2 N',\n",
       " 'got: R IY0 F L EH1 K T AH0 N',\n",
       " '',\n",
       " 'wrd: cherishing',\n",
       " 'exp: CH EH1 R IH0 SH IH0 NG',\n",
       " 'got: CH EH1 R IH0 SH IH0 NG',\n",
       " '',\n",
       " 'wrd: necessitate',\n",
       " 'exp: N AH0 S EH1 S AH0 T EY2 T',\n",
       " 'got: N EH2 S AH0 S EH1 T IH0 T',\n",
       " '',\n",
       " 'wrd: swiatkowski',\n",
       " 'exp: S V IY0 AH0 T K AO1 F S K IY0',\n",
       " 'got: S W IH0 T AO1 K S W IH0 K',\n",
       " '',\n",
       " 'wrd: tendons',\n",
       " 'exp: T EH1 N D AH0 N Z',\n",
       " 'got: T EH1 N D AH0 N Z',\n",
       " '',\n",
       " 'wrd: nucleonic',\n",
       " 'exp: N UW2 K L IY0 AA1 N IH0 K',\n",
       " 'got: N AH0 K L EH1 N IH0 K',\n",
       " '',\n",
       " 'wrd: nutone',\n",
       " 'exp: N UW1 T OW2 N',\n",
       " 'got: N UW1 T OW2 N',\n",
       " '',\n",
       " 'wrd: demaree',\n",
       " 'exp: D EH0 M ER0 IY1',\n",
       " 'got: D IH0 M AA1 R IY0',\n",
       " '',\n",
       " 'wrd: soltau',\n",
       " 'exp: S OW1 L T AW0',\n",
       " 'got: S OW1 L T OW0',\n",
       " '',\n",
       " 'wrd: methodically',\n",
       " 'exp: M AH0 TH AA1 D IH0 K AH0 L IY0',\n",
       " 'got: M EH2 TH AH0 D AA1 K AH0 L IY0',\n",
       " '',\n",
       " 'wrd: ahoskie',\n",
       " 'exp: AH0 HH AO1 S K IY0',\n",
       " 'got: AH0 HH AO1 S K IY0',\n",
       " '',\n",
       " 'wrd: mcivor',\n",
       " 'exp: M AH0 K IH1 V ER0',\n",
       " 'got: M AH0 K V EH1 R',\n",
       " '',\n",
       " 'wrd: generalissimo',\n",
       " 'exp: JH EH2 N EH0 R AH0 L IH1 S IH0 M OW2',\n",
       " 'got: JH EH2 N ER0 AH0 L IH1 Z AH0 M IY0',\n",
       " '',\n",
       " \"wrd: kasinga's\",\n",
       " 'exp: K AH0 S IH1 NG G AH0 Z',\n",
       " 'got: K AH0 S IH1 N JH AH0 Z',\n",
       " '',\n",
       " 'wrd: currin',\n",
       " 'exp: K AO1 R IH0 N',\n",
       " 'got: K ER1 IH0 N',\n",
       " '',\n",
       " 'wrd: deregulatory',\n",
       " 'exp: D IY0 R EH1 G Y AH0 L AH0 T AO2 R IY0',\n",
       " 'got: D EH2 R AH0 G L Y AA1 T ER0 IY0',\n",
       " '',\n",
       " 'wrd: calbos',\n",
       " 'exp: K AA1 L B OW0 S',\n",
       " 'got: K AE1 L B OW0 Z',\n",
       " '',\n",
       " 'wrd: kreg',\n",
       " 'exp: K R EH1 G',\n",
       " 'got: K R EH1 G',\n",
       " '',\n",
       " 'wrd: dezarn',\n",
       " 'exp: D EY0 Z AA1 R N',\n",
       " 'got: D IH0 Z AA1 R N',\n",
       " '',\n",
       " 'wrd: rapprochement',\n",
       " 'exp: R AE2 P R OW2 SH M AA1 N',\n",
       " 'got: R AE2 P R AH0 K M EY1 N T',\n",
       " '',\n",
       " 'wrd: rosenshine',\n",
       " 'exp: R OW1 Z AH0 N SH AY2 N',\n",
       " 'got: R OW0 Z EH1 N SH IY0 N',\n",
       " '']"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "open('result', 'r').read().splitlines()[-100:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Seq2seq tutorial with g2p.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}