Module8/08_Word2vec_pytorch_empty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BZ88HxNEA9r0"
   },
   "source": [
    "# Word Embedding (word2vec)\n",
    "\n",
    "This notebook is a PyTorch adaptation of the [mxnet implementation of word2vec](http://www.d2l.ai/chapter_natural-language-processing-pretraining/word2vec.html) of the book [Dive into Deep Learning](http://www.d2l.ai/index.html). \n",
    "\n",
    "**Update (2023)** there is now a PyTorch implementation provided in [Pretraining word2vec](http://www.d2l.ai/chapter_natural-language-processing-pretraining/word2vec-pretraining.html#pretraining-word2vec) which is sligthly different form the one proposed here.\n",
    "\n",
    "A natural language is a complex system that we use to express meanings. In this system, words are the basic unit of linguistic meaning. As its name implies, a word vector is a vector used to represent a word. It can also be thought of as the feature vector of a word. The technique of mapping words to vectors of real numbers is also known as word embedding. Over the last few years, word embedding has gradually become basic knowledge in natural language processing.\n",
    "\n",
    "## Why not Use One-hot Vectors?\n",
    "\n",
    "\n",
    "Although one-hot word vectors are easy to construct, they are usually not a good choice. One of the major reasons is that the one-hot word vectors cannot accurately express the similarity between different words, such as the cosine similarity that we commonly use. For the vectors $\\boldsymbol{x}, \\boldsymbol{y} \\in \\mathbb{R}^d$, their cosine similarities are the cosines of the angles between them:\n",
    "\n",
    "$$\\frac{\\boldsymbol{x}^\\top \\boldsymbol{y}}{\\|\\boldsymbol{x}\\| \\|\\boldsymbol{y}\\|} \\in [-1, 1].$$\n",
    "\n",
    "Since the cosine similarity between the one-hot vectors of any two different words is 0, it is difficult to use the one-hot vector to accurately represent the similarity between multiple different words.\n",
    "\n",
    "[Word2vec](https://code.google.com/archive/p/word2vec/) is a tool to solve the problem above.  It represents each word with a fixed-length vector and uses these vectors to better indicate the similarity and analogy relationships between different words. The Word2vec tool contains two models: [skip-gram](Distributed representations of words and phrases and their compositionality.) and continuous bag of words [CBOW](Efficient estimation of word representations in vector space). Next, we will take a look at the two models and their training methods.\n",
    "\n",
    "\n",
    "## The Skip-Gram Model\n",
    "\n",
    "The skip-gram model assumes that a word can be used to generate the words that surround it in a text sequence. For example, we assume that the text sequence is \"the\", \"man\", \"loves\", \"his\", and \"son\". We use \"loves\" as the central target word and set the context window size to 2. As shown below, given the central target word \"loves\", the skip-gram model is concerned with the conditional probability for generating the context words, \"the\", \"man\", \"his\" and \"son\", that are within a distance of no more than 2 words, which is\n",
    "\n",
    "$$\\mathbb{P}(\\textrm{\"the\"},\\textrm{\"man\"},\\textrm{\"his\"},\\textrm{\"son\"}\\mid\\textrm{\"loves\"}).$$\n",
    "\n",
    "We assume that, given the central target word, the context words are generated independently of each other. In this case, the formula above can be rewritten as\n",
    "\n",
    "$$\\mathbb{P}(\\textrm{\"the\"}\\mid\\textrm{\"loves\"})\\cdot\\mathbb{P}(\\textrm{\"man\"}\\mid\\textrm{\"loves\"})\\cdot\\mathbb{P}(\\textrm{\"his\"}\\mid\\textrm{\"loves\"})\\cdot\\mathbb{P}(\\textrm{\"son\"}\\mid\\textrm{\"loves\"}).$$\n",
    "\n",
    "![The skip-gram model cares about the conditional probability of generating context words for a given central target word. ](https://www.di.ens.fr/~lelarge/skip-gram.svg)\n",
    "\n",
    "\n",
    "In the skip-gram model, each word is represented as two $d$-dimension vectors, which are used to compute the conditional probability. We assume that the word is indexed as $i$ in the dictionary, its vector is represented as $\\boldsymbol{v}_i\\in\\mathbb{R}^d$ when it is the central target word, and $\\boldsymbol{u}_i\\in\\mathbb{R}^d$ when it is a context word.  Let the central target word $w_c$ and context word $w_o$ be indexed as $c$ and $o$ respectively in the dictionary. The conditional probability of generating the context word for the given central target word can be obtained by performing a softmax operation on the vector inner product:\n",
    "\n",
    "$$\\mathbb{P}(w_o \\mid w_c) = \\frac{\\text{exp}(\\boldsymbol{u}_o^\\top \\boldsymbol{v}_c)}{ \\sum_{i \\in \\mathcal{V}} \\text{exp}(\\boldsymbol{u}_i^\\top \\boldsymbol{v}_c)},$$\n",
    "\n",
    "where vocabulary index set $\\mathcal{V} = \\{0, 1, \\ldots, |\\mathcal{V}|-1\\}$. Assume that a text sequence of length $T$ is given, where the word at time step $t$ is denoted as $w^{(t)}$. Assume that context words are independently generated given center words. When context window size is $m$, the likelihood function of the skip-gram model is the joint probability of generating all the context words given any center word\n",
    "\n",
    "$$ \\prod_{t=1}^{T} \\prod_{-m \\leq j \\leq m,\\ j \\neq 0} \\mathbb{P}(w^{(t+j)} \\mid w^{(t)}),$$\n",
    "\n",
    "After the training, for any word in the dictionary with index $i$, we are going to get its two word vector sets $\\boldsymbol{v}_i$ and $\\boldsymbol{u}_i$.  In applications of natural language processing (NLP), the central target word vector in the skip-gram model is generally used as the representation vector of a word.\n",
    "\n",
    "## The Continuous Bag Of Words (CBOW) Model\n",
    "\n",
    "The continuous bag of words (CBOW) model is similar to the skip-gram model. The biggest difference is that the CBOW model assumes that the central target word is generated based on the context words before and after it in the text sequence. With the same text sequence \"the\", \"man\", \"loves\", \"his\" and \"son\", in which \"loves\" is the central target word, given a context window size of 2, the CBOW model is concerned with the conditional probability of generating the target word \"loves\" based on the context words \"the\", \"man\", \"his\" and \"son\"(as shown below), such as\n",
    "\n",
    "$$\\mathbb{P}(\\textrm{\"loves\"}\\mid\\textrm{\"the\"},\\textrm{\"man\"},\\textrm{\"his\"},\\textrm{\"son\"}).$$\n",
    "\n",
    "![The CBOW model cares about the conditional probability of generating the central target word from given context words.  ](https://www.di.ens.fr/~lelarge/cbow.svg)\n",
    "\n",
    "Since there are multiple context words in the CBOW model, we will average their word vectors and then use the same method as the skip-gram model to compute the conditional probability. We assume that $\\boldsymbol{v_i}\\in\\mathbb{R}^d$ and $\\boldsymbol{u_i}\\in\\mathbb{R}^d$ are the context word vector and central target word vector of the word with index $i$ in the dictionary (notice that the symbols are opposite to the ones in the skip-gram model). Let central target word $w_c$ be indexed as $c$, and context words $w_{o_1}, \\ldots, w_{o_{2m}}$ be indexed as $o_1, \\ldots, o_{2m}$ in the dictionary. Thus, the conditional probability of generating a central target word from the given context word is\n",
    "\n",
    "$$\\mathbb{P}(w_c \\mid w_{o_1}, \\ldots, w_{o_{2m}}) = \\frac{\\text{exp}\\left(\\frac{1}{2m}\\boldsymbol{u}_c^\\top (\\boldsymbol{v}_{o_1} + \\ldots + \\boldsymbol{v}_{o_{2m}}) \\right)}{ \\sum_{i \\in \\mathcal{V}} \\text{exp}\\left(\\frac{1}{2m}\\boldsymbol{u}_i^\\top (\\boldsymbol{v}_{o_1} + \\ldots + \\boldsymbol{v}_{o_{2m}}) \\right)}.$$\n",
    "\n",
    "\n",
    "For brevity, denote $\\mathcal{W}_o= \\{w_{o_1}, \\ldots, w_{o_{2m}}\\}$, and $\\bar{\\boldsymbol{v}}_o = \\left(\\boldsymbol{v}_{o_1} + \\ldots + \\boldsymbol{v}_{o_{2m}} \\right)/(2m)$. The equation above can be simplified as\n",
    "\n",
    "$$\\mathbb{P}(w_c \\mid \\mathcal{W}_o) = \\frac{\\exp\\left(\\boldsymbol{u}_c^\\top \\bar{\\boldsymbol{v}}_o\\right)}{\\sum_{i \\in \\mathcal{V}} \\exp\\left(\\boldsymbol{u}_i^\\top \\bar{\\boldsymbol{v}}_o\\right)}.$$\n",
    "\n",
    "Given a text sequence of length $T$, we assume that the word at time step $t$ is $w^{(t)}$, and the context window size is $m$.  The likelihood function of the CBOW model is the probability of generating any central target word from the context words.\n",
    "\n",
    "$$ \\prod_{t=1}^{T}  \\mathbb{P}(w^{(t)} \\mid  w^{(t-m)}, \\ldots,  w^{(t-1)},  w^{(t+1)}, \\ldots,  w^{(t+m)}).$$\n",
    "\n",
    " Unlike the skip-gram model, we usually use the context word vector as the representation vector for a word in the CBOW model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "collapsed": true,
    "id": "6swivv2wA9r2"
   },
   "source": [
    "# Approximate Training\n",
    "\n",
    "The core feature of the skip-gram model is the use of softmax operations to compute the conditional probability of generating context word $w_o$ based on the given central target word $w_c$.\n",
    "\n",
    "$$\\mathbb{P}(w_o \\mid w_c) = \\frac{\\text{exp}(\\boldsymbol{u}_o^\\top \\boldsymbol{v}_c)}{ \\sum_{i \\in \\mathcal{V}} \\text{exp}(\\boldsymbol{u}_i^\\top \\boldsymbol{v}_c)}.$$\n",
    "\n",
    "The logarithmic loss corresponding to the conditional probability is given as\n",
    "\n",
    "$$-\\log \\mathbb{P}(w_o \\mid w_c) =\n",
    "-\\boldsymbol{u}_o^\\top \\boldsymbol{v}_c + \\log\\left(\\sum_{i \\in \\mathcal{V}} \\text{exp}(\\boldsymbol{u}_i^\\top \\boldsymbol{v}_c)\\right).$$\n",
    "\n",
    "\n",
    "Because the softmax operation has considered that the context word could be any word in the dictionary $\\mathcal{V}$, the loss mentioned above actually includes the sum of the number of items in the dictionary size. Hence, the gradient computation for each step contains the sum of the number of items in the dictionary size. For larger dictionaries with hundreds of thousands or even millions of words, the overhead for computing each gradient may be too high.  In order to reduce such computational complexity, we will introduce an approximate training method in this section: negative sampling. Since there is no major difference between the skip-gram model and the CBOW model, we will only use the skip-gram model as an example to introduce these two training methods in this section.\n",
    "\n",
    "\n",
    "\n",
    "## Negative Sampling\n",
    "\n",
    "Negative sampling modifies the original objective function. Given a context window for the central target word $w_c$, we will treat it as an event for context word $w_o$ to appear in the context window and compute the probability of this event from\n",
    "\n",
    "$$\\mathbb{P}(D=1\\mid w_c, w_o) = \\sigma(\\boldsymbol{u}_o^\\top \\boldsymbol{v}_c),$$\n",
    "\n",
    "Here, the $\\sigma$ function has the same definition as the sigmoid activation function:\n",
    "\n",
    "$$\\sigma(x) = \\frac{1}{1+\\exp(-x)}.$$\n",
    "\n",
    "We will first consider training the word vector by maximizing the joint probability of all events in the text sequence. Given a text sequence of length $T$, we assume that the word at time step $t$ is $w^{(t)}$ and the context window size is $m$. Now we consider maximizing the joint probability\n",
    "\n",
    "$$ \\prod_{t=1}^{T} \\prod_{-m \\leq j \\leq m,\\ j \\neq 0} \\mathbb{P}(D=1\\mid w^{(t)}, w^{(t+j)}).$$\n",
    "\n",
    "However, the events included in the model only consider positive examples. In this case, only when all the word vectors are equal and their values approach infinity can the joint probability above be maximized to 1. Obviously, such word vectors are meaningless. Negative sampling makes the objective function more meaningful by sampling with an addition of negative examples. Assume that event $P$ occurs when context word $w_o$ to appear in the context window of central target word $w_c$, and we sample $K$ words that do not appear in the context window according to the distribution $\\mathbb{P}(w)$ to act as noise words. We assume the event for noise word $w_k$($k=1, \\ldots, K$) to not appear in the context window of central target word $w_c$ is $N_k$. Suppose that events $P and N_1, \\ldots, N_K$ for both positive and negative examples are independent of each other. By considering negative sampling, we can rewrite the joint probability above, which only considers the positive examples, as\n",
    "\n",
    "\n",
    "$$ \\prod_{t=1}^{T} \\prod_{-m \\leq j \\leq m,\\ j \\neq 0} \\mathbb{P}(w^{(t+j)} \\mid w^{(t)}),$$\n",
    "\n",
    "Here, the conditional probability is approximated to be\n",
    "$$ \\mathbb{P}(w^{(t+j)} \\mid w^{(t)}) =\\mathbb{P}(D=1\\mid w^{(t)}, w^{(t+j)})\\prod_{k=1,\\ w_k \\sim \\mathbb{P}(w)}^K \\mathbb{P}(D=0\\mid w^{(t)}, w_k).$$\n",
    "\n",
    "\n",
    "Let the text sequence index of word $w^{(t)}$ at time step $t$ be $i_t$ and $h_k$ for noise word $w_k$ in the dictionary. The logarithmic loss for the conditional probability above is\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "-\\log\\mathbb{P}(w^{(t+j)} \\mid w^{(t)})\n",
    "=& -\\log\\mathbb{P}(D=1\\mid w^{(t)}, w^{(t+j)}) - \\sum_{k=1,\\ w_k \\sim \\mathbb{P}(w)}^K \\log\\mathbb{P}(D=0\\mid w^{(t)}, w_k)\\\\\n",
    "=&-  \\log\\, \\sigma\\left(\\boldsymbol{u}_{i_{t+j}}^\\top \\boldsymbol{v}_{i_t}\\right) - \\sum_{k=1,\\ w_k \\sim \\mathbb{P}(w)}^K \\log\\left(1-\\sigma\\left(\\boldsymbol{u}_{h_k}^\\top \\boldsymbol{v}_{i_t}\\right)\\right)\\\\\n",
    "=&-  \\log\\, \\sigma\\left(\\boldsymbol{u}_{i_{t+j}}^\\top \\boldsymbol{v}_{i_t}\\right) - \\sum_{k=1,\\ w_k \\sim \\mathbb{P}(w)}^K \\log\\sigma\\left(-\\boldsymbol{u}_{h_k}^\\top \\boldsymbol{v}_{i_t}\\right).\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Here, the gradient computation in each step of the training is no longer related to the dictionary size, but linearly related to $K$. When $K$ takes a smaller constant, the negative sampling has a lower computational overhead for each step.\n",
    "\n",
    "For more details see [On word embeddings - Part 2: Approximating the Softmax](https://www.ruder.io/word-embeddings-softmax/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3NaV02ejA9r4"
   },
   "source": [
    "# Implementation of Word2vec\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DTFsOQ8eA9r6"
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xw1fkYcTA9sB"
   },
   "source": [
    "## Process the Data Set\n",
    "\n",
    "Penn Tree Bank (PTB) is a small commonly-used [corpus](https://github.com/tomsercu/lstm/tree/master/data). It takes samples from Wall Street Journal articles and includes training sets, validation sets, and test sets. We will train the word embedding model on the PTB training set. Each line of the data set acts as a sentence. All the words in a sentence are separated by spaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YYN50OYqLzj-"
   },
   "outputs": [],
   "source": [
    "## colab SETUP\n",
    "!mkdir data\n",
    "%cd data\n",
    "!wget https://raw.githubusercontent.com/tomsercu/lstm/master/data/ptb.train.txt\n",
    "ROOT_DIR='content'\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "ROOT_DIR = Path.home()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "v4XKQvI3A9sD"
   },
   "outputs": [],
   "source": [
    "data_path = os.path.join(ROOT_DIR,'data/')\n",
    "file = 'ptb.train.txt'\n",
    "with open(data_path+file, 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    # st is the abbreviation of \"sentence\" in the loop\n",
    "    raw_dataset = [st.split() for st in lines]\n",
    "\n",
    "'# sentences: %d' % len(raw_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ndY8FkLVA9sJ"
   },
   "source": [
    "For the first three sentences of the data set, print the number of words and the first five words of each sentence. The end character of this data set is \"&lt;eos&gt;\", uncommon words are all represented by \"&lt;unk&gt;\", and numbers are replaced with \"N\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "enuzZ5ehA9sL"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nOLu1KYBA9sU"
   },
   "source": [
    "### Create Word Index\n",
    "\n",
    "For the sake of simplicity, we only keep words that appear at least 5 times in the data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "k0PQHu8NA9sW"
   },
   "outputs": [],
   "source": [
    "# tk is an abbreviation for \"token\" in the loop\n",
    "counter = collections.Counter([tk for st in raw_dataset for tk in st])\n",
    "counter = dict(filter(lambda x: x[1] >= 5, counter.items()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-xn9hltgA9sb"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FcqXd38FA9sm"
   },
   "outputs": [],
   "source": [
    "idx_to_token = [tk for tk, _ in counter.items()]\n",
    "token_to_idx = {tk: idx for idx, tk in enumerate(idx_to_token)}\n",
    "dataset = [[token_to_idx[tk] for tk in st if tk in token_to_idx]\n",
    "           for st in raw_dataset]\n",
    "num_tokens = sum([len(st) for st in dataset])\n",
    "'# tokens: %d' % num_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4r8_tQd_A9ss"
   },
   "outputs": [],
   "source": [
    "idx_to_token[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zTSoD5d2A9sz"
   },
   "outputs": [],
   "source": [
    "token_to_idx['consensus']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3a_eZugcA9s6"
   },
   "source": [
    "### Subsampling\n",
    "\n",
    "In text data, there are generally some words that appear at high frequencies, such \"the\", \"a\", and \"in\" in English. Generally speaking, in a context window, it is better to train the word embedding model when a word (such as \"chip\") and a lower-frequency word (such as \"microprocessor\") appear at the same time, rather than when a word appears with a higher-frequency word (such as \"the\"). Therefore, when training the word embedding model, we can perform subsampling[2] on the words. Specifically, each indexed word $w_i$ in the data set will drop out at a certain probability. The dropout probability is given as:\n",
    "\n",
    "$$ \\mathbb{P}(w_i) = \\max\\left(1 - \\sqrt{\\frac{t}{f(w_i)}}, 0\\right),$$\n",
    "\n",
    "Here, $f(w_i)$ is the ratio of the instances of word $w_i$ to the total number of words in the data set, and the constant $t$ is a hyper-parameter (set to $10^{-4}$ in this experiment). As we can see, it is only possible to drop out the word $w_i$ in subsampling when $f(w_i) > t$. The higher the word's frequency, the higher its dropout probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pAz0GiAYA9s7"
   },
   "outputs": [],
   "source": [
    "def discard(idx):\n",
    "    return random.uniform(0, 1) < 1 - math.sqrt(\n",
    "        1e-4 / counter[idx_to_token[idx]] * num_tokens)\n",
    "\n",
    "subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]\n",
    "'# tokens: %d' % sum([len(st) for st in subsampled_dataset])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Jv1nPR-NA9tA"
   },
   "outputs": [],
   "source": [
    "def compare_counts(token):\n",
    "    return '# %s: before=%d, after=%d' % (token, sum(\n",
    "        [st.count(token_to_idx[token]) for st in dataset]), sum(\n",
    "        [st.count(token_to_idx[token]) for st in subsampled_dataset]))\n",
    "\n",
    "compare_counts('the')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MNjCPHXrA9tE"
   },
   "outputs": [],
   "source": [
    "compare_counts('join')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aAsNCgldA9tJ"
   },
   "source": [
    "### Extract Central Target Words and Context Words\n",
    "\n",
    "We use words with a distance from the central target word not exceeding the context window size as the context words of the given center target word. The following definition function extracts all the central target words and their context words. It uniformly and randomly samples an integer to be used as the context window size between integer 1 and the `max_window_size` (maximum context window)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8huGW8DAA9tK"
   },
   "outputs": [],
   "source": [
    "def get_centers_and_contexts(dataset, max_window_size):\n",
    "    centers, contexts = [], []\n",
    "    for st in dataset:\n",
    "        # Each sentence needs at least 2 words to form a\n",
    "        # \"central target word - context word\" pair\n",
    "        if len(st) < 2:\n",
    "            continue\n",
    "        centers += st\n",
    "        for center_i in range(len(st)):\n",
    "            window_size = random.randint(1, max_window_size)\n",
    "            indices = list(range(max(0, center_i - window_size),\n",
    "                                 min(len(st), center_i + 1 + window_size)))\n",
    "            # Exclude the central target word from the context words\n",
    "            indices.remove(center_i)\n",
    "            contexts.append([st[idx] for idx in indices])\n",
    "    return centers, contexts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "N84nUVa-A9tO"
   },
   "outputs": [],
   "source": [
    "tiny_dataset = [list(range(7)), list(range(7, 10))]\n",
    "print('dataset', tiny_dataset)\n",
    "for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):\n",
    "    print('center', center, 'has contexts', context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0RK2JtgpA9tY"
   },
   "source": [
    "In the experiment, we set the maximum context window size to 5. The following extracts all the central target words and their context words in the data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Y2k9PtoqA9tZ"
   },
   "outputs": [],
   "source": [
    "all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "F4M3c2pbA9tf"
   },
   "source": [
    "## Negative Sampling\n",
    "\n",
    "We use negative sampling for approximate training. For a central and context word pair, we randomly sample $K$ noise words ($K=5$ in the experiment). According to the suggestion in the Word2vec paper, the noise word sampling probability $\\mathbb{P}(w)$ is the ratio of the word frequency of $w$ to the total word frequency raised to the power of 0.75."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Pfd8ukxhA9tl"
   },
   "outputs": [],
   "source": [
    "def get_negatives(all_contexts, sampling_weights, K):\n",
    "    all_negatives, neg_candidates, i = [], [], 0\n",
    "    population = list(range(len(sampling_weights)))\n",
    "    for contexts in all_contexts:\n",
    "        negatives = []\n",
    "        while len(negatives) < len(contexts) * K:\n",
    "            if i == len(neg_candidates):\n",
    "                # An index of k words is randomly generated as noise words\n",
    "                # based on the weight of each word (sampling_weights). For\n",
    "                # efficient calculation, k can be set slightly larger\n",
    "                i, neg_candidates = 0, random.choices(population, sampling_weights, k=int(1e5))\n",
    "            neg, i = neg_candidates[i], i + 1\n",
    "            # Noise words cannot be context words\n",
    "            if neg not in set(contexts):\n",
    "                negatives.append(neg)\n",
    "        all_negatives.append(negatives)\n",
    "    return all_negatives\n",
    "\n",
    "sampling_weights = [counter[w]**0.75 for w in idx_to_token]\n",
    "all_negatives = get_negatives(all_contexts, sampling_weights, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vRYt2t8wA9tq"
   },
   "outputs": [],
   "source": [
    "all_negatives[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Nd-w6UFdA9tz"
   },
   "outputs": [],
   "source": [
    "all_contexts[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xCAy09AqA9t2"
   },
   "source": [
    "## Reading Data\n",
    "\n",
    "We extracted all central target words `all_centers`, and the context words `all_contexts` and noise words `all_negatives` of each central target word from the data set. We will read them in random mini-batches.\n",
    "\n",
    "In a mini-batch of data, the $i$-th example includes a central word and its corresponding $n_i$ context words and $m_i$ noise words. Since the context window size of each example may be different, the sum of context words and noise words, $n_i+m_i$, will be different. When constructing a mini-batch, we concatenate the context words and noise words of each example, and add 0s for padding until the length of the concatenations are the same, that is, the length of all concatenations is $\\max_i n_i+m_i$(`max_len`). In order to avoid the effect of padding on the loss function calculation, we construct the mask variable `masks`, each element of which corresponds to an element in the concatenation of context and noise words, `contexts_negatives`. When an element in the variable `contexts_negatives` is a padding, the element in the mask variable `masks` at the same position will be 0. Otherwise, it takes the value 1. In order to distinguish between positive and negative examples, we also need to distinguish the context words from the noise words in the `contexts_negatives` variable. Based on the construction of the mask variable, we only need to create a label variable `labels` with the same shape as the `contexts_negatives` variable and set the elements corresponding to context words (positive examples) to 1, and the rest to 0.\n",
    "\n",
    "Next, we will implement the mini-batch reading function `batchify`. Its mini-batch input `data` is a list whose length is the batch size, each element of which contains central target words `center`, context words `context`, and noise words `negative`. The mini-batch data returned by this function conforms to the format we need, for example, it includes the mask variable. We wrap this function in a `Dataset` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uJs9ymetA9t3"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "t0GFenm_A9t7"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PuheXdUEA9t_"
   },
   "outputs": [],
   "source": [
    "class PTB_dataset(Dataset):\n",
    "    \n",
    "    def __init__(self, all_centers, all_contexts, all_negatives):\n",
    "        self.all_centers, self.all_contexts_negatives, self.all_masks, self.all_labels = self.batchify(list(zip(all_centers,all_contexts,all_negatives)))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.all_centers)\n",
    "    \n",
    "    def __getitem__(self,idx):\n",
    "        return self.all_centers[idx], self.all_contexts_negatives[idx], self.all_masks[idx], self.all_labels[idx]\n",
    "        \n",
    "    def batchify(self,data):\n",
    "        max_len = max(len(c) + len(n) for _, c, n in data)\n",
    "        centers, contexts_negatives, masks, labels = [], [], [], []\n",
    "        for center, context, negative in data:\n",
    "            cur_len = len(context) + len(negative)\n",
    "            centers += [center]\n",
    "            contexts_negatives += [context + negative + [0] * (max_len - cur_len)]\n",
    "            masks += [[1] * cur_len + [0] * (max_len - cur_len)]\n",
    "            labels += [[1] * len(context) + [0] * (max_len - len(context))]\n",
    "        return (torch.tensor(centers).view((-1, 1)), torch.tensor(np.array(contexts_negatives)),\n",
    "            torch.tensor(np.array(masks)), torch.tensor(np.array(labels)))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rjAdHiQTA9uC"
   },
   "outputs": [],
   "source": [
    "ptbdata = PTB_dataset(all_centers, all_contexts, all_negatives)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iQaI870eA9uE"
   },
   "outputs": [],
   "source": [
    "ptbdata[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VIIjdClMA9uH"
   },
   "source": [
    "And here is our dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mMUFDG2LA9uI"
   },
   "outputs": [],
   "source": [
    "batch_size = 512\n",
    "\n",
    "data_loader = DataLoader(ptbdata, batch_size, shuffle=True,\n",
    "                              num_workers=4)\n",
    "for batch in data_loader:\n",
    "    for name, data in zip(['centers', 'contexts_negatives', 'masks',\n",
    "                           'labels'], batch):\n",
    "        print(name, 'shape:', data.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A4YEjWTSA9uL"
   },
   "source": [
    "## The Skip-Gram Model\n",
    "\n",
    "You will implement the skip-gram model by using embedding layers and mini-batch multiplication [`torch.einsum`](https://pytorch.org/docs/stable/torch.html#torch.bmm). These methods are also often used to implement other natural language processing applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nLz15ca5A9uM"
   },
   "outputs": [],
   "source": [
    "# taken from the spotlight library, see\n",
    "# https://github.com/maciejkula/spotlight/blob/master/spotlight/layers.py\n",
    "class ScaledEmbedding(nn.Embedding):\n",
    "    \"\"\"\n",
    "    Embedding layer that initialises its values\n",
    "    to using a normal variable scaled by the inverse\n",
    "    of the emedding dimension.\n",
    "    \"\"\"\n",
    "    def reset_parameters(self):\n",
    "        \"\"\"\n",
    "        Initialize parameters.\n",
    "        \"\"\"\n",
    "        self.weight.data.normal_(0, 1.0 / self.embedding_dim)\n",
    "        if self.padding_idx is not None:\n",
    "            self.weight.data[self.padding_idx].fill_(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "o7bS-dFbA9uP"
   },
   "source": [
    "### Skip-gram Model Forward Calculation\n",
    "\n",
    "In forward calculation, the input of the skip-gram model contains the central target word index `center` and the concatenated context and noise word index `contexts_and_negatives`. In which, the `center` variable has the shape (batch size, 1), while the `contexts_and_negatives` variable has the shape (batch size, `max_len`). These two variables are first transformed from word indexes to word vectors by the word embedding layer, and then the output of shape (batch size, 1, `max_len`) is obtained by mini-batch multiplication. Each element in the output is the inner product of the central target word vector and the context word vector or noise word vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CTOeSlqmA9uP"
   },
   "outputs": [],
   "source": [
    "class Skip_gram(nn.Module):\n",
    "    def __init__(self, input_dim, embed_size = 100):\n",
    "        super(Skip_gram, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.embed_size = embed_size\n",
    "        self.central_emb = ScaledEmbedding(self.input_dim,self.embed_size)\n",
    "        self.context_emb = ScaledEmbedding(self.input_dim,self.embed_size)\n",
    "        \n",
    "    def forward(self, icent, icont):\n",
    "        #\n",
    "        # your code here (hint: dimensions are for icent (bs,1) for icont (bs,max_len) and output (bs,1,max_len))\n",
    "        #"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "z0Z7NPX1A9uT"
   },
   "outputs": [],
   "source": [
    "net = Skip_gram(len(idx_to_token))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xgSDQ6TeA9uV"
   },
   "outputs": [],
   "source": [
    "net(torch.tensor([0,1,3]).unsqueeze(1),torch.tensor([[0,0],[0,0],[0,0]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tuYrt-x4A9uY"
   },
   "source": [
    "### Loss function\n",
    "\n",
    "It is worth mentioning that we need use the mask variable to specify the partial predicted value and label that participate in loss function calculation in the mini-batch: when the mask is 1, the predicted value and label of the corresponding position will participate in the calculation of the loss function; When the mask is 0, the predicted value and label of the corresponding position do not participate in the calculation of the loss function. As we mentioned earlier, mask variables can be used to avoid the effect of padding on loss function calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Cb2V_wIEA9uY"
   },
   "outputs": [],
   "source": [
    "loss_fn = nn.BCEWithLogitsLoss(reduction='none')\n",
    "def criterion(pred, label, mask):\n",
    "    #\n",
    "    # your code here\n",
    "    #"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "J8OJMelWA9ua"
   },
   "outputs": [],
   "source": [
    "pred = torch.tensor([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]])\n",
    "# 1 and 0 in the label variables label represent context words and the noise\n",
    "# words, respectively\n",
    "label = torch.tensor([[1, 0, 0, 0], [1, 1, 0, 0]]).type(torch.FloatTensor)\n",
    "mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]]).type(torch.FloatTensor)  # Mask variable\n",
    "\n",
    "criterion(pred,label,mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vfx1o6c1A9ud"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(),lr=0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WONDLY8IA9uf"
   },
   "outputs": [],
   "source": [
    "def train(n_epochs):\n",
    "    \n",
    "    for epoch in range(n_epochs):\n",
    "        start, loss = time.time(), 0\n",
    "        for batch in data_loader:\n",
    "            #\n",
    "            # your code here\n",
    "            #\n",
    "            \n",
    "        print('epoch %d, loss %.2f, time %.2fs'\n",
    "              % (epoch + 1, loss, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sCSPXuRjA9ug"
   },
   "outputs": [],
   "source": [
    "train(6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fux1T7u-A9uj"
   },
   "source": [
    "## Applying the Word Embedding Model\n",
    "\n",
    "After training the word embedding model, we can represent similarity in meaning between words based on the cosine similarity of two word vectors. As we can see, when using the trained word embedding model, the words closest in meaning to the word \"chip\" are mostly related to chips."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jPQZOcdGA9uk"
   },
   "outputs": [],
   "source": [
    "def get_similar_tokens(query_token, k, W):\n",
    "    #\n",
    "    # your code here\n",
    "    #\n",
    "    \n",
    "    _,topk = torch.topk(cos, k=k+1,)\n",
    "    for i in topk[1:]:# Remove the input words\n",
    "        print('cosine sim=%.3f: %s' % (cos[i], (idx_to_token[i])))\n",
    "\n",
    "get_similar_tokens('chip', 3, net.central_emb.weight.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dWz6Cno3A9um"
   },
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": false,
   "name": "08_Word2vec_pytorch_empty.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "dldiy",
   "language": "python",
   "name": "dldiy"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}