vectorizer/11.elmo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "import tensorflow as tf\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['negative', 'positive']\n",
      "10662\n",
      "10662\n"
     ]
    }
   ],
   "source": [
    "trainset = sklearn.datasets.load_files(container_path = 'data', encoding = 'UTF-8')\n",
    "trainset.data, trainset.target = separate_dataset(trainset,1.0)\n",
    "print(trainset.target_names)\n",
    "print(len(trainset.data))\n",
    "print(len(trainset.target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "def build_dataset(words, n_words, atleast=1):\n",
    "    count = [['PAD', 0], ['GO', 1], ['EOS', 2], ['UNK', 3]]\n",
    "    counter = collections.Counter(words).most_common(n_words)\n",
    "    counter = [i for i in counter if i[1] >= atleast]\n",
    "    count.extend(counter)\n",
    "    dictionary = dict()\n",
    "    for word, _ in count:\n",
    "        dictionary[word] = len(dictionary)\n",
    "    data = list()\n",
    "    for word in words:\n",
    "        index = dictionary.get(word, 3)\n",
    "        data.append(index)\n",
    "    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))\n",
    "    return data, dictionary, reversed_dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = (' '.join(trainset.data)).split()\n",
    "vocabulary_size = len(list(set(split)))\n",
    "data, dictionary, rev_dictionary = build_dataset(split, vocabulary_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20336"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dictionary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class Vocabulary:\n",
    "    def __init__(self, dictionary, rev_dictionary):\n",
    "        self._dictionary = dictionary\n",
    "        self._rev_dictionary = rev_dictionary\n",
    "\n",
    "    @property\n",
    "    def start_string(self):\n",
    "        return self._dictionary['GO']\n",
    "\n",
    "    @property\n",
    "    def end_string(self):\n",
    "        return self._dictionary['EOS']\n",
    "\n",
    "    @property\n",
    "    def unk(self):\n",
    "        return self._dictionary['UNK']\n",
    "\n",
    "    @property\n",
    "    def size(self):\n",
    "        return len(self._dictionary)\n",
    "\n",
    "    def word_to_id(self, word):\n",
    "        return self._dictionary.get(word, self.unk)\n",
    "\n",
    "    def id_to_word(self, cur_id):\n",
    "        return self._rev_dictionary.get(cur_id, self._rev_dictionary[3])\n",
    "\n",
    "    def decode(self, cur_ids):\n",
    "        return ' '.join([self.id_to_word(cur_id) for cur_id in cur_ids])\n",
    "\n",
    "    def encode(self, sentence, reverse = False, split = True):\n",
    "\n",
    "        if split:\n",
    "            sentence = sentence.split()\n",
    "        word_ids = [self.word_to_id(cur_word) for cur_word in sentence]\n",
    "\n",
    "        if reverse:\n",
    "            return np.array(\n",
    "                [self.end_string] + word_ids + [self.start_string],\n",
    "                dtype = np.int32,\n",
    "            )\n",
    "        else:\n",
    "            return np.array(\n",
    "                [self.start_string] + word_ids + [self.end_string],\n",
    "                dtype = np.int32,\n",
    "            )\n",
    "\n",
    "\n",
    "class UnicodeCharsVocabulary(Vocabulary):\n",
    "    def __init__(self, dictionary, rev_dictionary, max_word_length, **kwargs):\n",
    "        super(UnicodeCharsVocabulary, self).__init__(\n",
    "            dictionary, rev_dictionary, **kwargs\n",
    "        )\n",
    "        self._max_word_length = max_word_length\n",
    "        self.bos_char = 256\n",
    "        self.eos_char = 257\n",
    "        self.bow_char = 258\n",
    "        self.eow_char = 259\n",
    "        self.pad_char = 260\n",
    "        num_words = self.size\n",
    "\n",
    "        self._word_char_ids = np.zeros(\n",
    "            [num_words, max_word_length], dtype = np.int32\n",
    "        )\n",
    "\n",
    "        def _make_bos_eos(c):\n",
    "            r = np.zeros([self._max_word_length], dtype = np.int32)\n",
    "            r[:] = self.pad_char\n",
    "            r[0] = self.bow_char\n",
    "            r[1] = c\n",
    "            r[2] = self.eow_char\n",
    "            return r\n",
    "\n",
    "        self.bos_chars = _make_bos_eos(self.bos_char)\n",
    "        self.eos_chars = _make_bos_eos(self.eos_char)\n",
    "\n",
    "        for i, word in enumerate(self._dictionary.keys()):\n",
    "            self._word_char_ids[i] = self._convert_word_to_char_ids(word)\n",
    "\n",
    "        self._word_char_ids[self.start_string] = self.bos_chars\n",
    "        self._word_char_ids[self.end_string] = self.eos_chars\n",
    "\n",
    "    @property\n",
    "    def word_char_ids(self):\n",
    "        return self._word_char_ids\n",
    "\n",
    "    @property\n",
    "    def max_word_length(self):\n",
    "        return self._max_word_length\n",
    "\n",
    "    def _convert_word_to_char_ids(self, word):\n",
    "        code = np.zeros([self.max_word_length], dtype = np.int32)\n",
    "        code[:] = self.pad_char\n",
    "        word_encoded = word.encode('utf-8', 'ignore')[\n",
    "            : (self.max_word_length - 2)\n",
    "        ]\n",
    "        code[0] = self.bow_char\n",
    "        for k, chr_id in enumerate(word_encoded, start = 1):\n",
    "            code[k] = chr_id\n",
    "\n",
    "        code[len(word_encoded) + 1] = self.eow_char\n",
    "        return code\n",
    "\n",
    "    def word_to_char_ids(self, word):\n",
    "        if word in self._dictionary:\n",
    "            return self._word_char_ids[self._dictionary[word]]\n",
    "        else:\n",
    "            return self._convert_word_to_char_ids(word)\n",
    "\n",
    "    def encode_chars(self, sentence, reverse = False, split = True):\n",
    "        if split:\n",
    "            sentence = sentence.split()\n",
    "        chars_ids = [self.word_to_char_ids(cur_word) for cur_word in sentence]\n",
    "\n",
    "        if reverse:\n",
    "            return np.vstack([self.eos_chars] + chars_ids + [self.bos_chars])\n",
    "        else:\n",
    "            return np.vstack([self.bos_chars] + chars_ids + [self.eos_chars])\n",
    "\n",
    "\n",
    "def _get_batch(generator, batch_size, num_steps, max_word_length):\n",
    "    cur_stream = [None] * batch_size\n",
    "\n",
    "    no_more_data = False\n",
    "    while True:\n",
    "        inputs = np.zeros([batch_size, num_steps], np.int32)\n",
    "        if max_word_length is not None:\n",
    "            char_inputs = np.zeros(\n",
    "                [batch_size, num_steps, max_word_length], np.int32\n",
    "            )\n",
    "        else:\n",
    "            char_inputs = None\n",
    "        targets = np.zeros([batch_size, num_steps], np.int32)\n",
    "        for i in range(batch_size):\n",
    "            cur_pos = 0\n",
    "\n",
    "            while cur_pos < num_steps:\n",
    "                if cur_stream[i] is None or len(cur_stream[i][0]) <= 1:\n",
    "                    try:\n",
    "                        cur_stream[i] = list(next(generator))\n",
    "                    except StopIteration:\n",
    "                        no_more_data = True\n",
    "                        break\n",
    "                how_many = min(len(cur_stream[i][0]) - 1, num_steps - cur_pos)\n",
    "                next_pos = cur_pos + how_many\n",
    "\n",
    "                inputs[i, cur_pos:next_pos] = cur_stream[i][0][:how_many]\n",
    "                if max_word_length is not None:\n",
    "                    char_inputs[i, cur_pos:next_pos] = cur_stream[i][1][\n",
    "                        :how_many\n",
    "                    ]\n",
    "                targets[i, cur_pos:next_pos] = cur_stream[i][0][\n",
    "                    1 : how_many + 1\n",
    "                ]\n",
    "\n",
    "                cur_pos = next_pos\n",
    "\n",
    "                cur_stream[i][0] = cur_stream[i][0][how_many:]\n",
    "                if max_word_length is not None:\n",
    "                    cur_stream[i][1] = cur_stream[i][1][how_many:]\n",
    "\n",
    "        if no_more_data:\n",
    "            break\n",
    "\n",
    "        X = {\n",
    "            'token_ids': inputs,\n",
    "            'tokens_characters': char_inputs,\n",
    "            'next_token_id': targets,\n",
    "        }\n",
    "\n",
    "        yield X\n",
    "\n",
    "\n",
    "class LMDataset:\n",
    "    def __init__(self, string, vocab, reverse = False):\n",
    "        self._vocab = vocab\n",
    "        self._string = string\n",
    "        self._reverse = reverse\n",
    "        self._use_char_inputs = hasattr(vocab, 'encode_chars')\n",
    "        self._i = 0\n",
    "        self._nids = len(self._string)\n",
    "\n",
    "    def _load_string(self, string):\n",
    "        if self._reverse:\n",
    "            string = string.split()\n",
    "            string.reverse()\n",
    "            string = ' '.join(string)\n",
    "        \n",
    "        ids = self._vocab.encode(string, self._reverse)\n",
    "\n",
    "        if self._use_char_inputs:\n",
    "            chars_ids = self._vocab.encode_chars(string, self._reverse)\n",
    "        else:\n",
    "            chars_ids = None\n",
    "\n",
    "        return list(zip([ids], [chars_ids]))[0]\n",
    "\n",
    "    def get_sentence(self):\n",
    "        while True:\n",
    "            if self._i == self._nids:\n",
    "                self._i = 0\n",
    "            ret = self._load_string(self._string[self._i])\n",
    "            self._i += 1\n",
    "            yield ret\n",
    "\n",
    "    @property\n",
    "    def max_word_length(self):\n",
    "        if self._use_char_inputs:\n",
    "            return self._vocab.max_word_length\n",
    "        else:\n",
    "            return None\n",
    "\n",
    "    def iter_batches(self, batch_size, num_steps):\n",
    "        for X in _get_batch(\n",
    "            self.get_sentence(), batch_size, num_steps, self.max_word_length\n",
    "        ):\n",
    "            yield X\n",
    "\n",
    "    @property\n",
    "    def vocab(self):\n",
    "        return self._vocab\n",
    "\n",
    "\n",
    "class BidirectionalLMDataset:\n",
    "    def __init__(self, string, vocab):\n",
    "        self._data_forward = LMDataset(string, vocab, reverse = False)\n",
    "        self._data_reverse = LMDataset(string, vocab, reverse = True)\n",
    "\n",
    "    def iter_batches(self, batch_size, num_steps):\n",
    "        max_word_length = self._data_forward.max_word_length\n",
    "\n",
    "        for X, Xr in zip(\n",
    "            _get_batch(\n",
    "                self._data_forward.get_sentence(),\n",
    "                batch_size,\n",
    "                num_steps,\n",
    "                max_word_length,\n",
    "            ),\n",
    "            _get_batch(\n",
    "                self._data_reverse.get_sentence(),\n",
    "                batch_size,\n",
    "                num_steps,\n",
    "                max_word_length,\n",
    "            ),\n",
    "        ):\n",
    "\n",
    "            for k, v in Xr.items():\n",
    "                X[k + '_reverse'] = v\n",
    "\n",
    "            yield X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni = UnicodeCharsVocabulary(dictionary, rev_dictionary, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([258, 116, 101, 115, 116,  51, 259, 260, 260, 260, 260, 260, 260,\n",
       "       260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260,\n",
       "       260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260,\n",
       "       260, 260, 260, 260, 260, 260, 260, 260, 260, 260, 260], dtype=int32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "uni.word_to_char_ids('test3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "bi = BidirectionalLMDataset(trainset.data, uni)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "n_train_tokens = len(dictionary)\n",
    "options = {\n",
    "    'bidirectional': True,\n",
    "    'char_cnn': {\n",
    "        'activation': 'relu',\n",
    "        'embedding': {'dim': 128},\n",
    "        'filters': [\n",
    "            [1, 32],\n",
    "            [2, 32],\n",
    "            [3, 64],\n",
    "            [4, 128],\n",
    "            [5, 256],\n",
    "            [6, 512],\n",
    "            [7, 1024],\n",
    "        ],\n",
    "        'max_characters_per_token': 50,\n",
    "        'n_characters': 261,\n",
    "        'n_highway': 2,\n",
    "    },\n",
    "    'dropout': 0.1,\n",
    "    'lstm': {\n",
    "        'cell_clip': 3,\n",
    "        'dim': 512,\n",
    "        'n_layers': 2,\n",
    "        'projection_dim': 256,\n",
    "        'proj_clip': 3,\n",
    "        'use_skip_connections': True,\n",
    "    },\n",
    "    'n_epochs': 100,\n",
    "    'n_train_tokens': n_train_tokens,\n",
    "    'batch_size': batch_size,\n",
    "    'n_tokens_vocab': uni.size,\n",
    "    'unroll_steps': 20,\n",
    "    'n_negative_samples_batch': 0.001,\n",
    "    'sample_softmax': True,\n",
    "    'share_embedding_softmax': False,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LanguageModel:\n",
    "    def __init__(self, options, is_training):\n",
    "        self.options = options\n",
    "        self.is_training = is_training\n",
    "        self.bidirectional = options.get('bidirectional', False)\n",
    "\n",
    "        self.char_inputs = 'char_cnn' in self.options\n",
    "\n",
    "        self.share_embedding_softmax = options.get(\n",
    "            'share_embedding_softmax', False\n",
    "        )\n",
    "        if self.char_inputs and self.share_embedding_softmax:\n",
    "            raise ValueError(\n",
    "                'Sharing softmax and embedding weights requires ' 'word input'\n",
    "            )\n",
    "\n",
    "        self.sample_softmax = options.get('sample_softmax', False)\n",
    "        self._build()\n",
    "        lr = options.get('learning_rate', 0.2)\n",
    "        self.optimizer = tf.train.AdagradOptimizer(\n",
    "            learning_rate = lr, initial_accumulator_value = 1.0\n",
    "        ).minimize(self.total_loss)\n",
    "\n",
    "    def _build_word_embeddings(self):\n",
    "        n_tokens_vocab = self.options['n_tokens_vocab']\n",
    "        batch_size = self.options['batch_size']\n",
    "        unroll_steps = self.options['unroll_steps']\n",
    "\n",
    "        projection_dim = self.options['lstm']['projection_dim']\n",
    "        self.token_ids = tf.placeholder(\n",
    "            tf.int32, shape = (None, unroll_steps), name = 'token_ids'\n",
    "        )\n",
    "        self.batch_size = tf.shape(self.token_ids)[0]\n",
    "        with tf.device('/cpu:0'):\n",
    "\n",
    "            self.embedding_weights = tf.get_variable(\n",
    "                'embedding',\n",
    "                [n_tokens_vocab, projection_dim],\n",
    "                dtype = tf.float32,\n",
    "                initializer = tf.random_uniform_initializer(-1.0, 1.0),\n",
    "            )\n",
    "            self.embedding = tf.nn.embedding_lookup(\n",
    "                self.embedding_weights, self.token_ids\n",
    "            )\n",
    "\n",
    "        if self.bidirectional:\n",
    "            self.token_ids_reverse = tf.placeholder(\n",
    "                tf.int32,\n",
    "                shape = (None, unroll_steps),\n",
    "                name = 'token_ids_reverse',\n",
    "            )\n",
    "            with tf.device('/cpu:0'):\n",
    "                self.embedding_reverse = tf.nn.embedding_lookup(\n",
    "                    self.embedding_weights, self.token_ids_reverse\n",
    "                )\n",
    "\n",
    "    def _build_word_char_embeddings(self):\n",
    "\n",
    "        batch_size = self.options['batch_size']\n",
    "        unroll_steps = self.options['unroll_steps']\n",
    "        projection_dim = self.options['lstm']['projection_dim']\n",
    "\n",
    "        cnn_options = self.options['char_cnn']\n",
    "        filters = cnn_options['filters']\n",
    "        n_filters = sum(f[1] for f in filters)\n",
    "        max_chars = cnn_options['max_characters_per_token']\n",
    "        char_embed_dim = cnn_options['embedding']['dim']\n",
    "        n_chars = cnn_options['n_characters']\n",
    "\n",
    "        if cnn_options['activation'] == 'tanh':\n",
    "            activation = tf.nn.tanh\n",
    "        elif cnn_options['activation'] == 'relu':\n",
    "            activation = tf.nn.relu\n",
    "\n",
    "        self.tokens_characters = tf.placeholder(\n",
    "            tf.int32,\n",
    "            shape = (None, unroll_steps, max_chars),\n",
    "            name = 'tokens_characters',\n",
    "        )\n",
    "        self.batch_size = tf.shape(self.tokens_characters)[0]\n",
    "        with tf.device('/cpu:0'):\n",
    "            self.embedding_weights = tf.get_variable(\n",
    "                'char_embed',\n",
    "                [n_chars, char_embed_dim],\n",
    "                dtype = tf.float32,\n",
    "                initializer = tf.random_uniform_initializer(-1.0, 1.0),\n",
    "            )\n",
    "            self.char_embedding = tf.nn.embedding_lookup(\n",
    "                self.embedding_weights, self.tokens_characters\n",
    "            )\n",
    "\n",
    "            if self.bidirectional:\n",
    "                self.tokens_characters_reverse = tf.placeholder(\n",
    "                    tf.int32,\n",
    "                    shape = (None, unroll_steps, max_chars),\n",
    "                    name = 'tokens_characters_reverse',\n",
    "                )\n",
    "                self.char_embedding_reverse = tf.nn.embedding_lookup(\n",
    "                    self.embedding_weights, self.tokens_characters_reverse\n",
    "                )\n",
    "\n",
    "        def make_convolutions(inp, reuse):\n",
    "            with tf.variable_scope('CNN', reuse = reuse) as scope:\n",
    "                convolutions = []\n",
    "                for i, (width, num) in enumerate(filters):\n",
    "                    if cnn_options['activation'] == 'relu':\n",
    "                        w_init = tf.random_uniform_initializer(\n",
    "                            minval = -0.05, maxval = 0.05\n",
    "                        )\n",
    "                    elif cnn_options['activation'] == 'tanh':\n",
    "                        w_init = tf.random_normal_initializer(\n",
    "                            mean = 0.0,\n",
    "                            stddev = np.sqrt(1.0 / (width * char_embed_dim)),\n",
    "                        )\n",
    "                    w = tf.get_variable(\n",
    "                        'W_cnn_%s' % i,\n",
    "                        [1, width, char_embed_dim, num],\n",
    "                        initializer = w_init,\n",
    "                        dtype = tf.float32,\n",
    "                    )\n",
    "                    b = tf.get_variable(\n",
    "                        'b_cnn_%s' % i,\n",
    "                        [num],\n",
    "                        dtype = tf.float32,\n",
    "                        initializer = tf.constant_initializer(0.0),\n",
    "                    )\n",
    "                    conv = (\n",
    "                        tf.nn.conv2d(\n",
    "                            inp, w, strides = [1, 1, 1, 1], padding = 'VALID'\n",
    "                        )\n",
    "                        + b\n",
    "                    )\n",
    "                    conv = tf.nn.max_pool(\n",
    "                        conv,\n",
    "                        [1, 1, max_chars - width + 1, 1],\n",
    "                        [1, 1, 1, 1],\n",
    "                        'VALID',\n",
    "                    )\n",
    "                    conv = activation(conv)\n",
    "                    conv = tf.squeeze(conv, squeeze_dims = [2])\n",
    "\n",
    "                    convolutions.append(conv)\n",
    "\n",
    "            return tf.concat(convolutions, 2)\n",
    "\n",
    "        reuse = tf.get_variable_scope().reuse\n",
    "        embedding = make_convolutions(self.char_embedding, reuse)\n",
    "        self.token_embedding_layers = [embedding]\n",
    "        if self.bidirectional:\n",
    "            embedding_reverse = make_convolutions(\n",
    "                self.char_embedding_reverse, True\n",
    "            )\n",
    "        n_highway = cnn_options.get('n_highway')\n",
    "        use_highway = n_highway is not None and n_highway > 0\n",
    "        use_proj = n_filters != projection_dim\n",
    "\n",
    "        if use_highway or use_proj:\n",
    "            embedding = tf.reshape(embedding, [-1, n_filters])\n",
    "            if self.bidirectional:\n",
    "                embedding_reverse = tf.reshape(\n",
    "                    embedding_reverse, [-1, n_filters]\n",
    "                )\n",
    "\n",
    "        if use_proj:\n",
    "            assert n_filters > projection_dim\n",
    "            with tf.variable_scope('CNN_proj') as scope:\n",
    "                W_proj_cnn = tf.get_variable(\n",
    "                    'W_proj',\n",
    "                    [n_filters, projection_dim],\n",
    "                    initializer = tf.random_normal_initializer(\n",
    "                        mean = 0.0, stddev = np.sqrt(1.0 / n_filters)\n",
    "                    ),\n",
    "                    dtype = tf.float32,\n",
    "                )\n",
    "                b_proj_cnn = tf.get_variable(\n",
    "                    'b_proj',\n",
    "                    [projection_dim],\n",
    "                    initializer = tf.constant_initializer(0.0),\n",
    "                    dtype = tf.float32,\n",
    "                )\n",
    "\n",
    "        def high(x, ww_carry, bb_carry, ww_tr, bb_tr):\n",
    "            carry_gate = tf.nn.sigmoid(tf.matmul(x, ww_carry) + bb_carry)\n",
    "            transform_gate = tf.nn.relu(tf.matmul(x, ww_tr) + bb_tr)\n",
    "            return carry_gate * transform_gate + (1.0 - carry_gate) * x\n",
    "\n",
    "        if use_highway:\n",
    "            highway_dim = n_filters\n",
    "\n",
    "            for i in range(n_highway):\n",
    "                with tf.variable_scope('CNN_high_%s' % i) as scope:\n",
    "                    W_carry = tf.get_variable(\n",
    "                        'W_carry',\n",
    "                        [highway_dim, highway_dim],\n",
    "                        initializer = tf.random_normal_initializer(\n",
    "                            mean = 0.0, stddev = np.sqrt(1.0 / highway_dim)\n",
    "                        ),\n",
    "                        dtype = tf.float32,\n",
    "                    )\n",
    "                    b_carry = tf.get_variable(\n",
    "                        'b_carry',\n",
    "                        [highway_dim],\n",
    "                        initializer = tf.constant_initializer(-2.0),\n",
    "                        dtype = tf.float32,\n",
    "                    )\n",
    "                    W_transform = tf.get_variable(\n",
    "                        'W_transform',\n",
    "                        [highway_dim, highway_dim],\n",
    "                        initializer = tf.random_normal_initializer(\n",
    "                            mean = 0.0, stddev = np.sqrt(1.0 / highway_dim)\n",
    "                        ),\n",
    "                        dtype = tf.float32,\n",
    "                    )\n",
    "                    b_transform = tf.get_variable(\n",
    "                        'b_transform',\n",
    "                        [highway_dim],\n",
    "                        initializer = tf.constant_initializer(0.0),\n",
    "                        dtype = tf.float32,\n",
    "                    )\n",
    "\n",
    "                embedding = high(\n",
    "                    embedding, W_carry, b_carry, W_transform, b_transform\n",
    "                )\n",
    "                if self.bidirectional:\n",
    "                    embedding_reverse = high(\n",
    "                        embedding_reverse,\n",
    "                        W_carry,\n",
    "                        b_carry,\n",
    "                        W_transform,\n",
    "                        b_transform,\n",
    "                    )\n",
    "                self.token_embedding_layers.append(\n",
    "                    tf.reshape(\n",
    "                        embedding, [self.batch_size, unroll_steps, highway_dim]\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        if use_proj:\n",
    "            embedding = tf.matmul(embedding, W_proj_cnn) + b_proj_cnn\n",
    "            if self.bidirectional:\n",
    "                embedding_reverse = (\n",
    "                    tf.matmul(embedding_reverse, W_proj_cnn) + b_proj_cnn\n",
    "                )\n",
    "            self.token_embedding_layers.append(\n",
    "                tf.reshape(\n",
    "                    embedding, [self.batch_size, unroll_steps, projection_dim]\n",
    "                )\n",
    "            )\n",
    "\n",
    "        if use_highway or use_proj:\n",
    "            shp = [self.batch_size, unroll_steps, projection_dim]\n",
    "            embedding = tf.reshape(embedding, shp)\n",
    "            if self.bidirectional:\n",
    "                embedding_reverse = tf.reshape(embedding_reverse, shp)\n",
    "\n",
    "        self.embedding = embedding\n",
    "        if self.bidirectional:\n",
    "            self.embedding_reverse = embedding_reverse\n",
    "\n",
    "    def _build(self):\n",
    "        n_tokens_vocab = self.options['n_tokens_vocab']\n",
    "        batch_size = self.options['batch_size']\n",
    "        unroll_steps = self.options['unroll_steps']\n",
    "\n",
    "        lstm_dim = self.options['lstm']['dim']\n",
    "        projection_dim = self.options['lstm']['projection_dim']\n",
    "        n_lstm_layers = self.options['lstm'].get('n_layers', 1)\n",
    "        dropout = self.options['dropout']\n",
    "        keep_prob = 1.0 - dropout\n",
    "\n",
    "        if self.char_inputs:\n",
    "            self._build_word_char_embeddings()\n",
    "        else:\n",
    "            self._build_word_embeddings()\n",
    "\n",
    "        self.init_lstm_state = []\n",
    "        self.final_lstm_state = []\n",
    "\n",
    "        if self.bidirectional:\n",
    "            lstm_inputs = [self.embedding, self.embedding_reverse]\n",
    "        else:\n",
    "            lstm_inputs = [self.embedding]\n",
    "\n",
    "        cell_clip = self.options['lstm'].get('cell_clip')\n",
    "        proj_clip = self.options['lstm'].get('proj_clip')\n",
    "\n",
    "        use_skip_connections = self.options['lstm'].get('use_skip_connections')\n",
    "\n",
    "        lstm_outputs = []\n",
    "        for lstm_num, lstm_input in enumerate(lstm_inputs):\n",
    "            lstm_cells = []\n",
    "            for i in range(n_lstm_layers):\n",
    "                lstm_cell = tf.nn.rnn_cell.LSTMCell(\n",
    "                    lstm_dim,\n",
    "                    num_proj = lstm_dim // 2,\n",
    "                    cell_clip = cell_clip,\n",
    "                    proj_clip = proj_clip,\n",
    "                )\n",
    "\n",
    "                if use_skip_connections:\n",
    "                    if i == 0:\n",
    "                        pass\n",
    "                    else:\n",
    "                        lstm_cell = tf.nn.rnn_cell.ResidualWrapper(lstm_cell)\n",
    "\n",
    "                if self.is_training:\n",
    "                    lstm_cell = tf.nn.rnn_cell.DropoutWrapper(\n",
    "                        lstm_cell, input_keep_prob = keep_prob\n",
    "                    )\n",
    "\n",
    "                lstm_cells.append(lstm_cell)\n",
    "\n",
    "            if n_lstm_layers > 1:\n",
    "                lstm_cell = tf.nn.rnn_cell.MultiRNNCell(lstm_cells)\n",
    "            else:\n",
    "                lstm_cell = lstm_cells[0]\n",
    "\n",
    "            with tf.control_dependencies([lstm_input]):\n",
    "                self.init_lstm_state.append(\n",
    "                    lstm_cell.zero_state(self.batch_size, tf.float32)\n",
    "                )\n",
    "                if self.bidirectional:\n",
    "                    with tf.variable_scope('RNN_%s' % lstm_num):\n",
    "                        _lstm_output_unpacked, final_state = tf.nn.static_rnn(\n",
    "                            lstm_cell,\n",
    "                            tf.unstack(lstm_input, axis = 1),\n",
    "                            initial_state = self.init_lstm_state[-1],\n",
    "                        )\n",
    "                else:\n",
    "                    _lstm_output_unpacked, final_state = tf.nn.static_rnn(\n",
    "                        lstm_cell,\n",
    "                        tf.unstack(lstm_input, axis = 1),\n",
    "                        initial_state = self.init_lstm_state[-1],\n",
    "                    )\n",
    "                self.final_lstm_state.append(final_state)\n",
    "\n",
    "            lstm_output_flat = tf.reshape(\n",
    "                tf.stack(_lstm_output_unpacked, axis = 1), [-1, projection_dim]\n",
    "            )\n",
    "            tf.add_to_collection(\n",
    "                'lstm_output_embeddings', _lstm_output_unpacked\n",
    "            )\n",
    "\n",
    "            lstm_outputs.append(lstm_output_flat)\n",
    "        self._build_loss(lstm_outputs)\n",
    "\n",
    "    def _build_loss(self, lstm_outputs):\n",
    "        batch_size = self.options['batch_size']\n",
    "        unroll_steps = self.options['unroll_steps']\n",
    "\n",
    "        n_tokens_vocab = self.options['n_tokens_vocab']\n",
    "\n",
    "        def _get_next_token_placeholders(suffix):\n",
    "            name = 'next_token_id' + suffix\n",
    "            id_placeholder = tf.placeholder(\n",
    "                tf.int32, shape = (None, unroll_steps), name = name\n",
    "            )\n",
    "            return id_placeholder\n",
    "\n",
    "        self.next_token_id = _get_next_token_placeholders('')\n",
    "        if self.bidirectional:\n",
    "            self.next_token_id_reverse = _get_next_token_placeholders(\n",
    "                '_reverse'\n",
    "            )\n",
    "        softmax_dim = self.options['lstm']['projection_dim']\n",
    "        if self.share_embedding_softmax:\n",
    "            self.softmax_W = self.embedding_weights\n",
    "\n",
    "        with tf.variable_scope('softmax'), tf.device('/cpu:0'):\n",
    "            softmax_init = tf.random_normal_initializer(\n",
    "                0.0, 1.0 / np.sqrt(softmax_dim)\n",
    "            )\n",
    "            if not self.share_embedding_softmax:\n",
    "                self.softmax_W = tf.get_variable(\n",
    "                    'W',\n",
    "                    [n_tokens_vocab, softmax_dim],\n",
    "                    dtype = tf.float32,\n",
    "                    initializer = softmax_init,\n",
    "                )\n",
    "            self.softmax_b = tf.get_variable(\n",
    "                'b',\n",
    "                [n_tokens_vocab],\n",
    "                dtype = tf.float32,\n",
    "                initializer = tf.constant_initializer(0.0),\n",
    "            )\n",
    "\n",
    "        self.individual_losses = []\n",
    "\n",
    "        if self.bidirectional:\n",
    "            next_ids = [self.next_token_id, self.next_token_id_reverse]\n",
    "        else:\n",
    "            next_ids = [self.next_token_id]\n",
    "\n",
    "        self.output_scores = tf.identity(lstm_outputs, name = 'softmax_score')\n",
    "\n",
    "        for id_placeholder, lstm_output_flat in zip(next_ids, lstm_outputs):\n",
    "            next_token_id_flat = tf.reshape(id_placeholder, [-1, 1])\n",
    "            with tf.control_dependencies([lstm_output_flat]):\n",
    "                if self.is_training and self.sample_softmax:\n",
    "                    losses = tf.nn.sampled_softmax_loss(\n",
    "                        self.softmax_W,\n",
    "                        self.softmax_b,\n",
    "                        next_token_id_flat,\n",
    "                        lstm_output_flat,\n",
    "                        int(\n",
    "                            self.options['n_negative_samples_batch']\n",
    "                            * self.options['n_tokens_vocab']\n",
    "                        ),\n",
    "                        self.options['n_tokens_vocab'],\n",
    "                        num_true = 1,\n",
    "                    )\n",
    "\n",
    "                else:\n",
    "                    output_scores = (\n",
    "                        tf.matmul(\n",
    "                            lstm_output_flat, tf.transpose(self.softmax_W)\n",
    "                        )\n",
    "                        + self.softmax_b\n",
    "                    )\n",
    "\n",
    "                    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "                        logits = self.output_scores,\n",
    "                        labels = tf.squeeze(\n",
    "                            next_token_id_flat, squeeze_dims = [1]\n",
    "                        ),\n",
    "                    )\n",
    "\n",
    "            self.individual_losses.append(tf.reduce_mean(losses))\n",
    "\n",
    "        if self.bidirectional:\n",
    "            self.total_loss = 0.5 * (\n",
    "                self.individual_losses[0] + self.individual_losses[1]\n",
    "            )\n",
    "        else:\n",
    "            self.total_loss = self.individual_losses[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-12-718029c1ea2f>:140: calling squeeze (from tensorflow.python.ops.array_ops) with squeeze_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use the `axis` argument instead\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = LanguageModel(options, True)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def _get_feed_dict_from_X(X, model, char_inputs, bidirectional):\n",
    "    feed_dict = {}\n",
    "    if not char_inputs:\n",
    "        token_ids = X['token_ids']\n",
    "        feed_dict[model.token_ids] = token_ids\n",
    "    else:\n",
    "        char_ids = X['tokens_characters']\n",
    "        feed_dict[model.tokens_characters] = char_ids\n",
    "    if bidirectional:\n",
    "        if not char_inputs:\n",
    "            feed_dict[model.token_ids_reverse] = X['token_ids_reverse']\n",
    "        else:\n",
    "            feed_dict[model.tokens_characters_reverse] = X['tokens_characters_reverse']\n",
    "    next_id_placeholders = [[model.next_token_id, '']]\n",
    "    if bidirectional:\n",
    "        next_id_placeholders.append([model.next_token_id_reverse, '_reverse'])\n",
    "\n",
    "    for id_placeholder, suffix in next_id_placeholders:\n",
    "        name = 'next_token_id' + suffix\n",
    "        feed_dict[id_placeholder] = X[name]\n",
    "\n",
    "    return feed_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "bidirectional = options.get('bidirectional', False)\n",
    "batch_size = options['batch_size']\n",
    "unroll_steps = options['unroll_steps']\n",
    "n_train_tokens = options.get('n_train_tokens')\n",
    "n_tokens_per_batch = batch_size * unroll_steps\n",
    "n_batches_per_epoch = int(n_train_tokens / n_tokens_per_batch)\n",
    "n_batches_total = options['n_epochs'] * n_batches_per_epoch\n",
    "\n",
    "init_state_tensors = model.init_lstm_state\n",
    "final_state_tensors = model.final_lstm_state\n",
    "\n",
    "char_inputs = 'char_cnn' in options\n",
    "if char_inputs:\n",
    "    max_chars = options['char_cnn']['max_characters_per_token']\n",
    "    feed_dict = {\n",
    "        model.tokens_characters: np.zeros(\n",
    "            [batch_size, unroll_steps, max_chars], dtype = np.int32\n",
    "        )\n",
    "    }\n",
    "\n",
    "else:\n",
    "    feed_dict = {model.token_ids: np.zeros([batch_size, unroll_steps])}\n",
    "\n",
    "if bidirectional:\n",
    "    if char_inputs:\n",
    "        feed_dict.update(\n",
    "            {\n",
    "                model.tokens_characters_reverse: np.zeros(\n",
    "                    [batch_size, unroll_steps, max_chars], dtype = np.int32\n",
    "                )\n",
    "            }\n",
    "        )\n",
    "    else:\n",
    "        feed_dict.update(\n",
    "            {\n",
    "                model.token_ids_reverse: np.zeros(\n",
    "                    [batch_size, unroll_steps], dtype = np.int32\n",
    "                )\n",
    "            }\n",
    "        )\n",
    "\n",
    "init_state_values = sess.run(init_state_tensors, feed_dict = feed_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 6300/6300 [28:51<00:00,  3.64it/s, cost=3.01]\n"
     ]
    }
   ],
   "source": [
    "data_gen = bi.iter_batches(batch_size, unroll_steps)\n",
    "pbar = tqdm(range(n_batches_total), desc = 'train minibatch loop')\n",
    "for p in pbar:\n",
    "    batch = next(data_gen)\n",
    "    feed_dict = {t: v for t, v in zip(init_state_tensors, init_state_values)}\n",
    "    feed_dict.update(_get_feed_dict_from_X(batch, model, char_inputs, bidirectional))\n",
    "    score, loss, _, init_state_values = sess.run([model.output_scores,\n",
    "                                           model.total_loss, model.optimizer, final_state_tensors],\n",
    "            feed_dict = feed_dict)\n",
    "    pbar.set_postfix(cost = loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_embed = model.softmax_W.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import cdist\n",
    "from sklearn.neighbors import NearestNeighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['comedy', 0.42140477895736694],\n",
       " ['film', 0.3776022791862488],\n",
       " ['one', 0.3622022271156311],\n",
       " ['story', 0.36113226413726807],\n",
       " ['watch', 0.3598611354827881],\n",
       " ['movie', 0.3553805351257324],\n",
       " ['like', 0.354833722114563],\n",
       " ['almost', 0.3544100522994995],\n",
       " ['time', 0.35355472564697266]]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word = 'beautiful'\n",
    "nn = NearestNeighbors(10, metric = 'cosine').fit(word_embed)\n",
    "distances, idx = nn.kneighbors(word_embed[dictionary[word]].reshape((1, -1)))\n",
    "word_list = []\n",
    "for i in range(1, idx.shape[1]):\n",
    "    word_list.append([rev_dictionary[idx[0, i]], 1 - distances[0, i]])\n",
    "word_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}