word2vec-embeddings/Negative_Sampling_Exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Skip-gram Word2Vec\n",
    "\n",
    "In this notebook, I'll lead you through using PyTorch to implement the [Word2Vec algorithm](https://en.wikipedia.org/wiki/Word2vec) 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 machine translation.\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",
    "* [Neural Information Processing Systems, 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",
    "\n",
    "---\n",
    "## Word embeddings\n",
    "\n",
    "When you're dealing with words in text, you end up with tens of thousands of word classes to analyze; one for each word in a vocabulary. Trying to one-hot encode these words is massively inefficient because most values in a one-hot vector will be set to zero. So, the matrix multiplication that happens in between a one-hot input vector and a first, hidden layer will result in mostly zero-valued hidden outputs.\n",
    "\n",
    "To solve this problem and greatly increase the efficiency of our networks, we use what are called **embeddings**. Embeddings are just a fully connected layer like you've seen before. We call this layer the embedding layer and the weights are embedding weights. We skip the multiplication into the embedding layer by instead directly grabbing the hidden layer values from the weight matrix. We can do this because the multiplication of a one-hot encoded vector with a matrix returns the row of the matrix corresponding the index of the \"on\" input unit.\n",
    "\n",
    "<img src='assets/lookup_matrix.png' width=50%>\n",
    "\n",
    "Instead of doing the matrix multiplication, we use the weight matrix as a lookup table. We encode the words as integers, for example \"heart\" is encoded as 958, \"mind\" as 18094. Then to get hidden layer values for \"heart\", you just take the 958th row of the embedding matrix. This process is called an **embedding lookup** and the number of hidden units is the **embedding dimension**.\n",
    " \n",
    "There is nothing magical going on here. The embedding lookup table is just a weight matrix. The embedding layer is just a hidden layer. The lookup is just a shortcut for the matrix multiplication. The lookup table is trained just like any weight matrix.\n",
    "\n",
    "Embeddings aren't only used for words of course. You can use them for any model where you have a massive number of classes. A particular type of model called **Word2Vec** uses the embedding layer to find vector representations of words that contain semantic meaning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Word2Vec\n",
    "\n",
    "The Word2Vec algorithm finds much more efficient representations by finding vectors that represent the words. These vectors also contain semantic information about the words.\n",
    "\n",
    "<img src=\"assets/context_drink.png\" width=40%>\n",
    "\n",
    "Words that show up in similar **contexts**, such as \"coffee\", \"tea\", and \"water\" will have vectors near each other. Different words will be further away from one another, and relationships can be represented by distance in vector space.\n",
    "\n",
    "\n",
    "There are two architectures for implementing Word2Vec:\n",
    ">* CBOW (Continuous Bag-Of-Words) and \n",
    "* Skip-gram\n",
    "\n",
    "<img src=\"assets/word2vec_architectures.png\" width=60%>\n",
    "\n",
    "In this implementation, we'll be using the **skip-gram architecture** with **negative sampling** because it performs better than CBOW and trains faster with negative sampling. 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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Loading Data\n",
    "\n",
    "Next, we'll ask you to load in data and place it in the `data` directory\n",
    "\n",
    "1. Load the [text8 dataset](https://s3.amazonaws.com/video.udacity-data.com/topher/2018/October/5bbe6499_text8/text8.zip); a file of cleaned up *Wikipedia article text* from Matt Mahoney. \n",
    "2. Place that data in the `data` folder in the home directory.\n",
    "3. Then you can extract it and delete the archive, zip file to save storage space.\n",
    "\n",
    "After following these steps, you should have one file in your data directory: `data/text8`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# read in the extracted text file      \n",
    "with open('data/text8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# print out the first 100 characters\n",
    "print(text[:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-processing\n",
    "\n",
    "Here I'm fixing up the text to make training easier. This comes from the `utils.py` file. The `preprocess` function does a few things:\n",
    ">* It converts 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. \n",
    "* It removes 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. \n",
    "* It returns a list of words in the text.\n",
    "\n",
    "This may take a few seconds to run, since our text file is quite large. If you want to write your own functions for this stuff, go for it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import utils\n",
    "\n",
    "# get list of words\n",
    "words = utils.preprocess(text)\n",
    "print(words[:30])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# print some stats about this word data\n",
    "print(\"Total words in text: {}\".format(len(words)))\n",
    "print(\"Unique words: {}\".format(len(set(words)))) # `set` removes any duplicate words"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dictionaries\n",
    "\n",
    "Next, I'm creating two dictionaries to convert words to integers and back again (integers to words). This is again done with a function in the `utils.py` file. `create_lookup_tables` takes in a list of words in a text and returns two dictionaries.\n",
    ">* 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. \n",
    "\n",
    "Once we have our dictionaries, the words are converted to integers and stored in the list `int_words`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": 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]\n",
    "\n",
    "print(int_words[:30])"
   ]
  },
  {
   "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",
    "> 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": true
   },
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "threshold = 1e-5\n",
    "word_counts = Counter(int_words)\n",
    "#print(list(word_counts.items())[0])  # dictionary of int_words, how many times they appear\n",
    "\n",
    "total_count = len(int_words)\n",
    "freqs = {word: count/total_count for word, count in word_counts.items()}\n",
    "p_drop = {word: 1 - np.sqrt(threshold/freqs[word]) for word in word_counts}\n",
    "# discard some frequent words, according to the subsampling equation\n",
    "# create a new list of words for training\n",
    "train_words = [word for word in int_words if random.random() < (1 - p_drop[word])]\n",
    "\n",
    "print(train_words[:30])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making batches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 define a surrounding _context_ and 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 chose a random number of words to from the window.\n",
    "\n",
    "Say, we have an input and we're interested in the idx=2 token, `741`: \n",
    "```\n",
    "[5233, 58, 741, 10571, 27349, 0, 15067, 58112, 3580, 58, 10712]\n",
    "```\n",
    "\n",
    "For `R=2`, `get_target` should return a list of four values:\n",
    "```\n",
    "[5233, 58, 10571, 27349]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": 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",
    "    R = np.random.randint(1, window_size+1)\n",
    "    start = idx - R if (idx - R) > 0 else 0\n",
    "    stop = idx + R\n",
    "    target_words = words[start:idx] + words[idx+1:stop+1]\n",
    "    \n",
    "    return list(target_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# test your code!\n",
    "\n",
    "# run this cell multiple times to check for random window selection\n",
    "int_text = [i for i in range(10)]\n",
    "print('Input: ', int_text)\n",
    "idx=5 # word index of interest\n",
    "\n",
    "target = get_target(int_text, idx=idx, window_size=5)\n",
    "print('Target: ', target)  # you should get some indices around the idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating Batches \n",
    "\n",
    "Here's a generator function that returns batches of input and target data for our model, using the `get_target` function from above. The idea is that it grabs `batch_size` words from a words list. Then for each of those batches, it gets the target words in a window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": 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": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "int_text = [i for i in range(20)]\n",
    "x,y = next(get_batches(int_text, batch_size=4, window_size=5))\n",
    "\n",
    "print('x\\n', x)\n",
    "print('y\\n', y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Validation\n",
    "\n",
    "Here, I'm creating a function that will help us observe our model as it learns. We're going to choose a few common words and few uncommon words. Then, we'll print out the closest words to them using the cosine similarity: \n",
    "\n",
    "<img src=\"assets/two_vectors.png\" width=30%>\n",
    "\n",
    "$$\n",
    "\\mathrm{similarity} = \\cos(\\theta) = \\frac{\\vec{a} \\cdot \\vec{b}}{|\\vec{a}||\\vec{b}|}\n",
    "$$\n",
    "\n",
    "\n",
    "We can encode the validation words as vectors $\\vec{a}$ using the embedding table, then calculate the similarity with each word vector $\\vec{b}$ in the embedding table. With the similarities, we can print out the validation words and words in our embedding table semantically similar to those words. 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": [
    "def cosine_similarity(embedding, valid_size=16, valid_window=100, device='cpu'):\n",
    "    \"\"\" Returns the cosine similarity of validation words with words in the embedding matrix.\n",
    "        Here, embedding should be a PyTorch embedding module.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Here we're calculating the cosine similarity between some random words and \n",
    "    # our embedding vectors. With the similarities, we can look at what words are\n",
    "    # close to our random words.\n",
    "    \n",
    "    # sim = (a . b) / |a||b|\n",
    "    \n",
    "    embed_vectors = embedding.weight\n",
    "    \n",
    "    # magnitude of embedding vectors, |b|\n",
    "    magnitudes = embed_vectors.pow(2).sum(dim=1).sqrt().unsqueeze(0)\n",
    "    \n",
    "    # pick N words from our ranges (0,window) and (1000,1000+window). 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",
    "    valid_examples = torch.LongTensor(valid_examples).to(device)\n",
    "    \n",
    "    valid_vectors = embedding(valid_examples)\n",
    "    similarities = torch.mm(valid_vectors, embed_vectors.t())/magnitudes\n",
    "        \n",
    "    return valid_examples, similarities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# SkipGram model\n",
    "\n",
    "Define and train the SkipGram model. \n",
    "> You'll need to define an [embedding layer](https://pytorch.org/docs/stable/nn.html#embedding) and a final, softmax output layer.\n",
    "\n",
    "An Embedding layer takes in a number of inputs, importantly:\n",
    "* **num_embeddings** – the size of the dictionary of embeddings, or how many rows you'll want in the embedding weight matrix\n",
    "* **embedding_dim** – the size of each embedding vector; the embedding dimension\n",
    "\n",
    "Below is an approximate diagram of the general structure of our network.\n",
    "<img src=\"assets/skip_gram_arch.png\" width=60%>\n",
    "\n",
    ">* The input words are passed in as batches of input word tokens. \n",
    "* This will go into a hidden layer of linear units (our embedding layer). \n",
    "* Then, finally into a softmax output layer. \n",
    "\n",
    "We'll use the softmax layer to make a prediction about the context words by sampling, as usual."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Negative Sampling\n",
    "\n",
    "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 example, but only a small number of incorrect, or noise, examples. This is called [\"negative sampling\"](http://papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf). \n",
    "\n",
    "There are two modifications we need to make. First, since we're not taking the softmax output over all the words, we're really only concerned with one output word at a time. Similar to how we use an embedding table to map the input word to the hidden layer, we can now use another embedding table to map the hidden layer to the output word. Now we have two embedding layers, one for input words and one for output words. Secondly, we use a modified loss function where we only care about the true example and a small subset of noise examples.\n",
    "\n",
    "$$\n",
    "- \\large \\log{\\sigma\\left(u_{w_O}\\hspace{0.001em}^\\top v_{w_I}\\right)} -\n",
    "\\sum_i^N \\mathbb{E}_{w_i \\sim P_n(w)}\\log{\\sigma\\left(-u_{w_i}\\hspace{0.001em}^\\top v_{w_I}\\right)}\n",
    "$$\n",
    "\n",
    "This is a little complicated so I'll go through it bit by bit. $u_{w_O}\\hspace{0.001em}^\\top$ is the embedding vector for our \"output\" target word (transposed, that's the $^\\top$ symbol) and $v_{w_I}$ is the embedding vector for the \"input\" word. Then the first term \n",
    "\n",
    "$$\\large \\log{\\sigma\\left(u_{w_O}\\hspace{0.001em}^\\top v_{w_I}\\right)}$$\n",
    "\n",
    "says we take the log-sigmoid of the inner product of the output word vector and the input word vector. Now the second term, let's first look at \n",
    "\n",
    "$$\\large \\sum_i^N \\mathbb{E}_{w_i \\sim P_n(w)}$$ \n",
    "\n",
    "This means we're going to take a sum over words $w_i$ drawn from a noise distribution $w_i \\sim P_n(w)$. The noise distribution is basically our vocabulary of words that aren't in the context of our input word. In effect, we can randomly sample words from our vocabulary to get these words. $P_n(w)$ is an arbitrary probability distribution though, which means we get to decide how to weight the words that we're sampling. This could be a uniform distribution, where we sample all words with equal probability. Or it could be according to the frequency that each word shows up in our text corpus, the unigram distribution $U(w)$. The authors found the best distribution to be $U(w)^{3/4}$, empirically. \n",
    "\n",
    "Finally, in \n",
    "\n",
    "$$\\large \\log{\\sigma\\left(-u_{w_i}\\hspace{0.001em}^\\top v_{w_I}\\right)},$$ \n",
    "\n",
    "we take the log-sigmoid of the negated inner product of a noise vector with the input vector. \n",
    "\n",
    "<img src=\"assets/neg_sampling_loss.png\" width=50%>\n",
    "\n",
    "To give you an intuition for what we're doing here, remember that the sigmoid function returns a probability between 0 and 1. The first term in the loss pushes the probability that our network will predict the correct word $w_O$ towards 1. In the second term, since we are negating the sigmoid input, we're pushing the probabilities of the noise words towards 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SkipGramNeg(nn.Module):\n",
    "    def __init__(self, n_vocab, n_embed, noise_dist=None):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_vocab = n_vocab\n",
    "        self.n_embed = n_embed\n",
    "        self.noise_dist = noise_dist\n",
    "        \n",
    "        # define embedding layers for input and output words\n",
    "        self.in_embed = \n",
    "        self.out_embed = \n",
    "        \n",
    "        # Initialize both embedding tables with uniform distribution\n",
    "        \n",
    "    def forward_input(self, input_words):\n",
    "        # return input vector embeddings\n",
    "\n",
    "        return None\n",
    "    \n",
    "    def forward_output(self, output_words):\n",
    "        # return output vector embeddings\n",
    "\n",
    "        return None\n",
    "    \n",
    "    def forward_noise(self, batch_size, n_samples):\n",
    "        \"\"\" Generate noise vectors with shape (batch_size, n_samples, n_embed)\"\"\"\n",
    "        if self.noise_dist is None:\n",
    "            # Sample words uniformly\n",
    "            noise_dist = torch.ones(self.n_vocab)\n",
    "        else:\n",
    "            noise_dist = self.noise_dist\n",
    "            \n",
    "        # Sample words from our noise distribution\n",
    "        noise_words = torch.multinomial(noise_dist,\n",
    "                                        batch_size * n_samples,\n",
    "                                        replacement=True)\n",
    "        \n",
    "        device = \"cuda\" if model.out_embed.weight.is_cuda else \"cpu\"\n",
    "        noise_words = noise_words.to(device)\n",
    "        \n",
    "        ## TODO: get the noise embeddings\n",
    "        # reshape the embeddings so that they have dims (batch_size, n_samples, n_embed)\n",
    "\n",
    "        \n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NegativeSamplingLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, input_vectors, output_vectors, noise_vectors):\n",
    "        \n",
    "        batch_size, embed_size = input_vectors.shape\n",
    "        \n",
    "        # Input vectors should be a batch of column vectors\n",
    "        input_vectors = input_vectors.view(batch_size, embed_size, 1)\n",
    "        \n",
    "        # Output vectors should be a batch of row vectors\n",
    "        output_vectors = output_vectors.view(batch_size, 1, embed_size)\n",
    "        \n",
    "        # bmm = batch matrix multiplication\n",
    "        # correct log-sigmoid loss\n",
    "        out_loss = torch.bmm(output_vectors, input_vectors).sigmoid().log()\n",
    "        out_loss = out_loss.squeeze()\n",
    "        \n",
    "        # incorrect log-sigmoid loss\n",
    "        noise_loss = torch.bmm(noise_vectors.neg(), input_vectors).sigmoid().log()\n",
    "        noise_loss = noise_loss.squeeze().sum(1)  # sum the losses over the sample of noise vectors\n",
    "\n",
    "        # negate and sum correct and noisy log-sigmoid losses\n",
    "        # return average batch loss\n",
    "        return -(out_loss + noise_loss).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "Below is our training loop, and I recommend that you train on GPU, if available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# Get our noise distribution\n",
    "# Using word frequencies calculated earlier in the notebook\n",
    "word_freqs = np.array(sorted(freqs.values(), reverse=True))\n",
    "unigram_dist = word_freqs/word_freqs.sum()\n",
    "noise_dist = torch.from_numpy(unigram_dist**(0.75)/np.sum(unigram_dist**(0.75)))\n",
    "\n",
    "# instantiating the model\n",
    "embedding_dim = 300\n",
    "model = SkipGramNeg(len(vocab_to_int), embedding_dim, noise_dist=noise_dist).to(device)\n",
    "\n",
    "# using the loss that we defined\n",
    "criterion = NegativeSamplingLoss() \n",
    "optimizer = optim.Adam(model.parameters(), lr=0.003)\n",
    "\n",
    "print_every = 1500\n",
    "steps = 0\n",
    "epochs = 5\n",
    "\n",
    "# train for some number of epochs\n",
    "for e in range(epochs):\n",
    "    \n",
    "    # get our input, target batches\n",
    "    for input_words, target_words in get_batches(train_words, 512):\n",
    "        steps += 1\n",
    "        inputs, targets = torch.LongTensor(input_words), torch.LongTensor(target_words)\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        \n",
    "        # input, outpt, and noise vectors\n",
    "        input_vectors = model.forward_input(inputs)\n",
    "        output_vectors = model.forward_output(targets)\n",
    "        noise_vectors = model.forward_noise(inputs.shape[0], 5)\n",
    "\n",
    "        # negative sampling loss\n",
    "        loss = criterion(input_vectors, output_vectors, noise_vectors)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # loss stats\n",
    "        if steps % print_every == 0:\n",
    "            print(\"Epoch: {}/{}\".format(e+1, epochs))\n",
    "            print(\"Loss: \", loss.item()) # avg batch loss at this point in training\n",
    "            valid_examples, valid_similarities = cosine_similarity(model.in_embed, device=device)\n",
    "            _, closest_idxs = valid_similarities.topk(6)\n",
    "\n",
    "            valid_examples, closest_idxs = valid_examples.to('cpu'), closest_idxs.to('cpu')\n",
    "            for ii, valid_idx in enumerate(valid_examples):\n",
    "                closest_words = [int_to_vocab[idx.item()] for idx in closest_idxs[ii]][1:]\n",
    "                print(int_to_vocab[valid_idx.item()] + \" | \" + ', '.join(closest_words))\n",
    "            print(\"...\\n\")"
   ]
  },
  {
   "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": 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": true
   },
   "outputs": [],
   "source": [
    "# getting embeddings from the embedding layer of our model, by name\n",
    "embeddings = model.in_embed.weight.to('cpu').data.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "viz_words = 380\n",
    "tsne = TSNE()\n",
    "embed_tsne = tsne.fit_transform(embeddings[:viz_words, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(16, 16))\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}