rnn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy\n",
    "from tensorflow import keras\n",
    "\n",
    "_MODEL_DIRECTORY = '.'\n",
    "_BATCH_SIZE = 128\n",
    "_EMBEDDING_SIZE = 256\n",
    "_HIDDEN_LAYER_SIZE = 256\n",
    "_PROJECTION_LAYER_SIZE = 256\n",
    "_VALIDATION_SPLIT = 0.05\n",
    "_LOOKBACK = 100\n",
    "\n",
    "def create_model(input_length, input_vocab_size, output_vocab_size):\n",
    "    model = keras.models.Sequential()\n",
    "    model.add(keras.layers.Embedding(input_vocab_size, _EMBEDDING_SIZE, input_length=input_length))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.GRU(_HIDDEN_LAYER_SIZE, unroll=True))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.Dense(_PROJECTION_LAYER_SIZE, activation='linear'))\n",
    "    model.add(keras.layers.Dropout(0.5))\n",
    "    model.add(keras.layers.Dense(output_vocab_size, activation='softmax'))\n",
    "    model.compile(optimizer=keras.optimizers.SGD(momentum=0.9, nesterov=True, clipnorm=1.0),\n",
    "        loss='categorical_crossentropy',\n",
    "        metrics=['accuracy', 'top_k_categorical_accuracy'])\n",
    "    return model\n",
    "                            \n",
    "class Sequence(keras.utils.Sequence):\n",
    "    def __init__(self, X, Y, vocab_size):\n",
    "        self._X = X\n",
    "        self._Y = Y\n",
    "        self._vocab_size = vocab_size\n",
    "        self.on_epoch_end()\n",
    "    \n",
    "    def on_epoch_end(self):\n",
    "        randomize = numpy.arange(len(self._X))\n",
    "        numpy.random.shuffle(randomize)\n",
    "        self._X = self._X[randomize]\n",
    "        self._Y = self._Y[randomize]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        start = idx * _BATCH_SIZE\n",
    "        end = start + _BATCH_SIZE\n",
    "        X = []\n",
    "        for token_list in self._X[start:end]:\n",
    "          row = numpy.zeros(_LOOKBACK)\n",
    "          for idx, token in enumerate(token_list[-_LOOKBACK:]):\n",
    "            row[idx] = token\n",
    "          X.append(row)\n",
    "        return (\n",
    "            numpy.array(X),\n",
    "            keras.utils.to_categorical(\n",
    "                self._Y[start:end], num_classes=self._vocab_size))\n",
    "    \n",
    "    def __len__(self):\n",
    "        return int(numpy.floor(self._Y.size / _BATCH_SIZE))\n",
    "\n",
    "class RNN:\n",
    "    def __init__(self, vocabulary, word2idx, idx2word, lookback):\n",
    "        self._vocabulary = vocabulary\n",
    "        self._word2idx = word2idx\n",
    "        self._lookback = lookback\n",
    "        self._idx2word = idx2word\n",
    "        self._model = None\n",
    "\n",
    "    def fit(self, X, Y):\n",
    "        split = int(len(Y) * _VALIDATION_SPLIT)\n",
    "        X_train = X[split:]\n",
    "        Y_train = Y[split:]\n",
    "        X_validate = X[:split]\n",
    "        Y_validate = Y[:split]\n",
    "        input_vocab_size = len(self._vocabulary)\n",
    "        output_vocab_size = len(self._idx2word) + 100\n",
    "        self._model = create_model(self._lookback, input_vocab_size, output_vocab_size)\n",
    "        self._model.summary()\n",
    "        generator = Sequence(X_train, Y_train, output_vocab_size)\n",
    "        validation_generator = Sequence(X_validate, Y_validate, output_vocab_size)\n",
    "        self._model.fit_generator(generator, validation_data=validation_generator, verbose=2, epochs=1)\n",
    "  \n",
    "    def predict(self, x):\n",
    "        return self._model.predict(x)\n",
    "    \n",
    "    def save(self, output_directory):\n",
    "        self._model.save('%s/model.h5' % output_directory)\n",
    "        print('Loading training data...')\n",
    "\n",
    "def main():\n",
    "    with open(_MODEL_DIRECTORY + '/x.json', 'r') as file:\n",
    "        inputs = numpy.array(json.load(file))\n",
    "    with open(_MODEL_DIRECTORY + '/y.json', 'r') as file:\n",
    "        outputs = numpy.array(json.load(file))\n",
    "    with open(_MODEL_DIRECTORY + '/word2idx.json', 'r') as file:\n",
    "        word2idx = json.load(file)\n",
    "        vocabulary = set(word2idx.keys())\n",
    "    with open(_MODEL_DIRECTORY + '/idx2word.json', 'r') as file:\n",
    "        idx2word = json.load(file)\n",
    "    model = RNN(vocabulary, word2idx, idx2word, _LOOKBACK)\n",
    "    model.fit(inputs, outputs)\n",
    "    print('Saving model to %s...' % _MODEL_DIRECTORY)\n",
    "    model.save(_MODEL_DIRECTORY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture cap\n",
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cap.stdout"
   ]
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}