word2vec.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "from scipy import spatial\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project 3: Word2Vec (70 pt)\n",
    "The goal of this project is to obtain the vector representations for words from text.\n",
    "\n",
    "The main idea is that words appearing in similar contexts have similar meanings. Because of that, word vectors of similar words should be close together. Models that use word vectors can utilize these properties, e.g., in sentiment analysis a model will learn that \"good\" and \"great\" are positive words, but will also generalize to other words that it has not seen (e.g. \"amazing\") because they should be close together in the vector space.\n",
    "\n",
    "Vectors can keep other language properties as well, like analogies. The question \"a is to b as c is to ...?\", where the answer is d, can be answered by looking into word vector space and calculating $\\mathbf{u}_b - \\mathbf{u}_a + \\mathbf{u}_c$, and finding the word vector that is the closest to the result.\n",
    "\n",
    "## Your task\n",
    "Complete the missing code in this notebook. Make sure that all the functions follow the provided specification, i.e. the output of the function exactly matches the description in the docstring. \n",
    "\n",
    "We are given a text that contains $N$ unique words $\\{ x_1, ..., x_N \\}$. We will focus on the Skip-Gram model in which the goal is to predict the context window $S = \\{ x_{i-l}, ..., x_{i-1}, x_{i+1}, ..., x_{i+l} \\}$ from current word $x_i$, where $l$ is the window size. \n",
    "\n",
    "We get a word embedding $\\mathbf{u}_i$ by multiplying the matrix $\\mathbf{U}$ with a one-hot representation $\\mathbf{x}_i$ of a word $x_i$. Then, to get output probabilities for context window, we multiply this embedding with another matrix $\\mathbf{V}$ and apply softmax. The objective is to minimize the loss: $-\\mathop{\\mathbb{E}}[P(S|x_i;\\mathbf{U}, \\mathbf{V})]$.\n",
    "\n",
    "You are given a dataset with positive and negative reviews. Your task is to:\n",
    "+ Construct input-output pairs corresponding to the current word and a word in the context window\n",
    "+ Implement forward and backward propagation with parameter updates for Skip-Gram model\n",
    "+ Train the model\n",
    "+ Test it on word analogies and sentiment analysis task\n",
    "\n",
    "## General remarks\n",
    "\n",
    "Fill in the missing code at the markers\n",
    "```python\n",
    "###########################\n",
    "# YOUR CODE HERE\n",
    "###########################\n",
    "```\n",
    "Do not add or modify any code at other places in the notebook except where otherwise explicitly stated.\n",
    "After you fill in all the missing code, restart the kernel and re-run all the cells in the notebook.\n",
    "\n",
    "The following things are **NOT** allowed:\n",
    "- Using additional `import` statements\n",
    "- Copying / reusing code from other sources (e.g. code by other students)\n",
    "\n",
    "If you plagiarise even for a single project task, you won't be eligible for the bonus this semester."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load data (5 pts)\n",
    "\n",
    "We'll be working with a subset of reviews for restaurants in Las Vegas. The reviews that we'll be working with are either 1-star or 5-star."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = np.load(\"task03_data.npy\", allow_pickle=True).item()\n",
    "reviews_1star = [[x.lower() for x in s] for s in data[\"reviews_1star\"]]\n",
    "reviews_5star = [[x.lower() for x in s] for s in data[\"reviews_5star\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate the vocabulary by taking the top 200 words by their frequency from both positive and negative sentences. We could also use the whole vocabulary, but that would be slower. Other words are represented with \"unk\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "corpus = reviews_1star + reviews_5star\n",
    "vocabulary = [w for s in corpus for w in s]\n",
    "vocabulary, counts = zip(*Counter(vocabulary).most_common(200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = [[w if w in vocabulary else 'unk' for w in s] for s in corpus]\n",
    "vocabulary += ('unk',) # Add \"unk\" to vocabulary\n",
    "counts += (sum([w == 'unk' for s in corpus for w in s]),) # Add count for \"unk\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "VOCABULARY_SIZE = len(vocabulary)\n",
    "EMBEDDING_DIM = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of positive reviews: 1000\n",
      "Number of negative reviews: 2000\n",
      "Number of unique words: 201\n"
     ]
    }
   ],
   "source": [
    "print('Number of positive reviews:', len(reviews_1star))\n",
    "print('Number of negative reviews:', len(reviews_5star))\n",
    "print('Number of unique words:', VOCABULARY_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3000\n"
     ]
    }
   ],
   "source": [
    "print(len(corpus))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have to create two dictionaries: `word_to_ind` and `ind_to_word` so we can go from text to numerical representation and vice versa. The input into the model will be the index of the word denoting the position in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Implement\n",
    "---------\n",
    "word_to_ind: dict\n",
    "    The keys are words (str) and the value is the corresponding position in the vocabulary\n",
    "ind_to_word: dict\n",
    "    The keys are indices (int) and the value is the corresponding word from the vocabulary\n",
    "ind_to_freq: dict\n",
    "    The keys are indices (int) and the value is the corresponding count in the vocabulary\n",
    "\"\"\"\n",
    "\n",
    "###########################\n",
    "# YOUR CODE HERE\n",
    "word_to_ind = {}\n",
    "ind_to_word = {}\n",
    "ind_to_freq = {}\n",
    "frequency = 0\n",
    "for index in range(len(vocabulary)):\n",
    "    words = vocabulary[index]\n",
    "    word_to_ind[words] = index\n",
    "    ind_to_word[index] = words\n",
    "    for all_word in corpus:\n",
    "        frequency += Counter(all_word)[words]\n",
    "    ind_to_freq[index] = frequency\n",
    "###########################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Word \"the\" is at position 0 appearing 2017 times\n"
     ]
    }
   ],
   "source": [
    "print(f'Word \\\"the\\\" is at position {word_to_ind[\"the\"]} appearing {ind_to_freq[word_to_ind[\"the\"]]} times')\n",
    "# Should output: \n",
    "# Word \"the\" is at position 0 appearing 2017 times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Create word pairs (5 pts)\n",
    "\n",
    "We need all the word pairs $\\{ x_i, x_j \\}$, where $x_i$ is the current word and $x_j$ is from its context window. These will correspond to input-output pairs. We want them to be represented numericaly so you should use `word_to_ind` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_window(sentence, window_size):\n",
    "    \"\"\"\n",
    "    Collect all ordered word pairs from a sentence that are at most window_size apart.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    sentence: list\n",
    "        A list of words of str type forming a sentence\n",
    "    window_size: int\n",
    "        A positive scalar\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    pairs: list\n",
    "        A list of tuple (word index, word index from its context) of int type\n",
    "    \"\"\"\n",
    "    \n",
    "    pairs = []\n",
    "\n",
    "    ###########################\n",
    "    # YOUR CODE HERE\n",
    "#     if window_size > np.floor(len(sentence) / 2):\n",
    "#         raise ValueError('Window_size is too big')\n",
    "\n",
    "    for index in range(len(sentence)):\n",
    "        if index - window_size < 0 and index + window_size <= len(sentence) - 1:\n",
    "            for label in range(0, index):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "            for label in range(index + 1, index + window_size + 1):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "        elif index - window_size >= 0 and index + window_size <= len(sentence) - 1:\n",
    "            for label in range(index - window_size, index):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "            for label in range(index + 1, index + window_size + 1):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "        elif index - window_size >= 0 and index + window_size > len(sentence) - 1:\n",
    "            for label in range(index - window_size, index):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "            for label in range(index + 1, len(sentence)):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "        else:\n",
    "            for label in range(0, index):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "            for label in range(index + 1, len(sentence)):\n",
    "                pairs += [[word_to_ind[sentence[index]], word_to_ind[sentence[label]]]]\n",
    "                #      \n",
    "    ###########################\n",
    "\n",
    "    return pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First 5 pairs: [[10, 200], [10, 6], [10, 64], [200, 10], [200, 6]]\n",
      "Total number of pairs: 207462\n"
     ]
    }
   ],
   "source": [
    "data = []\n",
    "for x in corpus:\n",
    "    data += get_window(x, window_size=3)\n",
    "data = np.array(data)\n",
    "\n",
    "print('First 5 pairs:', data[:5].tolist())\n",
    "print('Total number of pairs:', data.shape[0])\n",
    "# Should output:\n",
    "# First 5 pairs: [[10, 200], [10, 6], [10, 64], [200, 10], [200, 6]]\n",
    "# Total number of pairs: 207462"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We calculate a weighting score to counter the imbalance between the rare and frequent words. Rare words will be sampled more frequently. See https://arxiv.org/pdf/1310.4546.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4.81989821e-06 4.81989821e-06 4.81989821e-06]\n"
     ]
    }
   ],
   "source": [
    "probabilities = [1 - np.sqrt(1e-3 / ind_to_freq[x]) for x in data[:,0]]\n",
    "probabilities /= np.sum(probabilities)\n",
    "print(probabilities[:3])\n",
    "# Should output: \n",
    "# [4.8206203e-06 4.8206203e-06 4.8206203e-06]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Model definition (55 pts)\n",
    "\n",
    "In this part you should implement forward and backward propagation together with update of the parameters i.e.:\n",
    "+ One-hot encoding of the words(5 pts)\n",
    "+ Softmax (5 pts)\n",
    "+ Loss implementation & computation (5 pts)\n",
    "+ Forward pass (15 pts)\n",
    "+ Backward pass (15 pts)\n",
    "+ Optimizer (10 pts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Embedding():\n",
    "    \"\"\"\n",
    "    Word embedding model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    N: int\n",
    "        Number of unique words in the vocabulary\n",
    "    D: int\n",
    "        Dimension of the word vector embedding\n",
    "    \"\"\"\n",
    "    def __init__(self, N, D):\n",
    "        self.N = N\n",
    "        self.D = D\n",
    "\n",
    "        self.ctx = None # Used to store values for backpropagation\n",
    "\n",
    "        self.U = None\n",
    "        self.V = None\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        We initialize weight matrices U and V of dimension (D, N) and (N, D), respectively\n",
    "        \"\"\"\n",
    "        self.ctx = None\n",
    "        self.U = np.random.normal(0, np.sqrt(6. / (self.D + self.N)), (self.D, self.N))\n",
    "        self.V = np.random.normal(0, np.sqrt(6. / (self.D + self.N)), (self.N, self.D))\n",
    "\n",
    "    def one_hot(self, x, N):\n",
    "        \"\"\"\n",
    "        Given a vector returns a matrix with rows corresponding to one-hot encoding.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        x: array\n",
    "            M-dimensional vector containing integers from [0, N - 1]\n",
    "        N: int\n",
    "            Number of posible classes\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        one_hot: array\n",
    "            (N, M) matrix where each column is N-dimensional one-hot encoding of elements from x \n",
    "        \"\"\"\n",
    "\n",
    "        ###########################\n",
    "        # YOUR CODE HERE\n",
    "        M = x.shape[0]\n",
    "        x_one_hot = np.zeros((N,M))\n",
    "        x_one_hot[x,np.arange(M)] = 1 \n",
    "        one_hot = x_one_hot\n",
    "        ###########################\n",
    "\n",
    "        assert one_hot.shape == (N, x.shape[0]), 'Incorrect one-hot embedding shape'\n",
    "        return one_hot\n",
    "\n",
    "    def softmax(self, x, axis):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: array\n",
    "            A non-empty matrix of any dimension\n",
    "        axis: int\n",
    "            Dimension on which softmax is performed\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        y: array\n",
    "            Matrix of same dimension as x with softmax applied to 'axis' dimension\n",
    "        \"\"\"\n",
    "        \n",
    "        # Note! You should implement a numerically stable version of softmax\n",
    "        \n",
    "        ###########################\n",
    "        # YOUR CODE HERE\n",
    "        x_exp = np.exp(x - np.max(x, axis=axis, keepdims=True))\n",
    "        y = x_exp / np.sum(x_exp, axis=axis, keepdims=True)\n",
    "        ###########################\n",
    "        \n",
    "        assert x.shape == y.shape, 'Output should have the same shape is input'\n",
    "        return y\n",
    "\n",
    "    def loss(self, y, prob):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        y: array\n",
    "            (N, M) matrix of M samples where columns are one-hot vectors for true values\n",
    "        prob: array\n",
    "            (N, M) column of M samples where columns are probability vectors after softmax\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        loss: int\n",
    "            Cross-entropy loss calculated as: 1 / M * sum_i(sum_j(y_ij * log(prob_ij)))\n",
    "        \"\"\"\n",
    "\n",
    "        prob = np.clip(prob, 1e-8, None)\n",
    "        \n",
    "        ###########################\n",
    "        # YOUR CODE HERE\n",
    "        M = y.shape[1]\n",
    "        inside = y * np.log(prob)\n",
    "        loss = - 1 / M * np.sum(inside)\n",
    "        ###########################\n",
    "        \n",
    "        assert isinstance(loss, float), 'Loss must be a scalar'\n",
    "        return loss\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        \"\"\"\n",
    "        Performs forward and backward propagation and updates weights\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        x: array\n",
    "            M-dimensional mini-batched vector containing input word indices of int type\n",
    "        y: array\n",
    "            Output words, same dimension and type as 'x'\n",
    "        learning_rate: float\n",
    "            A positive scalar determining the update rate\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        loss: float\n",
    "            Cross-entropy loss\n",
    "        d_U: array\n",
    "            Partial derivative of loss w.r.t. U\n",
    "        d_V: array\n",
    "            Partial derivative of loss w.r.t. V\n",
    "        \"\"\"\n",
    "        \n",
    "        # Input transformation\n",
    "        \"\"\"\n",
    "        Input is represented with M-dimensional vectors\n",
    "        We convert them to (N, M) matrices such that columns are one-hot \n",
    "        representations of the input\n",
    "        \"\"\"\n",
    "        x = self.one_hot(x, self.N)\n",
    "        y = self.one_hot(y, self.N)\n",
    "        \n",
    "        # Forward propagation\n",
    "        \"\"\"\n",
    "        Returns\n",
    "        -------\n",
    "        embedding: array\n",
    "            (D, M) matrix where columns are word embedding from U matrix\n",
    "        logits: array\n",
    "            (N, M) matrix where columns are output logits\n",
    "        prob: array\n",
    "            (N, M) matrix where columns are output probabilities\n",
    "        \"\"\"\n",
    "        \n",
    "        ###########################\n",
    "        # YOUR CODE HERE\n",
    "        # x:N*M;  U:D*N ;  V:N*D\n",
    "        embedding = self.U @ x  # D*M\n",
    "        logits = self.V @ embedding  # N*M\n",
    "        prob = self.softmax(logits, axis=0)  #N*M\n",
    "        ###########################\n",
    "\n",
    "        assert embedding.shape == (self.D, x.shape[1])\n",
    "        assert logits.shape == (self.N, x.shape[1])\n",
    "        assert prob.shape == (self.N, x.shape[1])\n",
    "        \n",
    "        # Save values for backpropagation\n",
    "        self.ctx = (embedding, logits, prob, x, y)\n",
    "        \n",
    "        # Loss calculation\n",
    "        loss = self.loss(y, prob)\n",
    "        \n",
    "        return loss\n",
    "        \n",
    "    def backward(self):\n",
    "        \"\"\"\n",
    "        Given parameters from forward propagation, returns gradient of U and V.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        d_V: array\n",
    "            (D, N) matrix of partial derivatives of loss w.r.t. V\n",
    "        d_U: array\n",
    "            (N, D) matrix of partial derivatives of loss w.r.t. U\n",
    "        \"\"\"\n",
    "\n",
    "        embedding, logits, prob, x, y = self.ctx\n",
    "\n",
    "        ###########################\n",
    "        # YOUR CODE HERE\n",
    "        M = y.shape[1]\n",
    "        d_U = np.matmul(self.V.T, (prob - y)) @ x.T / M\n",
    "        d_V = np.matmul((prob - y), embedding.T) / M\n",
    "        ###########################\n",
    "\n",
    "        assert d_V.shape == (self.N, self.D)\n",
    "        assert d_U.shape == (self.D, self.N)\n",
    "\n",
    "        return { 'V': d_V, 'U': d_U }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Optimizer():\n",
    "    \"\"\"\n",
    "    Stochastic gradient descent with momentum optimizer.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: object\n",
    "        Model as defined above\n",
    "    learning_rate: float\n",
    "        Learning rate\n",
    "    momentum: float (optional)\n",
    "        Momentum factor (default: 0)\n",
    "    \"\"\"\n",
    "    def __init__(self, model, learning_rate, momentum=0):\n",
    "        self.model = model\n",
    "        self.learning_rate = learning_rate\n",
    "        self.momentum = momentum\n",
    "        \n",
    "        self.previous = None # Previous gradients\n",
    "    \n",
    "    def _init_previous(self, grad):\n",
    "        # Initialize previous gradients to zero\n",
    "        self.previous = { k: np.zeros_like(v) for k,v in grad.items() }\n",
    "    \n",
    "    def step(self, grad):\n",
    "        if self.previous is None:\n",
    "            self._init_previous(grad)\n",
    "            \n",
    "        for name, dw in grad.items():\n",
    "            dw_prev = self.previous[name]\n",
    "            w = getattr(self.model, name)\n",
    "\n",
    "            \"\"\"\n",
    "            Given weight w, previous gradients dw_prev and current \n",
    "            gradients dw, performs an update of weight w.\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            dw_new: array\n",
    "                New gradients calculated as combination of previous and\n",
    "                current, weighted with momentum factor.\n",
    "            w_new: array\n",
    "                New weights calculated with a single step of gradient\n",
    "                descent using dw_new.\n",
    "            \"\"\"\n",
    "            ###########################\n",
    "            # YOUR CODE HERE\n",
    "            dw_new = self.momentum * dw_prev + (1 - self.momentum) * dw\n",
    "            w_new = w - self.learning_rate * dw_new  #iterative rule of SGDmomentum\n",
    "            ###########################\n",
    "\n",
    "            self.previous[name] = dw_new\n",
    "            setattr(self.model, name, w_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Gradient check\n",
    "\n",
    "The following code checks whether the updates for weights are implemented correctly. It should run without an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradients checked - all good!\n"
     ]
    }
   ],
   "source": [
    "def get_loss(model, old, variable, epsilon, x, y, i, j):\n",
    "    np.random.seed(123)\n",
    "    model.reset_parameters() # reset weights\n",
    "    \n",
    "    delta = np.zeros_like(old)\n",
    "    delta[i, j] = epsilon\n",
    "    \n",
    "    setattr(model, variable, old + delta) # change one weight by a small amount\n",
    "    loss = model.forward(x, y)\n",
    "\n",
    "    return loss\n",
    "\n",
    "def gradient_check_for_weight(model, variable, i, j, k, l):\n",
    "    x, y = np.array([i, i]), np.array([j, j]) # set input and output\n",
    "    \n",
    "    np.random.seed(123)\n",
    "    model.reset_parameters() # reset weights\n",
    "\n",
    "    old = getattr(model, variable)\n",
    "    \n",
    "    loss = model.forward(x, y)\n",
    "    grad = model.backward()\n",
    "    \n",
    "    eps = 1e-4\n",
    "    loss_positive = get_loss(model, old, variable, eps, x, y, k, l) # loss for positive change on one weight\n",
    "    loss_negative = get_loss(model, old, variable, -eps, x, y, k, l) # loss for negative change on one weight\n",
    "    \n",
    "    true_gradient = (loss_positive - loss_negative) / 2 / eps # calculate true derivative wrt one weight\n",
    "\n",
    "    assert abs(true_gradient - grad[variable][k, l]) < 1e-5, 'Incorrect gradient'\n",
    "\n",
    "def gradient_check():\n",
    "    N, D = VOCABULARY_SIZE, EMBEDDING_DIM\n",
    "    model = Embedding(N, D)\n",
    "\n",
    "    # check for V\n",
    "    for _ in range(20):\n",
    "        i, j, k = [np.random.randint(0, d) for d in [N, N, D]] # get random indices for input and weights\n",
    "        gradient_check_for_weight(model, 'V', i, j, i, k)\n",
    "\n",
    "    # check for U\n",
    "    for _ in range(20):\n",
    "        i, j, k = [np.random.randint(0, d) for d in [N, N, D]]\n",
    "        gradient_check_for_weight(model, 'U', i, j, k, i)\n",
    "\n",
    "    print('Gradients checked - all good!')\n",
    "\n",
    "gradient_check()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train our model using stochastic gradient descent. At every step we sample a mini-batch from data and update the weights.\n",
    "\n",
    "The following function samples words from data and creates mini-batches. It subsamples frequent words based on previously calculated probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(123)\n",
    "def get_batch(data, size, prob):\n",
    "    x = rng.choice(data, size, p=prob)\n",
    "    return x[:,0], x[:,1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training the model can take some time so plan accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration: 1000, Avg. training loss: 3.7228\n",
      "Iteration: 2000, Avg. training loss: 3.5667\n",
      "Iteration: 3000, Avg. training loss: 3.5492\n",
      "Iteration: 4000, Avg. training loss: 3.5353\n",
      "Iteration: 5000, Avg. training loss: 3.5241\n",
      "Iteration: 6000, Avg. training loss: 3.5186\n",
      "Iteration: 7000, Avg. training loss: 3.5112\n",
      "Iteration: 8000, Avg. training loss: 3.5051\n",
      "Iteration: 9000, Avg. training loss: 3.4965\n",
      "Iteration: 10000, Avg. training loss: 3.4915\n",
      "Iteration: 11000, Avg. training loss: 3.4860\n",
      "Iteration: 12000, Avg. training loss: 3.4848\n",
      "Iteration: 13000, Avg. training loss: 3.4829\n",
      "Iteration: 14000, Avg. training loss: 3.4782\n",
      "Iteration: 15000, Avg. training loss: 3.4806\n"
     ]
    }
   ],
   "source": [
    "model = Embedding(N=VOCABULARY_SIZE, D=EMBEDDING_DIM)\n",
    "optim = Optimizer(model, learning_rate=1.0, momentum=0.5)\n",
    "\n",
    "losses = []\n",
    "\n",
    "MAX_ITERATIONS = 15000\n",
    "PRINT_EVERY = 1000\n",
    "BATCH_SIZE = 1000\n",
    "\n",
    "for i in range(MAX_ITERATIONS):\n",
    "    x, y = get_batch(data, BATCH_SIZE, probabilities)\n",
    "    \n",
    "    loss = model.forward(x, y)\n",
    "    grad = model.backward()\n",
    "    optim.step(grad)\n",
    "    \n",
    "    assert not np.isnan(loss)\n",
    "    \n",
    "    losses.append(loss)\n",
    "\n",
    "    if (i + 1) % PRINT_EVERY == 0:\n",
    "        print(f'Iteration: {i + 1}, Avg. training loss: {np.mean(losses[-PRINT_EVERY:]):.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The embedding matrix is given by $\\mathbf{U}^T$, where the $i$th row is the vector for $i$th word in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.40381486 -0.1623919   0.30170626 ... -0.12495809 -0.19067946\n",
      "  -0.01389494]\n",
      " [-0.48401953  0.14918166  0.10936935 ...  0.73559488  0.22574966\n",
      "   0.16230905]\n",
      " [ 0.09782615 -0.2265617   0.15842491 ...  0.10891626  0.4361375\n",
      "  -0.3954269 ]\n",
      " ...\n",
      " [-0.10627612 -0.14266463  0.20496442 ...  0.49341252  0.07193523\n",
      "   0.15084159]\n",
      " [-0.17454098  0.15453431  0.20148213 ... -0.08285212  0.14985704\n",
      "  -0.09214063]\n",
      " [-0.14876761 -0.05720596  0.08581473 ...  0.16177278  0.04756354\n",
      "   0.28996516]]\n"
     ]
    }
   ],
   "source": [
    "emb_matrix = model.U.T\n",
    "print(emb_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Analogies (5 pts)\n",
    "\n",
    "As mentioned before, vectors can keep some language properties like analogies. Given a relation a:b and a query c, we can find d such that c:d follows the same relation. We hope to find d by using vector operations. In this case, finding the real word vector $\\mathbf{u}_d$ closest to $\\mathbf{u}_b - \\mathbf{u}_a + \\mathbf{u}_c$ gives us d. \n",
    "\n",
    "**Note that the quality of the analy results is not expected to be excellent.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "`is` is to `was` as [are, rice, sauce, way, at] is to `were`\n",
      "`lunch` is to `day` as [dinner, great, salad, experience, what] is to `night`\n",
      "`i` is to `my` as [you, not, don't, what, if] is to `your`\n"
     ]
    }
   ],
   "source": [
    "triplets = [['is', 'was', 'were'], ['lunch', 'day', 'night'], ['i', 'my', 'your']]\n",
    "\n",
    "for triplet in triplets:\n",
    "    a, b, d = triplet\n",
    "\n",
    "    \"\"\"\n",
    "    Returns\n",
    "    \n",
    "    Example: Paris (a) is to France (b) as _____ (c) is to Germany (d)\n",
    "    \n",
    "    -------\n",
    "    result: array\n",
    "        The embedding vector for word (c): w_a - w_b + w_d\n",
    "    \"\"\"\n",
    "\n",
    "    ###########################\n",
    "    # YOUR CODE HERE\n",
    "    index_a = word_to_ind[a]\n",
    "    index_b = word_to_ind[b]\n",
    "    index_d = word_to_ind[d]\n",
    "    w_a = emb_matrix[index_a]\n",
    "    w_b = emb_matrix[index_b]\n",
    "    w_d = emb_matrix[index_d]\n",
    "    result = w_a - w_b + w_d\n",
    "    ###########################\n",
    "\n",
    "    distances = [spatial.distance.cosine(x, result) for x in emb_matrix]\n",
    "    candidates = [ind_to_word[i] for i in np.argsort(distances)]\n",
    "    candidates = [x for x in candidates if x not in [a, b, d]][:5]\n",
    "\n",
    "    print(f'`{a}` is to `{b}` as [{\", \".join(candidates)}] is to `{d}`')"
   ]
  },
  {
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}