{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-1 fw=all -->\n",
    "\n",
    "# Word Similarity and Analogy\n",
    "\n",
    "In that section, \n",
    "we trained a word2vec model on a small dataset, \n",
    "and applied it\n",
    "to find semantically similar words \n",
    "for an input word.\n",
    "In practice,\n",
    "word vectors that are pretrained\n",
    "on large corpora can be\n",
    "applied to downstream\n",
    "natural language processing tasks,\n",
    "which will be covered later\n",
    "in that section.\n",
    "To demonstrate \n",
    "semantics of pretrained word vectors\n",
    "from large corpora in a straightforward way,\n",
    "let's apply them\n",
    "in the word similarity and analogy tasks."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-and-analogy"
   },
   "source": [
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import os"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-2 fw=all -->\n",
    "\n",
    "## Loading Pretrained Word Vectors\n",
    "\n",
    "Below lists pretrained GloVe embeddings of dimension 50, 100, and 300,\n",
    "which can be downloaded from the [GloVe website](https://nlp.stanford.edu/projects/glove/).\n",
    "The pretrained fastText embeddings are available in multiple languages.\n",
    "Here we consider one English version (300-dimensional \"wiki.en\") that can be downloaded from the\n",
    "[fastText website](https://fasttext.cc/)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-loading-pretrained-word-vectors-1"
   },
   "source": [
    "\n",
    "d2l.DATA_HUB['glove.6b.50d'] = (d2l.DATA_URL + 'glove.6B.50d.zip',\n",
    "                                '0b8703943ccdb6eb788e6f091b8946e82231bc4d')\n",
    "\n",
    "\n",
    "d2l.DATA_HUB['glove.6b.100d'] = (d2l.DATA_URL + 'glove.6B.100d.zip',\n",
    "                                 'cd43bfb07e44e6f27cbcc7bc9ae3d80284fdaf5a')\n",
    "\n",
    "\n",
    "d2l.DATA_HUB['glove.42b.300d'] = (d2l.DATA_URL + 'glove.42B.300d.zip',\n",
    "                                  'b5116e234e9eb9076672cfeabf5469f3eec904fa')\n",
    "\n",
    "\n",
    "d2l.DATA_HUB['wiki.en'] = (d2l.DATA_URL + 'wiki.en.zip',\n",
    "                           'c1816da3821ae9f43899be655002f6c723e91b88')"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-3 fw=all -->\n",
    "\n",
    "To load these pretrained GloVe and fastText embeddings, we define the following `TokenEmbedding` class."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-loading-pretrained-word-vectors-2"
   },
   "source": [
    "\n",
    "class TokenEmbedding:\n",
    "    \"\"\"Token Embedding.\"\"\"\n",
    "    def __init__(self, embedding_name):\n",
    "        self.idx_to_token, self.idx_to_vec = self._load_embedding(\n",
    "            embedding_name)\n",
    "        self.unknown_idx = 0\n",
    "        self.token_to_idx = {token: idx for idx, token in\n",
    "                             enumerate(self.idx_to_token)}\n",
    "\n",
    "    def _load_embedding(self, embedding_name):\n",
    "        idx_to_token, idx_to_vec = ['<unk>'], []\n",
    "        data_dir = d2l.download_extract(embedding_name)\n",
    "        # GloVe website: https://nlp.stanford.edu/projects/glove/\n",
    "        # fastText website: https://fasttext.cc/\n",
    "        with open(os.path.join(data_dir, 'vec.txt'), 'r') as f:\n",
    "            for line in f:\n",
    "                elems = line.rstrip().split(' ')\n",
    "                token, elems = elems[0], [float(elem) for elem in elems[1:]]\n",
    "                # Skip header information, such as the top row in fastText\n",
    "                if len(elems) > 1:\n",
    "                    idx_to_token.append(token)\n",
    "                    idx_to_vec.append(elems)\n",
    "        idx_to_vec = [[0] * len(idx_to_vec[0])] + idx_to_vec\n",
    "        return idx_to_token, d2l.tensor(idx_to_vec)\n",
    "\n",
    "    def __getitem__(self, tokens):\n",
    "        indices = [self.token_to_idx.get(token, self.unknown_idx)\n",
    "                   for token in tokens]\n",
    "        vecs = self.idx_to_vec[d2l.tensor(indices)]\n",
    "        return vecs\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.idx_to_token)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-4 fw=all -->\n",
    "\n",
    "Below we load the\n",
    "50-dimensional GloVe embeddings\n",
    "(pretrained on a Wikipedia subset).\n",
    "When creating the `TokenEmbedding` instance,\n",
    "the specified embedding file has to be downloaded if it\n",
    "was not yet."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-loading-pretrained-word-vectors-3"
   },
   "source": [
    "glove_6b50d = TokenEmbedding('glove.6b.50d')"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-5 fw=all -->\n",
    "\n",
    "Output the vocabulary size. The vocabulary contains 400000 words (tokens) and a special unknown token."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-loading-pretrained-word-vectors-4"
   },
   "source": [
    "len(glove_6b50d)"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "400001"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-6 fw=all -->\n",
    "\n",
    "We can get the index of a word in the vocabulary, and vice versa."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-loading-pretrained-word-vectors-5"
   },
   "source": [
    "glove_6b50d.token_to_idx['beautiful'], glove_6b50d.idx_to_token[3367]"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(3367, 'beautiful')"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-7 fw=all -->\n",
    "\n",
    "## Applying Pretrained Word Vectors\n",
    "\n",
    "Using the loaded GloVe vectors,\n",
    "we will demonstrate their semantics\n",
    "by applying them\n",
    "in the following word similarity and analogy tasks.\n",
    "\n",
    "\n",
    "### Word Similarity\n",
    "\n",
    "Similar to that section,\n",
    "in order to find semantically similar words\n",
    "for an input word\n",
    "based on cosine similarities between\n",
    "word vectors,\n",
    "we implement the following `knn`\n",
    "($k$-nearest neighbors) function."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-1"
   },
   "source": [
    "def knn(W, x, k):\n",
    "    # Add 1e-9 for numerical stability\n",
    "    cos = jnp.dot(W, x.reshape(-1,)) / (\n",
    "        jnp.sqrt(jnp.sum(W * W, axis=1) + 1e-9) *\n",
    "        jnp.sqrt((x * x).sum()))\n",
    "    topk = jnp.argsort(-cos)[:k]\n",
    "    return topk, [cos[int(i)] for i in topk]"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-8 fw=all -->\n",
    "\n",
    "Then, we \n",
    "search for similar words\n",
    "using the pretrained word vectors \n",
    "from the `TokenEmbedding` instance `embed`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-2"
   },
   "source": [
    "def get_similar_tokens(query_token, k, embed):\n",
    "    topk, cos = knn(embed.idx_to_vec, embed[[query_token]], k + 1)\n",
    "    for i, c in zip(topk[1:], cos[1:]):  # Exclude the input word\n",
    "        print(f'cosine sim={float(c):.3f}: {embed.idx_to_token[int(i)]}')"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-9 fw=all -->\n",
    "\n",
    "The vocabulary of the pretrained word vectors\n",
    "in `glove_6b50d` contains 400000 words and a special unknown token. \n",
    "Excluding the input word and unknown token,\n",
    "among this vocabulary\n",
    "let's find \n",
    "three most semantically similar words\n",
    "to word \"chip\"."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-3"
   },
   "source": [
    "get_similar_tokens('chip', 3, glove_6b50d)"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cosine sim=0.856: chips\ncosine sim=0.749: intel\ncosine sim=0.749: electronics\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-10 fw=all -->\n",
    "\n",
    "Below outputs similar words\n",
    "to \"baby\" and \"beautiful\"."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-4"
   },
   "source": [
    "get_similar_tokens('baby', 3, glove_6b50d)"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cosine sim=0.839: babies\ncosine sim=0.800: boy\ncosine sim=0.792: girl\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-similarity-5"
   },
   "source": [
    "get_similar_tokens('beautiful', 3, glove_6b50d)"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cosine sim=0.921: lovely\ncosine sim=0.893: gorgeous\ncosine sim=0.830: wonderful\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-11 fw=all -->\n",
    "\n",
    "### Word Analogy\n",
    "\n",
    "Besides finding similar words,\n",
    "we can also apply word vectors\n",
    "to word analogy tasks.\n",
    "For example,\n",
    "“man”:“woman”::“son”:“daughter”\n",
    "is the form of a word analogy:\n",
    "“man” is to “woman” as “son” is to “daughter”.\n",
    "Specifically,\n",
    "the word analogy completion task\n",
    "can be defined as:\n",
    "for a word analogy \n",
    "$a : b :: c : d$, given the first three words $a$, $b$ and $c$, find $d$. \n",
    "Denote the vector of word $w$ by $\\textrm{vec}(w)$. \n",
    "To complete the analogy,\n",
    "we will find the word \n",
    "whose vector is most similar\n",
    "to the result of $\\textrm{vec}(c)+\\textrm{vec}(b)-\\textrm{vec}(a)$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-analogy-1"
   },
   "source": [
    "def get_analogy(token_a, token_b, token_c, embed):\n",
    "    vecs = embed[[token_a, token_b, token_c]]\n",
    "    x = vecs[1] - vecs[0] + vecs[2]\n",
    "    topk, cos = knn(embed.idx_to_vec, x, 1)\n",
    "    return embed.idx_to_token[int(topk[0])]  # Remove unknown words"
   ],
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-12 fw=all -->\n",
    "\n",
    "Let's verify the \"male-female\" analogy using the loaded word vectors."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-analogy-2"
   },
   "source": [
    "get_analogy('man', 'woman', 'son', glove_6b50d)"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'daughter'"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-13 fw=all -->\n",
    "\n",
    "Below completes a\n",
    "“capital-country” analogy: \n",
    "“beijing”:“china”::“tokyo”:“japan”.\n",
    "This demonstrates \n",
    "semantics in the pretrained word vectors."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-analogy-3"
   },
   "source": [
    "get_analogy('beijing', 'china', 'tokyo', glove_6b50d)"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'japan'"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-14 fw=all -->\n",
    "\n",
    "For the\n",
    "“adjective-superlative adjective” analogy\n",
    "such as \n",
    "“bad”:“worst”::“big”:“biggest”,\n",
    "we can see that the pretrained word vectors\n",
    "may capture the syntactic information."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-analogy-4"
   },
   "source": [
    "get_analogy('bad', 'worst', 'big', glove_6b50d)"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'biggest'"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-15 fw=all -->\n",
    "\n",
    "To show the captured notion\n",
    "of past tense in the pretrained word vectors,\n",
    "we can test the syntax using the\n",
    "\"present tense-past tense\" analogy: “do”:“did”::“go”:“went”."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "similarity-analogy-word-analogy-5"
   },
   "source": [
    "get_analogy('do', 'did', 'go', glove_6b50d)"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'went'"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-16 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* In practice, word vectors that are pretrained on large corpora can be applied to downstream natural language processing tasks.\n",
    "* Pretrained word vectors can be applied to the word similarity and analogy tasks.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Test the fastText results using `TokenEmbedding('wiki.en')`.\n",
    "1. When the vocabulary is extremely large, how can we find similar words or complete a word analogy faster?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=similarity-analogy-md-17 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1336)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
