{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "# Skip-gram word2vec\n",
    "\n",
    "In this notebook, I'll lead you through using TensorFlow to implement the word2vec algorithm using the skip-gram architecture. By implementing this, you'll learn about embedding words for use in natural language processing. This will come in handy when dealing with things like translations.\n",
    "\n",
    "## Readings\n",
    "\n",
    "Here are the resources I used to build this notebook. I suggest reading these either beforehand or while you're working on this material.\n",
    "\n",
    "* A really good [conceptual overview](http://mccormickml.com/2016/04/19/word2vec-tutorial-the-skip-gram-model/) of word2vec from Chris McCormick \n",
    "* [First word2vec paper](https://arxiv.org/pdf/1301.3781.pdf) from Mikolov et al.\n",
    "* [NIPS paper](http://papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf) with improvements for word2vec also from Mikolov et al.\n",
    "* An [implementation of word2vec](http://www.thushv.com/natural_language_processing/word2vec-part-1-nlp-with-deep-learning-with-tensorflow-skip-gram/) from Thushan Ganegedara\n",
    "* TensorFlow [word2vec tutorial](https://www.tensorflow.org/tutorials/word2vec)\n",
    "\n",
    "## Word embeddings\n",
    "\n",
    "When you're dealing with language and words, you end up with tens of thousands of classes to predict, one for each word. Trying to one-hot encode these words is massively inefficient, you'll have one element set to 1 and the other 50,000 set to 0. The word2vec algorithm finds much more efficient representations by finding vectors that represent the words. These vectors also contain semantic information about the words. Words that show up in similar contexts, such as \"black\", \"white\", and \"red\" will have vectors near each other. There are two architectures for implementing word2vec, CBOW (Continuous Bag-Of-Words) and Skip-gram.\n",
    "\n",
    "<img src=\"assets/word2vec_architectures.png\" width=\"500\">\n",
    "\n",
    "In this implementation, we'll be using the skip-gram architecture because it performs better than CBOW. Here, we pass in a word and try to predict the words surrounding it in the text. In this way, we can train the network to learn representations for words that show up in similar contexts.\n",
    "\n",
    "First up, importing packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the [text8 dataset](http://mattmahoney.net/dc/textdata.html), a file of cleaned up Wikipedia articles from Matt Mahoney. The next cell will download the data set to the `data` folder. Then you can extract it and delete the archive file to save storage space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from urllib.request import urlretrieve\n",
    "from os.path import isfile, isdir\n",
    "from tqdm import tqdm\n",
    "import zipfile\n",
    "\n",
    "dataset_folder_path = 'data'\n",
    "dataset_filename = 'text8.zip'\n",
    "dataset_name = 'Text8 Dataset'\n",
    "\n",
    "class DLProgress(tqdm):\n",
    "    last_block = 0\n",
    "\n",
    "    def hook(self, block_num=1, block_size=1, total_size=None):\n",
    "        self.total = total_size\n",
    "        self.update((block_num - self.last_block) * block_size)\n",
    "        self.last_block = block_num\n",
    "\n",
    "if not isfile(dataset_filename):\n",
    "    with DLProgress(unit='B', unit_scale=True, miniters=1, desc=dataset_name) as pbar:\n",
    "        urlretrieve(\n",
    "            'http://mattmahoney.net/dc/text8.zip',\n",
    "            dataset_filename,\n",
    "            pbar.hook)\n",
    "\n",
    "if not isdir(dataset_folder_path):\n",
    "    with zipfile.ZipFile(dataset_filename) as zip_ref:\n",
    "        zip_ref.extractall(dataset_folder_path)\n",
    "        \n",
    "with open('data/text8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "\n",
    "Here I'm fixing up the text to make training easier. This comes from the `utils` module I wrote. The `preprocess` function coverts any punctuation into tokens, so a period is changed to ` <PERIOD> `. In this data set, there aren't any periods, but it will help in other NLP problems. I'm also removing all words that show up five or fewer times in the dataset. This will greatly reduce issues due to noise in the data and improve the quality of the vector representations. If you want to write your own functions for this stuff, go for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "words = utils.preprocess(text)\n",
    "print(words[:30])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"Total words: {}\".format(len(words)))\n",
    "print(\"Unique words: {}\".format(len(set(words))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here I'm creating dictionaries to covert words to integers and backwards, integers to words. The integers are assigned in descending frequency order, so the most frequent word (\"the\") is given the integer 0 and the next most frequent is 1 and so on. The words are converted to integers and stored in the list `int_words`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "vocab_to_int, int_to_vocab = utils.create_lookup_tables(words)\n",
    "int_words = [vocab_to_int[word] for word in words]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subsampling\n",
    "\n",
    "Words that show up often such as \"the\", \"of\", and \"for\" don't provide much context to the nearby words. If we discard some of them, we can remove some of the noise from our data and in return get faster training and better representations. This process is called subsampling by Mikolov. For each word $w_i$ in the training set, we'll discard it with probability given by \n",
    "\n",
    "$$ P(w_i) = 1 - \\sqrt{\\frac{t}{f(w_i)}} $$\n",
    "\n",
    "where $t$ is a threshold parameter and $f(w_i)$ is the frequency of word $w_i$ in the total dataset.\n",
    "\n",
    "I'm going to leave this up to you as an exercise. This is more of a programming challenge, than about deep learning specifically. But, being able to prepare your data for your network is an important skill to have. Check out my solution to see how I did it.\n",
    "\n",
    "> **Exercise:** Implement subsampling for the words in `int_words`. That is, go through `int_words` and discard each word given the probablility $P(w_i)$ shown above. Note that $P(w_i)$ is the probability that a word is discarded. Assign the subsampled data to `train_words`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "## Your code here\n",
    "train_words = # The final subsampled word list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making batches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Now that our data is in good shape, we need to get it into the proper form to pass it into our network. With the skip-gram architecture, for each word in the text, we want to grab all the words in a window around that word, with size $C$. \n",
    "\n",
    "From [Mikolov et al.](https://arxiv.org/pdf/1301.3781.pdf): \n",
    "\n",
    "\"Since the more distant words are usually less related to the current word than those close to it, we give less weight to the distant words by sampling less from those words in our training examples... If we choose $C = 5$, for each training word we will select randomly a number $R$ in range $< 1; C >$, and then use $R$ words from history and $R$ words from the future of the current word as correct labels.\"\n",
    "\n",
    "> **Exercise:** Implement a function `get_target` that receives a list of words, an index, and a window size, then returns a list of words in the window around the index. Make sure to use the algorithm described above, where you choose a random number of words from the window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def get_target(words, idx, window_size=5):\n",
    "    ''' Get a list of words in a window around an index. '''\n",
    "    \n",
    "    # Your code here\n",
    "    \n",
    "    return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a function that returns batches for our network. The idea is that it grabs `batch_size` words from a words list. Then for each of those words, it gets the target words in the window. I haven't found a way to pass in a random number of target words and get it to work with the architecture, so I make one row per input-target pair. This is a generator function by the way, helps save memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def get_batches(words, batch_size, window_size=5):\n",
    "    ''' Create a generator of word batches as a tuple (inputs, targets) '''\n",
    "    \n",
    "    n_batches = len(words)//batch_size\n",
    "    \n",
    "    # only full batches\n",
    "    words = words[:n_batches*batch_size]\n",
    "    \n",
    "    for idx in range(0, len(words), batch_size):\n",
    "        x, y = [], []\n",
    "        batch = words[idx:idx+batch_size]\n",
    "        for ii in range(len(batch)):\n",
    "            batch_x = batch[ii]\n",
    "            batch_y = get_target(batch, ii, window_size)\n",
    "            y.extend(batch_y)\n",
    "            x.extend([batch_x]*len(batch_y))\n",
    "        yield x, y\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Building the graph\n",
    "\n",
    "From Chris McCormick's blog, we can see the general structure of our network.\n",
    "![embedding_network](./assets/skip_gram_net_arch.png)\n",
    "\n",
    "The input words are passed in as one-hot encoded vectors. This will go into a hidden layer of linear units, then into a softmax layer. We'll use the softmax layer to make a prediction like normal.\n",
    "\n",
    "The idea here is to train the hidden layer weight matrix to find efficient representations for our words. This weight matrix is usually called the embedding matrix or embedding look-up table. We can discard the softmax layer becuase we don't really care about making predictions with this network. We just want the embedding matrix so we can use it in other networks we build from the dataset.\n",
    "\n",
    "I'm going to have you build the graph in stages now. First off, creating the `inputs` and `labels` placeholders like normal.\n",
    "\n",
    "> **Exercise:** Assign `inputs` and `labels` using `tf.placeholder`. We're going to be passing in integers, so set the data types to `tf.int32`. The batches we're passing in will have varying sizes, so set the batch sizes to [`None`]. To make things work later, you'll need to set the second dimension of `labels` to `None` or `1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "train_graph = tf.Graph()\n",
    "with train_graph.as_default():\n",
    "    inputs = \n",
    "    labels = "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "The embedding matrix has a size of the number of words by the number of neurons in the hidden layer. So, if you have 10,000 words and 300 hidden units, the matrix will have size $10,000 \\times 300$. Remember that we're using one-hot encoded vectors for our inputs. When you do the matrix multiplication of the one-hot vector with the embedding matrix, you end up selecting only one row out of the entire matrix:\n",
    "\n",
    "![one-hot matrix multiplication](assets/matrix_mult_w_one_hot.png)\n",
    "\n",
    "You don't actually need to do the matrix multiplication, you just need to select the row in the embedding matrix that corresponds to the input word. Then, the embedding matrix becomes a lookup table, you're looking up a vector the size of the hidden layer that represents the input word.\n",
    "\n",
    "<img src=\"assets/word2vec_weight_matrix_lookup_table.png\" width=500>\n",
    "\n",
    "\n",
    "> **Exercise:** Tensorflow provides a convenient function [`tf.nn.embedding_lookup`](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup) that does this lookup for us. You pass in the embedding matrix and a tensor of integers, then it returns rows in the matrix corresponding to those integers. Below, set the number of embedding features you'll use (200 is a good start), create the embedding matrix variable, and use [`tf.nn.embedding_lookup`](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup) to get the embedding tensors. For the embedding matrix, I suggest you initialize it with a uniform random numbers between -1 and 1 using [tf.random_uniform](https://www.tensorflow.org/api_docs/python/tf/random_uniform). This [TensorFlow tutorial](https://www.tensorflow.org/tutorials/word2vec) will help if you get stuck."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_vocab = len(int_to_vocab)\n",
    "n_embedding =  # Number of embedding features \n",
    "with train_graph.as_default():\n",
    "    embedding = # create embedding weight matrix here\n",
    "    embed = # use tf.nn.embedding_lookup to get the hidden layer output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Negative sampling\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For every example we give the network, we train it using the output from the softmax layer. That means for each input, we're making very small changes to millions of weights even though we only have one true example. This makes training the network very inefficient. We can approximate the loss from the softmax layer by only updating a small subset of all the weights at once. We'll update the weights for the correct label, but only a small number of incorrect labels. This is called [\"negative sampling\"](http://papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf). Tensorflow has a convenient function to do this, [`tf.nn.sampled_softmax_loss`](https://www.tensorflow.org/api_docs/python/tf/nn/sampled_softmax_loss).\n",
    "\n",
    "> **Exercise:** Below, create weights and biases for the softmax layer. Then, use [`tf.nn.sampled_softmax_loss`](https://www.tensorflow.org/api_docs/python/tf/nn/sampled_softmax_loss) to calculate the loss. Be sure to read the documentation to figure out how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# Number of negative labels to sample\n",
    "n_sampled = 100\n",
    "with train_graph.as_default():\n",
    "    softmax_w = # create softmax weight matrix here\n",
    "    softmax_b = # create softmax biases here\n",
    "    \n",
    "    # Calculate the loss using negative sampling\n",
    "    loss = tf.nn.sampled_softmax_loss \n",
    "    \n",
    "    cost = tf.reduce_mean(loss)\n",
    "    optimizer = tf.train.AdamOptimizer().minimize(cost)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation\n",
    "\n",
    "This code is from Thushan Ganegedara's implementation. Here we're going to choose a few common words and few uncommon words. Then, we'll print out the closest words to them. It's a nice way to check that our embedding table is grouping together words with similar semantic meanings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with train_graph.as_default():\n",
    "    ## From Thushan Ganegedara's implementation\n",
    "    valid_size = 16 # Random set of words to evaluate similarity on.\n",
    "    valid_window = 100\n",
    "    # pick 8 samples from (0,100) and (1000,1100) each ranges. lower id implies more frequent \n",
    "    valid_examples = np.array(random.sample(range(valid_window), valid_size//2))\n",
    "    valid_examples = np.append(valid_examples, \n",
    "                               random.sample(range(1000,1000+valid_window), valid_size//2))\n",
    "\n",
    "    valid_dataset = tf.constant(valid_examples, dtype=tf.int32)\n",
    "    \n",
    "    # We use the cosine distance:\n",
    "    norm = tf.sqrt(tf.reduce_sum(tf.square(embedding), 1, keep_dims=True))\n",
    "    normalized_embedding = embedding / norm\n",
    "    valid_embedding = tf.nn.embedding_lookup(normalized_embedding, valid_dataset)\n",
    "    similarity = tf.matmul(valid_embedding, tf.transpose(normalized_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# If the checkpoints directory doesn't exist:\n",
    "!mkdir checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Below is the code to train the network. Every 100 batches it reports the training loss. Every 1000 batches, it'll print out the validation words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "epochs = 10\n",
    "batch_size = 1000\n",
    "window_size = 10\n",
    "\n",
    "with train_graph.as_default():\n",
    "    saver = tf.train.Saver()\n",
    "\n",
    "with tf.Session(graph=train_graph) as sess:\n",
    "    iteration = 1\n",
    "    loss = 0\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    for e in range(1, epochs+1):\n",
    "        batches = get_batches(train_words, batch_size, window_size)\n",
    "        start = time.time()\n",
    "        for x, y in batches:\n",
    "            \n",
    "            feed = {inputs: x,\n",
    "                    labels: np.array(y)[:, None]}\n",
    "            train_loss, _ = sess.run([cost, optimizer], feed_dict=feed)\n",
    "            \n",
    "            loss += train_loss\n",
    "            \n",
    "            if iteration % 100 == 0: \n",
    "                end = time.time()\n",
    "                print(\"Epoch {}/{}\".format(e, epochs),\n",
    "                      \"Iteration: {}\".format(iteration),\n",
    "                      \"Avg. Training loss: {:.4f}\".format(loss/100),\n",
    "                      \"{:.4f} sec/batch\".format((end-start)/100))\n",
    "                loss = 0\n",
    "                start = time.time()\n",
    "            \n",
    "            if iteration % 1000 == 0:\n",
    "                ## From Thushan Ganegedara's implementation\n",
    "                # note that this is expensive (~20% slowdown if computed every 500 steps)\n",
    "                sim = similarity.eval()\n",
    "                for i in range(valid_size):\n",
    "                    valid_word = int_to_vocab[valid_examples[i]]\n",
    "                    top_k = 8 # number of nearest neighbors\n",
    "                    nearest = (-sim[i, :]).argsort()[1:top_k+1]\n",
    "                    log = 'Nearest to %s:' % valid_word\n",
    "                    for k in range(top_k):\n",
    "                        close_word = int_to_vocab[nearest[k]]\n",
    "                        log = '%s %s,' % (log, close_word)\n",
    "                    print(log)\n",
    "            \n",
    "            iteration += 1\n",
    "    save_path = saver.save(sess, \"checkpoints/text8.ckpt\")\n",
    "    embed_mat = sess.run(normalized_embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Restore the trained network if you need to:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with train_graph.as_default():\n",
    "    saver = tf.train.Saver()\n",
    "\n",
    "with tf.Session(graph=train_graph) as sess:\n",
    "    saver.restore(sess, tf.train.latest_checkpoint('checkpoints'))\n",
    "    embed_mat = sess.run(embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the word vectors\n",
    "\n",
    "Below we'll use T-SNE to visualize how our high-dimensional word vectors cluster together. T-SNE is used to project these vectors into two dimensions while preserving local stucture. Check out [this post from Christopher Olah](http://colah.github.io/posts/2014-10-Visualizing-MNIST/) to learn more about T-SNE and other ways to visualize high-dimensional data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.manifold import TSNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "viz_words = 500\n",
    "tsne = TSNE()\n",
    "embed_tsne = tsne.fit_transform(embed_mat[:viz_words, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 14))\n",
    "for idx in range(viz_words):\n",
    "    plt.scatter(*embed_tsne[idx, :], color='steelblue')\n",
    "    plt.annotate(int_to_vocab[idx], (embed_tsne[idx, 0], embed_tsne[idx, 1]), alpha=0.7)"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
