{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-1 fw=all -->\n",
    "\n",
    "# Subword Embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-2 fw=all -->\n",
    "\n",
    "In English,\n",
    "words such as\n",
    "\"helps\", \"helped\", and \"helping\" are \n",
    "inflected forms of the same word \"help\".\n",
    "The relationship \n",
    "between \"dog\" and \"dogs\"\n",
    "is the same as \n",
    "that between \"cat\" and \"cats\",\n",
    "and \n",
    "the relationship \n",
    "between \"boy\" and \"boyfriend\"\n",
    "is the same as \n",
    "that between \"girl\" and \"girlfriend\".\n",
    "In other languages\n",
    "such as French and Spanish,\n",
    "many verbs have over 40 inflected forms,\n",
    "while in Finnish,\n",
    "a noun may have up to 15 cases.\n",
    "In linguistics,\n",
    "morphology studies word formation and word relationships.\n",
    "However,\n",
    "the internal structure of words\n",
    "was neither explored in word2vec\n",
    "nor in GloVe.\n",
    "\n",
    "## The fastText Model\n",
    "\n",
    "Recall how words are represented in word2vec.\n",
    "In both the skip-gram model\n",
    "and the continuous bag-of-words model,\n",
    "different inflected forms of the same word\n",
    "are directly represented by different vectors\n",
    "without shared parameters.\n",
    "To use morphological information,\n",
    "the *fastText* model\n",
    "proposed a *subword embedding* approach,\n",
    "where a subword is a character $n$-gram [@Bojanowski.Grave.Joulin.ea.2017].\n",
    "Instead of learning word-level vector representations,\n",
    "fastText can be considered as\n",
    "the subword-level skip-gram,\n",
    "where each *center word* is represented by the sum of \n",
    "its subword vectors.\n",
    "\n",
    "Let's illustrate how to obtain \n",
    "subwords for each center word in fastText\n",
    "using the word \"where\".\n",
    "First, add special characters “&lt;” and “&gt;” \n",
    "at the beginning and end of the word to distinguish prefixes and suffixes from other subwords. \n",
    "Then, extract character $n$-grams from the word.\n",
    "For example, when $n=3$,\n",
    "we obtain all subwords of length 3: \"&lt;wh\", \"whe\", \"her\", \"ere\", \"re&gt;\", and the special subword \"&lt;where&gt;\".\n",
    "\n",
    "\n",
    "In fastText, for any word $w$,\n",
    "denote by $\\mathcal{G}_w$\n",
    "the union of all its subwords of length between 3 and 6\n",
    "and its special subword.\n",
    "The vocabulary \n",
    "is the union of the subwords of all words.\n",
    "Letting $\\mathbf{z}_g$\n",
    "be the vector of subword $g$ in the dictionary,\n",
    "the vector $\\mathbf{v}_w$ for \n",
    "word $w$ as a center word\n",
    "in the skip-gram model\n",
    "is the sum of its subword vectors:\n",
    "\n",
    "$$\\mathbf{v}_w = \\sum_{g\\in\\mathcal{G}_w} \\mathbf{z}_g.$$\n",
    "\n",
    "The rest of fastText is the same as the skip-gram model. Compared with the skip-gram model, \n",
    "the vocabulary in fastText is larger,\n",
    "resulting in more model parameters. \n",
    "Besides, \n",
    "to calculate the representation of a word,\n",
    "all its subword vectors\n",
    "have to be summed,\n",
    "leading to higher computational complexity.\n",
    "However,\n",
    "thanks to shared parameters from subwords among words with similar structures,\n",
    "rare words and even out-of-vocabulary words\n",
    "may obtain better vector representations in fastText.\n",
    "\n",
    "\n",
    "## Byte Pair Encoding\n",
    "\n",
    "In fastText, all the extracted subwords have to be of the specified lengths, such as $3$ to $6$, thus the vocabulary size cannot be predefined.\n",
    "To allow for variable-length subwords in a fixed-size vocabulary,\n",
    "we can apply a compression algorithm\n",
    "called *byte pair encoding* (BPE) to extract subwords [@Sennrich.Haddow.Birch.2015].\n",
    "\n",
    "Byte pair encoding performs a statistical analysis of the training dataset to discover common symbols within a word,\n",
    "such as consecutive characters of arbitrary length.\n",
    "Starting from symbols of length 1,\n",
    "byte pair encoding iteratively merges the most frequent pair of consecutive symbols to produce new longer symbols.\n",
    "Note that for efficiency, pairs crossing word boundaries are not considered.\n",
    "In the end, we can use such symbols as subwords to segment words.\n",
    "Byte pair encoding and its variants has been used for input representations in popular natural language processing pretraining models such as GPT-2 [@Radford.Wu.Child.ea.2019] and RoBERTa [@Liu.Ott.Goyal.ea.2019].\n",
    "In the following, we will illustrate how byte pair encoding works.\n",
    "\n",
    "First, we initialize the vocabulary of symbols as all the English lowercase characters, a special end-of-word symbol `'_'`, and a special unknown symbol `'[UNK]'`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-1"
   },
   "source": [
    "import collections\n",
    "\n",
    "symbols = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',\n",
    "           'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',\n",
    "           '_', '[UNK]']"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-3 fw=all -->\n",
    "\n",
    "Since we do not consider symbol pairs that cross boundaries of words,\n",
    "we only need a dictionary `raw_token_freqs` that maps words to their frequencies (number of occurrences)\n",
    "in a dataset.\n",
    "Note that the special symbol `'_'` is appended to each word so that\n",
    "we can easily recover a word sequence (e.g., \"a taller man\")\n",
    "from a sequence of output symbols ( e.g., \"a_ tall er_ man\").\n",
    "Since we start the merging process from a vocabulary of only single characters and special symbols, space is inserted between every pair of consecutive characters within each word (keys of the dictionary `token_freqs`).\n",
    "In other words, space is the delimiter between symbols within a word."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-2"
   },
   "source": [
    "raw_token_freqs = {'fast_': 4, 'faster_': 3, 'tall_': 5, 'taller_': 4}\n",
    "token_freqs = {}\n",
    "for token, freq in raw_token_freqs.items():\n",
    "    token_freqs[' '.join(list(token))] = raw_token_freqs[token]\n",
    "token_freqs"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'f a s t _': 4, 'f a s t e r _': 3, 't a l l _': 5, 't a l l e r _': 4}"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-4 fw=all -->\n",
    "\n",
    "We define the following `get_max_freq_pair` function that\n",
    "returns the most frequent pair of consecutive symbols within a word,\n",
    "where words come from keys of the input dictionary `token_freqs`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-3"
   },
   "source": [
    "def get_max_freq_pair(token_freqs):\n",
    "    pairs = collections.defaultdict(int)\n",
    "    for token, freq in token_freqs.items():\n",
    "        symbols = token.split()\n",
    "        for i in range(len(symbols) - 1):\n",
    "            # Key of `pairs` is a tuple of two consecutive symbols\n",
    "            pairs[symbols[i], symbols[i + 1]] += freq\n",
    "    return max(pairs, key=pairs.get)  # Key of `pairs` with the max value"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-5 fw=all -->\n",
    "\n",
    "As a greedy approach based on frequency of consecutive symbols,\n",
    "byte pair encoding will use the following `merge_symbols` function to merge the most frequent pair of consecutive symbols to produce new symbols."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-4"
   },
   "source": [
    "def merge_symbols(max_freq_pair, token_freqs, symbols):\n",
    "    symbols.append(''.join(max_freq_pair))\n",
    "    new_token_freqs = dict()\n",
    "    for token, freq in token_freqs.items():\n",
    "        new_token = token.replace(' '.join(max_freq_pair),\n",
    "                                  ''.join(max_freq_pair))\n",
    "        new_token_freqs[new_token] = token_freqs[token]\n",
    "    return new_token_freqs"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-6 fw=all -->\n",
    "\n",
    "Now we iteratively perform the byte pair encoding algorithm over the keys of the dictionary `token_freqs`. In the first iteration, the most frequent pair of consecutive symbols are `'t'` and `'a'`, thus byte pair encoding merges them to produce a new symbol `'ta'`. In the second iteration, byte pair encoding continues to merge `'ta'` and `'l'` to result in another new symbol `'tal'`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-5"
   },
   "source": [
    "num_merges = 10\n",
    "for i in range(num_merges):\n",
    "    max_freq_pair = get_max_freq_pair(token_freqs)\n",
    "    token_freqs = merge_symbols(max_freq_pair, token_freqs, symbols)\n",
    "    print(f'merge #{i + 1}:', max_freq_pair)"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "merge #1: ('t', 'a')\nmerge #2: ('ta', 'l')\nmerge #3: ('tal', 'l')\nmerge #4: ('f', 'a')\nmerge #5: ('fa', 's')\nmerge #6: ('fas', 't')\nmerge #7: ('e', 'r')\nmerge #8: ('er', '_')\nmerge #9: ('tall', '_')\nmerge #10: ('fast', '_')\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-7 fw=all -->\n",
    "\n",
    "After 10 iterations of byte pair encoding, we can see that list `symbols` now contains 10 more symbols that are iteratively merged from other symbols."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-6"
   },
   "source": [
    "print(symbols)"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '_', '[UNK]', 'ta', 'tal', 'tall', 'fa', 'fas', 'fast', 'er', 'er_', 'tall_', 'fast_']\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-8 fw=all -->\n",
    "\n",
    "For the same dataset specified in the keys of the dictionary `raw_token_freqs`,\n",
    "each word in the dataset is now segmented by subwords \"fast_\", \"fast\", \"er_\", \"tall_\", and \"tall\"\n",
    "as a result of the byte pair encoding algorithm.\n",
    "For instance, words \"faster_\" and \"taller_\" are segmented as \"fast er_\" and \"tall er_\", respectively."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-7"
   },
   "source": [
    "print(list(token_freqs.keys()))"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "['fast_', 'fast er_', 'tall_', 'tall er_']\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-9 fw=all -->\n",
    "\n",
    "Note that the result of byte pair encoding depends on the dataset being used.\n",
    "We can also use the subwords learned from one dataset\n",
    "to segment words of another dataset.\n",
    "As a greedy approach, the following `segment_BPE` function tries to break words into the longest possible subwords from the input argument `symbols`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-8"
   },
   "source": [
    "def segment_BPE(tokens, symbols):\n",
    "    outputs = []\n",
    "    for token in tokens:\n",
    "        start, end = 0, len(token)\n",
    "        cur_output = []\n",
    "        # Segment token with the longest possible subwords from symbols\n",
    "        while start < len(token) and start < end:\n",
    "            if token[start: end] in symbols:\n",
    "                cur_output.append(token[start: end])\n",
    "                start = end\n",
    "                end = len(token)\n",
    "            else:\n",
    "                end -= 1\n",
    "        if start < len(token):\n",
    "            cur_output.append('[UNK]')\n",
    "        outputs.append(' '.join(cur_output))\n",
    "    return outputs"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-10 fw=all -->\n",
    "\n",
    "In the following, we use the subwords in list `symbols`, which is learned from the aforementioned dataset,\n",
    "to segment `tokens` that represent another dataset."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "subword-embedding-byte-pair-encoding-9"
   },
   "source": [
    "tokens = ['tallest_', 'fatter_']\n",
    "print(segment_BPE(tokens, symbols))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "['tall e s t _', 'fa t t er_']\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-11 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* The fastText model proposes a subword embedding approach. Based on the skip-gram model in word2vec, it represents a center word as the sum of its subword vectors.\n",
    "* Byte pair encoding performs a statistical analysis of the training dataset to discover common symbols within a word. As a greedy approach, byte pair encoding iteratively merges the most frequent pair of consecutive symbols.\n",
    "* Subword embedding may improve the quality of representations of rare words and out-of-dictionary words.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. As an example, there are about $3\\times 10^8$ possible  $6$-grams in English. What is the issue when there are too many subwords? How to address the issue? Hint: refer to the end of Section 3.2 of the fastText paper [@Bojanowski.Grave.Joulin.ea.2017].\n",
    "1. How to design a subword embedding model based on the continuous bag-of-words model?\n",
    "1. To get a vocabulary of size $m$, how many merging operations are needed when the initial symbol vocabulary size is $n$?\n",
    "1. How to extend the idea of byte pair encoding to extract phrases?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=subword-embedding-md-12 fw=tensorflow,jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/4587)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
