kanakanji_convert_4_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "L9Srn_cEx1fa"
   },
   "outputs": [],
   "source": [
    "# !pip install tf-nightly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BYk_Je9Be2sS"
   },
   "outputs": [],
   "source": [
    "# prepare dataset\n",
    "\n",
    "from utility import load_train_data\n",
    "from collections import Counter\n",
    "import os\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ct8_vGQFe2sX"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf version:  1.13.0-dev20190111\n"
     ]
    }
   ],
   "source": [
    "# import tensorflow\n",
    "\n",
    "import tensorflow as tf\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "print(\"tf version: \", tf.VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pcc8Q5KVZgLi"
   },
   "outputs": [],
   "source": [
    "# for reproducability\n",
    "\n",
    "import random as rn\n",
    "\n",
    "np.random.seed(42)\n",
    "rn.seed(12345)\n",
    "tf.set_random_seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TufQU1Nee2sa"
   },
   "outputs": [],
   "source": [
    "# prepare train data\n",
    "\n",
    "def parse_file(file):\n",
    "    for line in file:\n",
    "        line = line.rstrip('\\n')\n",
    "        sentence = line.split(' ')\n",
    "        yield sentence\n",
    "\n",
    "# TODO: current method does not allow the model to learn boundary beyond bigram.\n",
    "def adjust_size(sentences, sentence_size):\n",
    "    # Increment sentence size for shifting output later\n",
    "    sentence_size_plus = sentence_size + 1\n",
    "\n",
    "    for sentence in sentences:\n",
    "        # Insert BOS = Beginning Of Sentence\n",
    "        sentence.insert(0, '_BOS/_BOS')\n",
    "\n",
    "        # Split long sentence allowing overlap of 1 word\n",
    "        while len(sentence) >= sentence_size_plus:\n",
    "            yield sentence[:sentence_size_plus]\n",
    "            sentence = sentence[sentence_size:]\n",
    "\n",
    "        # Do not yield EOS-only sentence\n",
    "        if sentence:\n",
    "            # Insert EOS = End Of Sentence\n",
    "            sentence.append('_EOS/_EOS')\n",
    "\n",
    "            if len(sentence) < sentence_size_plus:\n",
    "                # Padding sentence to make its size sentence_size_plus\n",
    "                sentence += ['_PAD/_PAD'] * (sentence_size_plus - len(sentence))\n",
    "            yield sentence\n",
    "        \n",
    "\n",
    "def create_vocabulary(sentences, vocabulary_size):\n",
    "    # Create list of words indexed by word ID\n",
    "    counter = Counter(word for words in sentences for word in words)\n",
    "    most_common = counter.most_common(vocabulary_size - 1)\n",
    "    vocabulary = [word for word, count in most_common]\n",
    "    vocabulary.insert(0, '_UNK/_UNK')\n",
    "    return vocabulary\n",
    "\n",
    "\n",
    "def convert_to_ids(sentences, vocabulary):\n",
    "    dictionary = dict((word, word_id) for word_id, word in enumerate(vocabulary))\n",
    "\n",
    "    for sentence in sentences:\n",
    "        word_ids = []\n",
    "\n",
    "        for word in sentence:\n",
    "            if word in dictionary:\n",
    "                word_id = dictionary[word]\n",
    "            else:\n",
    "                word_id = dictionary['_UNK/_UNK']\n",
    "            word_ids.append(word_id)\n",
    "\n",
    "        yield word_ids\n",
    "\n",
    "        \n",
    "# TODO: current batching ignores sentences that does't fit into last batch.\n",
    "def create_batches(sentences, batch_size):\n",
    "    all_batches = int(len(sentences) / batch_size)\n",
    "\n",
    "    for i in range(all_batches):\n",
    "        batch_sentences = sentences[i * batch_size:(i + 1) * batch_size]\n",
    "        batch_input = []\n",
    "        batch_output = []\n",
    "\n",
    "        for sentence in batch_sentences:\n",
    "            # Shift sentence by 1 time step\n",
    "            input_ = sentence[:-1]\n",
    "            output_ = sentence[1:]\n",
    "\n",
    "            batch_input.append(input_)\n",
    "            batch_output.append(output_)\n",
    "\n",
    "        yield batch_input, batch_output    \n",
    "\n",
    "        \n",
    "def create_pair(sentences):\n",
    "    \n",
    "    print(\"sentences count: \", len(sentences))\n",
    "#     print(sentences)\n",
    "    \n",
    "    input_list = []\n",
    "    output_list = []\n",
    "    \n",
    "    i=0\n",
    "    for sentence in sentences:\n",
    "#         print(\"aaa i: {}\", i)\n",
    "        input_ = sentence[:-1]\n",
    "        output_ = sentence[1:]\n",
    "        \n",
    "        input_list.append(input_)\n",
    "        output_list.append(output_)\n",
    "        \n",
    "        i = i + 1\n",
    "        \n",
    "    return input_list, output_list\n",
    "        \n",
    "def save_metadata(model_directory, vocabulary):\n",
    "    # Create directory if not exists\n",
    "    if not os.path.exists(model_directory):\n",
    "        os.makedirs(model_directory)\n",
    "\n",
    "#     # Save settings\n",
    "#     settings_path = os.path.join(model_directory, 'settings.json')\n",
    "#     with open(settings_path, 'w') as settings_file:\n",
    "#         json.dump(vars(args), settings_file, indent=4)\n",
    "\n",
    "    # Save vocabulary\n",
    "    vocabulary_path = os.path.join(model_directory, 'vocabulary.txt')\n",
    "    with open(vocabulary_path, 'w') as vocabulary_file:\n",
    "        vocabulary_file.write('\\n'.join(vocabulary))\n",
    "        \n",
    "        \n",
    "def load_train_data(dataset_name, sentence_size, vocabulary_size, batch_size, model_directory ):\n",
    "\n",
    "    sentences = parse_file(open(dataset_name))\n",
    "    sentences = list(adjust_size(sentences, sentence_size))\n",
    "    vocabulary = create_vocabulary(sentences, vocabulary_size)\n",
    "    sentences = list(convert_to_ids(sentences, vocabulary))\n",
    "#     train_data = list(create_batches(sentences, batch_size))\n",
    "    save_metadata(model_directory, vocabulary)\n",
    "    \n",
    "    # target\n",
    "    input_x, target_y = create_pair(sentences)\n",
    "    \n",
    "    return input_x, target_y\n",
    "#     return train_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "ZZo5EHvme2sd",
    "outputId": "9f6c35b0-2191-4038-91db-eea77a298883"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sentences count:  80103\n"
     ]
    }
   ],
   "source": [
    "# Load and preprocess training data\n",
    "\n",
    "dataset_name = \"wiki_dataset_mecab_80000.txt\"\n",
    "sentence_size = 30\n",
    "BATCH_SIZE = 64\n",
    "batch_size = BATCH_SIZE\n",
    "vocabulary_size = 50000\n",
    "model_directory = \"models\"\n",
    "\n",
    "input_x, target_y = load_train_data(dataset_name, sentence_size, vocabulary_size, BATCH_SIZE, model_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 89
    },
    "colab_type": "code",
    "id": "NhTydzl-e2sj",
    "outputId": "33de1813-31c2-4a25-b62f-9f4df12f5af3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_x, count: {} 80103\n",
      "target_y, count: {} 80103\n",
      "input_x.0:  [3, 37098, 12, 9, 7833, 16536, 12, 6, 131, 16, 1471, 10, 19, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "target_y.0:  [37098, 12, 9, 7833, 16536, 12, 6, 131, 16, 1471, 10, 19, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n"
     ]
    }
   ],
   "source": [
    "# check data\n",
    "\n",
    "print(\"input_x, count: {}\", len(input_x))\n",
    "print(\"target_y, count: {}\", len(target_y))\n",
    "print(\"input_x.0: \", input_x[0])\n",
    "print(\"target_y.0: \", target_y[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0VWeO2Fwe2so"
   },
   "outputs": [],
   "source": [
    "# parepare dataset\n",
    "\n",
    "# BUFFER_SIZE = len(input_x)\n",
    "BUFFER_SIZE = 50000\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((input_x, target_y)).shuffle(BUFFER_SIZE)\n",
    "dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 125
    },
    "colab_type": "code",
    "id": "GbrQbX3ne2ss",
    "outputId": "fff96071-b819-4e84-d6f2-d72b1175a937"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/kouohhashi/anaconda3/envs/vui2/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py:532: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "input_x_0: [   3   48   42    5 2936   22 9774   16   33 1561    9 1040   11   24\n",
      "   27    4    1    1    1    1    1    1    1    1    1    1    1    1\n",
      "    1    1] (30,)\n",
      "output_y_0: [  48   42    5 2936   22 9774   16   33 1561    9 1040   11   24   27\n",
      "    4    1    1    1    1    1    1    1    1    1    1    1    1    1\n",
      "    1    1] (30,)\n"
     ]
    }
   ],
   "source": [
    "# check dataset again\n",
    "\n",
    "for input_example, target_example in  dataset.take(1):\n",
    "    input_x_0 = input_example.numpy()[0]\n",
    "    print(\"input_x_0:\", input_x_0, input_x_0.shape)\n",
    "    output_y_0 = target_example.numpy()[0]\n",
    "    print(\"output_y_0:\", output_y_0, output_y_0.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "04AyUiMK3EPr"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "O-oglZ4Ka5vN"
   },
   "outputs": [],
   "source": [
    "# prepare model\n",
    "\n",
    "def gru(units, backword_flg):\n",
    "    # should use GRU anyway....\n",
    "    return tf.keras.layers.GRU(units, \n",
    "                               return_sequences=True, \n",
    "                               return_state=True, \n",
    "                               recurrent_activation='sigmoid', \n",
    "                               recurrent_initializer='glorot_uniform',\n",
    "                               go_backwards=backword_flg)\n",
    "    \n",
    "    \n",
    "class KanaKanjiModel(tf.keras.Model):\n",
    "    def __init__(self, vocab_size, embedding_dim, rnn_units, batch_size):\n",
    "        super(KanaKanjiModel, self).__init__()\n",
    "        self.batch_size = batch_size\n",
    "        self.rnn_units = rnn_units\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim, input_shape=(30,))\n",
    "        self.gru_f = gru(self.rnn_units, False)\n",
    "        self.gru_b = gru(self.rnn_units, True)\n",
    "        self.vocab_size = vocab_size\n",
    "        self.fc = tf.keras.layers.Dense(vocab_size)\n",
    "        self.dropout = tf.keras.layers.Dropout(0.5)\n",
    "        \n",
    "        \n",
    "    def call(self, x, hidden, training):\n",
    "        \n",
    "        x = self.embedding(x)\n",
    "        output, state = self.gru_f(x, initial_state = hidden)\n",
    "        output = self.dropout(output, training=training)\n",
    "        output = self.fc(output)\n",
    "\n",
    "        return output, state\n",
    "    \n",
    "    def initialize_hidden_state(self):\n",
    "        return tf.zeros((self.batch_size, self.rnn_units))\n",
    "      \n",
    "      \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        # You need to override this function if you want to use the subclassed model\n",
    "        # as part of a functional-style model.\n",
    "        # Otherwise, this method is optional.\n",
    "        \n",
    "        # maybe we don not need this function...\n",
    "        \n",
    "        shape = tf.TensorShape(input_shape).as_list()\n",
    "        shape[-1] = self.num_classes\n",
    "        \n",
    "        return tf.TensorShape([tf.TensorShape([-1, 30, 50000]), tf.TensorShape([-1, 30, 400])])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AOE5rPcje2sy"
   },
   "outputs": [],
   "source": [
    "# create model\n",
    "\n",
    "hidden_size = 400\n",
    "embedding_dim = hidden_size\n",
    "\n",
    "model = KanaKanjiModel(vocabulary_size, embedding_dim, hidden_size, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eXmv4IuH3_yx"
   },
   "outputs": [],
   "source": [
    "# model.summary() is not working..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Pl15bZZ3e2s5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# speed up\n",
    "model.call = tf.contrib.eager.defun(model.call)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q0IFgafQe2s8"
   },
   "outputs": [],
   "source": [
    "optimizer3 = tf.train.AdamOptimizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yxDg8eSbKUhK"
   },
   "outputs": [],
   "source": [
    "# check points\n",
    "\n",
    "checkpoint_dir = './ck_20190222'\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
    "checkpoint = tf.train.Checkpoint(optimizer=optimizer3,\n",
    "                                 model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 737
    },
    "colab_type": "code",
    "id": "MKhZJ9w4e2s_",
    "outputId": "30262be4-f740-4ab5-b565-7b50e3137445"
   },
   "outputs": [],
   "source": [
    "# training\n",
    "\n",
    "import time\n",
    "\n",
    "# Training step\n",
    "EPOCHS = 10\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    start = time.time()\n",
    "    \n",
    "    # initializing the hidden state at the start of every epoch\n",
    "    # initally hidden is None\n",
    "    \n",
    "    hidden_f = model.initialize_hidden_state()\n",
    "    # hidden_f = tf.zeros((BATCH_SIZE, hidden_size))\n",
    "    \n",
    "    hidden_b = hidden_f\n",
    "    \n",
    "    loss2 = 0\n",
    "    \n",
    "    for (batch_n, (inp, target)) in enumerate(dataset):\n",
    "        \n",
    "        loss1 = 0\n",
    "        \n",
    "        with tf.GradientTape() as tape:\n",
    "            # feeding the hidden state back into the model\n",
    "            # This is the interesting step\n",
    "            \n",
    "            predictions, hidden_f  = model(inp, hidden_f, True)\n",
    "            \n",
    "            target = tf.expand_dims(target, 2)\n",
    "            loss = tf.losses.sparse_softmax_cross_entropy(target, predictions)\n",
    "            loss1 = tf.reduce_mean(loss)\n",
    "\n",
    "        loss1_np = float(loss1.numpy())\n",
    "        batch_loss = (loss1_np / int(inp.shape[1]))\n",
    "        \n",
    "        grads = tape.gradient(loss1, model.variables)\n",
    "        optimizer3.apply_gradients(zip(grads, model.variables))\n",
    "        \n",
    "        loss2 = loss2 + batch_loss\n",
    "        \n",
    "        if batch_n % 1000 == 0:\n",
    "            template = 'Epoch {} Batch {} Loss {:.4f}'\n",
    "            print(template.format(epoch+1, batch_n, loss1))\n",
    "    \n",
    "    # saving (checkpoint) the model every 2 epochs\n",
    "    if (epoch + 1) % 2 == 0:\n",
    "        checkpoint.save(file_prefix = checkpoint_prefix)\n",
    "\n",
    "    print ('Epoch {} Loss {:.4f}'.format(epoch+1, loss2))\n",
    "    print ('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "RMLWUellCHNA",
    "outputId": "e817e2a6-a6fa-45cb-ca2f-655bfd345dd6"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./ck_20190209/ckpt-5'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# save weight\n",
    "\n",
    "checkpoint.save(file_prefix = checkpoint_prefix)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0k-6sTeCHVP1"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "r6OrRqUVe2tf"
   },
   "outputs": [],
   "source": [
    "# rebuild moel\n",
    "hidden_size = 400\n",
    "embedding_dim = hidden_size\n",
    "vocabulary_size = 50000\n",
    "\n",
    "model = KanaKanjiModel(vocabulary_size, embedding_dim, hidden_size, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "zG6ZrfTDER6K",
    "outputId": "58e1ac4d-91cb-4837-9eb2-c3ce108f4849"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./ck_20190209/ckpt-5'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load model\n",
    "# check latest checkpoints\n",
    "\n",
    "checkpoint_dir3 = './ck_20190222'\n",
    "tf.train.latest_checkpoint(checkpoint_dir3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "pLgiLjcogwwp",
    "outputId": "b09213f0-a16d-4929-d9c8-20af9ad362a9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "status:  <tensorflow.python.training.checkpointable.util.CheckpointLoadStatus object at 0x1288e8128>\n"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "# load weight\n",
    "\n",
    "checkpoint_dir = './ck_20190222'\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
    "checkpoint = tf.train.Checkpoint(optimizer=optimizer3,\n",
    "                                 model=model)\n",
    "\n",
    "status = checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))\n",
    "print(\"status: \", status)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GIxSTmtwN2cc"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SNotJvQHe2uJ"
   },
   "outputs": [],
   "source": [
    "# this is based on decode.py of the original scripts\n",
    "\n",
    "import collections\n",
    "import heapq\n",
    "import operator\n",
    "\n",
    "hiragana_added = []\n",
    "\n",
    "def load_dictionary(model_directory):\n",
    "    vocabulary_path = os.path.join(model_directory, 'vocabulary.txt')\n",
    "    vocabulary = []\n",
    "    for line in open(vocabulary_path):\n",
    "        line = line.rstrip('\\n')\n",
    "        target, source = line.split('/', 1)\n",
    "        vocabulary.append((target, source))\n",
    "\n",
    "    dictionary = collections.defaultdict(list)\n",
    "    for i, (target, source) in enumerate(vocabulary):\n",
    "        dictionary[source].append((target, i))\n",
    "\n",
    "    return dictionary\n",
    "\n",
    "def create_lattice(input_, dictionary):\n",
    "    lattice = [[[] for _ in range(len(input_) + 1)] for _ in range(len(input_) + 2)]\n",
    "    _, unk_id = dictionary['_UNK'][0]\n",
    "\n",
    "    for i in range(1, len(input_) + 1):\n",
    "        for j in range(i):\n",
    "            key = input_[j:i]\n",
    "            if key in dictionary:\n",
    "                for target, word_id in dictionary[key]:\n",
    "                    lattice[i][j].append((target, word_id))\n",
    "            elif len(key) == 1:\n",
    "                # Create _UNK node with verbatim target when single character key is not found in the dictionary.\n",
    "                lattice[i][j].append((key, unk_id))\n",
    "\n",
    "    _, eos_id = dictionary['_EOS'][0]\n",
    "    lattice[-1][-1].append(('', eos_id))\n",
    "    return lattice\n",
    "\n",
    "\n",
    "def initialize_queues(lattice, rnn_predictor, dictionary):\n",
    "    # Initialize priority queues for keeping hypotheses\n",
    "    # A hypothesis is a tuple of (cost, string, state, prediction)\n",
    "    # cost is total negative log probability\n",
    "    # state.shape == [hidden_size * layer_size]\n",
    "    # prediction.shape == [vocabulary_size]\n",
    "    \n",
    "    hiragana_added = []\n",
    "    \n",
    "    hidden_f = rnn_predictor.initialize_hidden_state()\n",
    "    hidden_b = hidden_f\n",
    "    \n",
    "    _, bos_id = dictionary['_BOS'][0]\n",
    "    \n",
    "    input_x0 = tf.expand_dims([bos_id], 0)\n",
    "    \n",
    "    bos_predictions, hidden_f = rnn_predictor(input_x0, hidden_f, False)\n",
    "    \n",
    "    \n",
    "    # logits to probability\n",
    "    bos_predictions = tf.squeeze(bos_predictions, 0)\n",
    "    bos_predictions = tf.squeeze(bos_predictions, 0)\n",
    "    bos_predictions = -1 * tf.nn.log_softmax(bos_predictions, axis=0)\n",
    "    \n",
    "    hidden_f = tf.expand_dims(hidden_f, 0)\n",
    "    \n",
    "    bos_hypothesis = (0.0, '', hidden_f[0], bos_predictions)\n",
    "    queues = [[] for _ in range(len(lattice))]\n",
    "    queues[0].append(bos_hypothesis)\n",
    "    return queues\n",
    "\n",
    "def search(lattice, queues, rnn_predictor, beam_size, viterbi_size):\n",
    "    # Breadth first search with beam pruning and viterbi-like pruning\n",
    "    for i in range(len(lattice)):\n",
    "        queue = []\n",
    "\n",
    "        # create hypotheses without predicting next word\n",
    "        for j in range(len(lattice[i])):\n",
    "            for target, word_id in lattice[i][j]:\n",
    "                \n",
    "                # if word_id == 350:\n",
    "                #    print(\"word_id: \", word_id)\n",
    "                \n",
    "                word_queue = []\n",
    "                for previous_cost, previous_string, previous_state_f, previous_prediction in queues[j]:\n",
    "                    # if logits is bigger, better.\n",
    "                    \n",
    "                    # seems to need give huge priority to first word\n",
    "                    cost = previous_cost + previous_prediction[word_id]\n",
    "                    \n",
    "                    string = previous_string + target\n",
    "                    hypothesis = (cost, string, word_id, previous_state_f)\n",
    "                    word_queue.append(hypothesis)\n",
    "                \n",
    "                # prune word_queue to viterbi size\n",
    "                if viterbi_size > 0:\n",
    "                    word_queue = heapq.nsmallest(viterbi_size, word_queue, key=operator.itemgetter(0))\n",
    "                    \n",
    "                queue += word_queue\n",
    "                \n",
    "        # prune queue to beam size\n",
    "        if beam_size > 0:\n",
    "            queue = heapq.nsmallest(beam_size, queue, key=operator.itemgetter(0))\n",
    "        \n",
    "        # predict next word and state before continue\n",
    "        for cost, string, word_id, previous_state_f in queue:\n",
    "            \n",
    "            input_x0 = tf.expand_dims([word_id], 0)\n",
    "              \n",
    "            predictions, state_f = rnn_predictor(input_x0, [previous_state_f], False)\n",
    "            \n",
    "            # logits to probability\n",
    "            predictions = tf.squeeze(predictions, 0)\n",
    "            predictions = tf.squeeze(predictions, 0)\n",
    "            predictions = -1 * tf.nn.log_softmax(predictions, axis=0)\n",
    "            # print(\"predictions.shape: \", predictions.shape)\n",
    "        \n",
    "            state_f = tf.expand_dims(state_f, 0)\n",
    "            hypothesis = (cost, string, state_f[0], predictions)\n",
    "            queues[i].append(hypothesis)\n",
    "\n",
    "    return queues\n",
    "\n",
    "def decode(source, dictionary, rnn_predictor, beam_size, viterbi_size):\n",
    "    lattice = create_lattice(source, dictionary)\n",
    "    queues = initialize_queues(lattice, rnn_predictor, dictionary)\n",
    "    queues = search(lattice, queues, rnn_predictor, beam_size, viterbi_size)\n",
    "\n",
    "    candidates = []\n",
    "    for cost, string, _, _ in queues[-1]:\n",
    "        candidates.append((string, cost))\n",
    "\n",
    "    top_result = candidates[0][0]\n",
    "    return top_result, candidates, lattice, queues\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QbKKlSO-e2uL"
   },
   "outputs": [],
   "source": [
    "def convert_kana_to_kanji(line, rnn_predictor):\n",
    "    \n",
    "    # Load settings and vocabulary\n",
    "    model_directory = \"models\"\n",
    "    dictionary = load_dictionary(model_directory)\n",
    "    \n",
    "    # debug\n",
    "    print_nbest = True \n",
    "    print_lattice = False\n",
    "    print_queue = False\n",
    "    \n",
    "    # parameters\n",
    "    beam_size = 5\n",
    "    viterbi_size = 50000\n",
    "    \n",
    "    \n",
    "    # Iterate input file line by line\n",
    "    line = line.rstrip('\\n')\n",
    "\n",
    "    # Decode - this might take ~10 seconds per line\n",
    "    result, candidates, lattice, queues = decode(line, dictionary, rnn_predictor, beam_size, viterbi_size)\n",
    "\n",
    "    # Print decoded results\n",
    "    if not print_nbest:\n",
    "        print(result)\n",
    "    else:\n",
    "        for string, cost in candidates:\n",
    "            print(string, cost)\n",
    "\n",
    "    # Print lattice for debug\n",
    "    if print_lattice:\n",
    "        for i in range(len(lattice)):\n",
    "            for j in range(len(lattice[i])):\n",
    "                print('i = {}, j = {}'.format(i, j))\n",
    "                for target, word_id in lattice[i][j]:\n",
    "                    print(target, word_id)\n",
    "\n",
    "    # Print queues for debug\n",
    "    if print_queue:\n",
    "        for i, queue in enumerate(queues):\n",
    "            print('queue', i)\n",
    "            for cost, string, state_f, prediction in queue:\n",
    "                # print(string, cost)\n",
    "                print(string, cost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 2591
    },
    "colab_type": "code",
    "id": "5oopc1Gde2uN",
    "outputId": "a6ce6042-c11c-402f-8467-d1ca4e86b30e"
   },
   "outputs": [],
   "source": [
    "# test the model\n",
    "\n",
    "convert_kana_to_kanji(\"ぱりのれきし\", model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4xTPs1dtBhLN"
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "kanakanji_convert_3_train.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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}