{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp nlp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NLP\n",
    "\n",
    "> API details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from collections import Counter\n",
    "from functools import partial\n",
    "import multiprocessing\n",
    "import numpy as np\n",
    "from sklearn.decomposition import PCA\n",
    "import spacy\n",
    "from textblob import TextBlob\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from htools import save, load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only needed for testing.\n",
    "import pandas as pd\n",
    "from string import ascii_lowercase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nonsense sample text.\n",
    "text = [\n",
    "    f\"Row {i}: I went, yesterday; she wasn't here after school? Today. --2\"\n",
    "    for i in range(25_000)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>a</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>24995</td>\n",
       "      <td>Row 24995: I went, yesterday; she wasn't here ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24996</td>\n",
       "      <td>Row 24996: I went, yesterday; she wasn't here ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24997</td>\n",
       "      <td>Row 24997: I went, yesterday; she wasn't here ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24998</td>\n",
       "      <td>Row 24998: I went, yesterday; she wasn't here ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24999</td>\n",
       "      <td>Row 24999: I went, yesterday; she wasn't here ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                       a\n",
       "24995  Row 24995: I went, yesterday; she wasn't here ...\n",
       "24996  Row 24996: I went, yesterday; she wasn't here ...\n",
       "24997  Row 24997: I went, yesterday; she wasn't here ...\n",
       "24998  Row 24998: I went, yesterday; she wasn't here ...\n",
       "24999  Row 24999: I went, yesterday; she wasn't here ..."
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(text, columns=['a'])\n",
    "df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "tokenizer = partial(spacy.load, name='en_core_web_sm',\n",
    "                    disable=('ner', 'parser', 'tagger'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "NLP = tokenizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def tokenize(text, nlp):\n",
    "    \"\"\"Word tokenize a single string. \n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: str\n",
    "        A piece of text to tokenize.\n",
    "    nlp: spacy tokenizer, e.g. spacy.lang.en.English\n",
    "        By default, a spacy tokenizer with a small English vocabulary \n",
    "        is used. NER, parsing, and tagging are disabled. Any spacy\n",
    "        tokenzer can be passed in, but keep in mind other configurations \n",
    "        may slow down this function dramatically.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[str]: List of word tokens from a single input string.\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in nlp(text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def tokenize_many(rows, nlp=None, chunk=1_000):\n",
    "    \"\"\"Word tokenize a sequence of strings using multiprocessing. The max\n",
    "    number of available processes are used.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    rows: Iterable[str]\n",
    "        A sequence of strings to tokenize. This could be a list, a column of\n",
    "        a DataFrame, etc.\n",
    "    nlp: spacy tokenizer, e.g. spacy.lang.en.English\n",
    "        By default, a spacy tokenizer with a small English vocabulary \n",
    "        is used. NER, parsing, and tagging are disabled. Any spacy\n",
    "        tokenzer can be passed in, but keep in mind other configurations \n",
    "        may slow down this function dramatically.\n",
    "    chunk: int\n",
    "        This determines how many items to send to multiprocessing at a time.\n",
    "        The default of 1,000 is usually fine, but if you have extremely\n",
    "        long pieces of text and memory is limited, you can always decrease it.\n",
    "        Very small chunk sizes may increase processing time. Note that larger\n",
    "        values will generally cause the progress bar to update more choppily.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[list[str]]: Each nested list of word tokens corresponds to one\n",
    "    of the input strings.\n",
    "    \"\"\"\n",
    "    tokenize_ = partial(tokenize, nlp=nlp or tokenizer())\n",
    "    length = len(rows)\n",
    "    with multiprocessing.Pool() as p:\n",
    "        res = list(tqdm(p.imap(tokenize_, rows, chunksize=chunk),\n",
    "                        total=length))\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ~5-6 seconds\n",
    "x = df.a.apply(tokenize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f4fb0b51fe44e148fb3b5c8a42afa9a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=25000), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# ~1-2 seconds\n",
    "x = tokenize_many(df.a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Vocabulary:\n",
    "\n",
    "    def __init__(self, w2idx, w2vec=None, idx_misc=None, corpus_counts=None,\n",
    "                 all_lower=True):\n",
    "        \"\"\"Defines a vocabulary object for NLP problems, allowing users to\n",
    "        encode text with indices or embeddings.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        w2idx: dict[str, int]\n",
    "            Dictionary mapping words to their integer index in a vocabulary.\n",
    "            The indices must allow for idx_misc to be added to the dictionary,\n",
    "            so in the default case this should have a minimum index of 2. If\n",
    "            a longer idx_misc is passed in, the minimum index would be larger.\n",
    "        w2vec: dict[str, np.array]\n",
    "            Dictionary mapping words to their embedding vectors stored as\n",
    "            numpy arrays (optional).\n",
    "        idx_misc: dict\n",
    "            A dictionary mapping non-word tokens to indices. If none is passed\n",
    "            in, a default version will be used with keys for unknown tokens\n",
    "            and padding. A customized version might pass in additional tokens\n",
    "            for repeated characters or all caps, for example.\n",
    "        corpus_counts: collections.Counter\n",
    "            Counter dict mapping words to their number of occurrences in a\n",
    "            corpus (optional).\n",
    "        all_lower: bool\n",
    "            Specifies whether the data you've passed in (w2idx, w2vec, i2w) is\n",
    "            all lowercase. Note that this will NOT change any of this data. If\n",
    "            True, it simply lowercases user-input words when looking up their\n",
    "            index or vector.\n",
    "        \"\"\"\n",
    "        if not idx_misc:\n",
    "            idx_misc = {'<PAD>': 0,\n",
    "                        '<UNK>': 1}\n",
    "        self.idx_misc = idx_misc\n",
    "        # Check that space has been left for misc keys.\n",
    "        assert len(idx_misc) == min(w2idx.values())\n",
    "\n",
    "        # Core data structures.\n",
    "        self.w2idx = {**self.idx_misc, **w2idx}\n",
    "        self.i2w = [word for word, idx in sorted(self.w2idx.items(),\n",
    "                                                 key=lambda x: x[1])]\n",
    "        self.w2vec = w2vec or dict()\n",
    "\n",
    "        # Miscellaneous other attributes.\n",
    "        if w2vec:\n",
    "            self.dim = len(w2vec[self[-1]])\n",
    "        else:\n",
    "            self.dim = 1\n",
    "        self.corpus_counts = corpus_counts\n",
    "        self.embedding_matrix = None\n",
    "        self.w2vec['<UNK>'] = np.zeros(self.dim)\n",
    "        self.all_lower = all_lower\n",
    "\n",
    "    @classmethod\n",
    "    def from_glove_file(cls, path, max_lines=float('inf'), idx_misc=None):\n",
    "        \"\"\"Create a new Vocabulary object by loading GloVe vectors from a text\n",
    "        file. The embeddings are all lowercase so the user does not have the\n",
    "        option to set the all_lower parameter.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        path: str\n",
    "            Path to file containing glove vectors.\n",
    "        max_lines: int, float (optional)\n",
    "            Loading the GloVe vectors can be slow, so for testing purposes\n",
    "            it can be helpful to read in a subset. If no value is provided,\n",
    "            all 400,000 lines in the file will be read in.\n",
    "        idx_misc: dict\n",
    "            Map non-standard tokens to indices. See constructor docstring.\n",
    "        \"\"\"\n",
    "        w2idx = dict()\n",
    "        w2vec = dict()\n",
    "        misc_len = 2 if not idx_misc else len(idx_misc)\n",
    "\n",
    "        with open(path, 'r') as f:\n",
    "            for i, line in enumerate(f):\n",
    "                if i >= max_lines:\n",
    "                    break\n",
    "                word, *values = line.strip().split(' ')\n",
    "                w2idx[word] = i + misc_len\n",
    "                w2vec[word] = np.array(values, dtype=np.float)\n",
    "\n",
    "        return cls(w2idx, w2vec, idx_misc)\n",
    "\n",
    "    @classmethod\n",
    "    def from_tokens(cls, tokens, idx_misc=None, all_lower=True):\n",
    "        \"\"\"Construct a Vocabulary object from a list or array of tokens.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        tokens: list[str]\n",
    "            The word-tokenized corpus.\n",
    "        idx_misc: dict\n",
    "            Map non-standard tokens to indices. See constructor docstring.\n",
    "        all_lower: bool\n",
    "            Specifies whether your tokens are all lowercase.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        Vocabulary\n",
    "        \"\"\"\n",
    "        misc_len = 2 if not idx_misc else len(idx_misc)\n",
    "        counts = Counter(tokens)\n",
    "        w2idx = {word: i for i, (word, freq)\n",
    "                 in enumerate(counts.most_common(), misc_len)}\n",
    "        return cls(w2idx, idx_misc=idx_misc, corpus_counts=counts,\n",
    "                   all_lower=all_lower)\n",
    "\n",
    "    @staticmethod\n",
    "    def from_pickle(path):\n",
    "        \"\"\"Load a previously saved Vocabulary object.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        path: str\n",
    "            Location of pickled Vocabulary file.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        Vocabulary\n",
    "        \"\"\"\n",
    "        return load(path)\n",
    "\n",
    "    def save(self, path, verbose=True):\n",
    "        \"\"\"Pickle Vocabulary object for later use. We can then quickly load\n",
    "        the object using torch.load(path), which can be much faster than\n",
    "        re-computing everything when the vocab size becomes large.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        path: str\n",
    "            Where to save the output file.\n",
    "        verbose: bool\n",
    "            If True, print message showing where the object was saved to.\n",
    "        \"\"\"\n",
    "        save(self, path, verbose)\n",
    "\n",
    "    def filter_tokens(self, tokens, max_words=None, min_freq=0, inplace=False,\n",
    "                      recompute=False):\n",
    "        \"\"\"Filter your vocabulary by specifying a max number of words or a min\n",
    "        frequency in the corpus. When done in place, this also sorts vocab by\n",
    "        frequency with more common words coming first (after idx_misc).\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        tokens: list[str]\n",
    "            A tokenized list of words in the corpus (must be all lowercase\n",
    "            when self.all_lower=True, such as when using GloVe vectors). There\n",
    "            is no need to hold out test data here since we are not using\n",
    "            labels.\n",
    "        max_words: int (optional)\n",
    "            Provides an upper threshold for the number of words in the\n",
    "            vocabulary. If no value is passed in, no maximum limit will be\n",
    "            enforced.\n",
    "        min_freq: int (optional)\n",
    "            Provides a lower threshold for the number of times a word must\n",
    "            appear in the corpus to remain in the vocabulary. If no value is\n",
    "            passed in, no minimum limit will be enforced.\n",
    "\n",
    "            Note that we can specify values for both max_words and min_freq\n",
    "            if desired. If no values are passed in for either, no pruning of\n",
    "            the vocabulary will be performed.\n",
    "        inplace: bool\n",
    "            If True, will change the object's attributes\n",
    "            (w2idx, w2vec, and i2w) to reflect the newly filtered vocabulary.\n",
    "            If False, will not change the object, but will simply compute word\n",
    "            counts and return what the new w2idx would be. This can be helpful\n",
    "            for experimentation, as we may want to try out multiple values of\n",
    "            min_freq to decide how many words to keep. After the first call,\n",
    "            the attribute corpus_counts can also be examined to help determine\n",
    "            the desired vocab size.\n",
    "        recompute: bool\n",
    "            If True, will calculate word counts from the given tokens. If\n",
    "            False (the default), this will use existing counts if there are\n",
    "            any.\n",
    "\n",
    "            The idea is that if we call this method, then realize we want\n",
    "            to change the corpus, we should calculate new word counts.\n",
    "            However, if we are simply calling this method multiple times on\n",
    "            the same corpus while deciding on the exact vocab size we want,\n",
    "            we should not recompute the word counts.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        dict or None: When called inplace, nothing is returned. When not\n",
    "        inplace,\n",
    "        \"\"\"\n",
    "        misc_len = len(self.idx_misc)\n",
    "        if recompute or not self.corpus_counts:\n",
    "            self.corpus_counts = Counter(tokens)\n",
    "        filtered = {word: i for i, (word, freq)\n",
    "                    in enumerate(self.corpus_counts.most_common(max_words),\n",
    "                                 misc_len)\n",
    "                    if freq >= min_freq}\n",
    "        filtered = {**self.idx_misc, **filtered}\n",
    "\n",
    "        if inplace:\n",
    "            # Relies on python3.7 dicts retaining insertion order.\n",
    "            self.i2w = list(filtered.keys())\n",
    "            self.w2idx = filtered\n",
    "            self.w2vec = {word: self.vector(word) for word in filtered}\n",
    "        else:\n",
    "            return filtered\n",
    "\n",
    "    def build_embedding_matrix(self, inplace=False):\n",
    "        \"\"\"Create a 2D numpy array of embedding vectors where row[i]\n",
    "        corresponds to word i in the vocabulary. This can be used to\n",
    "        initialize weights in the model's embedding layer.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        inplace: bool\n",
    "            If True, will store the output in the object's embedding_matrix\n",
    "            attribute. If False (default behavior), will simply return the\n",
    "            matrix without storing it as part of the object. In the\n",
    "            recommended case where inplace==False, we can store the output\n",
    "            in another variable which we can use to initialize the weights in\n",
    "            Torch, then delete the object and free up memory using\n",
    "            gc.collect().\n",
    "        \"\"\"\n",
    "        emb = np.zeros((len(self), self.dim))\n",
    "        for i, word in enumerate(self):\n",
    "            emb[i] = self.vector(word)\n",
    "\n",
    "        if inplace:\n",
    "            self.embedding_matrix = emb\n",
    "        else:\n",
    "            return emb\n",
    "\n",
    "    def idx(self, word):\n",
    "        \"\"\"This will map a word (str) to its index (int) in the vocabulary.\n",
    "        If a string is passed in and the word is not present, the index\n",
    "        corresponding to the <UNK> token is returned.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        word: str\n",
    "            A word that needs to be mapped to an integer index.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        int: The index of the given word in the vocabulary.\n",
    "\n",
    "        Examples\n",
    "        ---------\n",
    "        >>> vocab.idx('the')\n",
    "        2\n",
    "        \"\"\"\n",
    "        if self.all_lower and word not in self.idx_misc:\n",
    "            word = word.lower()\n",
    "        return self.w2idx.get(word, self.w2idx['<UNK>'])\n",
    "\n",
    "    def vector(self, word):\n",
    "        \"\"\"This maps a word to its corresponding embedding vector. If not\n",
    "        contained in the vocab, a vector of zeros will be returned.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        word: str\n",
    "            A word that needs to be mapped to a vector.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        np.array\n",
    "        \"\"\"\n",
    "        if self.all_lower and word not in self.idx_misc:\n",
    "            word = word.lower()\n",
    "        return self.w2vec.get(word, self.w2vec['<UNK>'])\n",
    "\n",
    "    def encode(self, text, nlp, max_len, pad_end=True, trim_start=True):\n",
    "        \"\"\"Encode text so that each token is replaced by its integer index in\n",
    "        the vocab.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        text: str\n",
    "            Raw text to be encoded.\n",
    "        nlp: spacy.lang.en.English\n",
    "            Spacy tokenizer. Typically want to disable 'parser', 'tagger', and\n",
    "            'ner' as they aren't used here and slow down the encoding process.\n",
    "        max_len: int\n",
    "            Length of output encoding. If text is shorter, it will be padded\n",
    "            to fit the specified length. If text is longer, it will be\n",
    "            trimmed.\n",
    "        pad_end: bool\n",
    "            If True, add padding to the end of short sentences. If False, pad\n",
    "            the start of these sentences.\n",
    "        trim_start: bool\n",
    "            If True, trim off the start of sentences that are too long. If\n",
    "            False, trim off the end.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        np.array[int]: Array of length max_len containing integer indices\n",
    "            corresponding to the words passed in.\n",
    "        \"\"\"\n",
    "        output = np.ones(max_len) * self.idx('<PAD>')\n",
    "        encoded = [self.idx(tok.text) for tok in nlp(text)]\n",
    "\n",
    "        # Trim sentence in case it's longer than max_len.\n",
    "        if len(encoded) > max_len:\n",
    "            if trim_start:\n",
    "                encoded = encoded[len(encoded) - max_len:]\n",
    "            else:\n",
    "                encoded = encoded[:max_len]\n",
    "\n",
    "        # Replace zeros at start or end, depending on choice of pad_end.\n",
    "        if pad_end:\n",
    "            output[:len(encoded)] = encoded\n",
    "        else:\n",
    "            output[max_len-len(encoded):] = encoded\n",
    "        return output.astype(int)\n",
    "\n",
    "    def decode(self, idx):\n",
    "        \"\"\"Convert a list of indices to a string of words/tokens.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        idx: list[int]\n",
    "            A list of integers indexing into the vocabulary. This will often\n",
    "            be the output of the encode() method.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        list[str]: A list of words/tokens reconstructed by indexing into the\n",
    "            vocabulary.\n",
    "        \"\"\"\n",
    "        return [self[i] for i in idx]\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        \"\"\"This will map an index (int) to a word (str).\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        i: int\n",
    "            Integer index for a word.\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        str: Word corresponding to the given index.\n",
    "\n",
    "        Examples\n",
    "        ---------\n",
    "        >>> vocab = Vocabulary(w2idx, w2vec)\n",
    "        >>> vocab[1]\n",
    "        '<UNK>'\n",
    "        \"\"\"\n",
    "        return self.i2w[i]\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Number of words in vocabulary.\"\"\"\n",
    "        return len(self.w2idx)\n",
    "\n",
    "    def __iter__(self):\n",
    "        for word in self.w2idx.keys():\n",
    "            yield word\n",
    "\n",
    "    def __contains__(self, word):\n",
    "        return word in self.w2idx.keys()\n",
    "\n",
    "    def __eq__(self, obj):\n",
    "        if not isinstance(obj, Vocabulary):\n",
    "            return False\n",
    "\n",
    "        ignore = {'w2vec', 'embedding_matrix'}\n",
    "        attrs = [k for k, v in hdir(vocab).items()\n",
    "                 if v == 'attribute' and k not in ignore]\n",
    "        return all([getattr(self, attr) == getattr(obj, attr)\n",
    "                    for attr in attrs])\n",
    "\n",
    "    def __repr__(self):\n",
    "        msg = f'Vocabulary({len(self)} words'\n",
    "        if self.dim > 1:\n",
    "            msg += f', {self.dim}-D embeddings'\n",
    "        return msg + ')'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Embeddings:\n",
    "    \"\"\"Embeddings object. Lets us easily map word to index, index to\n",
    "    word, and word to vector. We can use this to find similar words,\n",
    "    build analogies, or get 2D representations for cdting.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, mat, w2i, mat_2d=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        mat: str\n",
    "            Numpy array of embeddings where row i corresponds to ID i\n",
    "            in w2i.\n",
    "        w2i: dict[str, int]\n",
    "            Dictionary mapping word to its index in the vocabulary.\n",
    "        mat_2d: np.array\n",
    "            Matrix output of PCA after compressing mat to vectors of length 2.\n",
    "            If None, it will be computed from mat and cached.\n",
    "        \"\"\"\n",
    "        self.mat = mat\n",
    "        self.w2i = w2i\n",
    "        self.i2w = [w for w, i in sorted(self.w2i.items(), key=lambda x: x[1])]\n",
    "        self.mat_2d = mat_2d or PCA(n_components=2).fit_transform(self.mat)\n",
    "        self.n_embeddings, self.dim = self.mat.shape\n",
    "        \n",
    "    @classmethod\n",
    "    def from_text_file(cls, path, max_words=float('inf'), print_freq=10_000):\n",
    "        \"\"\"Create a new Embeddings object from a raw text file using the\n",
    "        GloVe format (each row contains a word and its embedding as \n",
    "        space-separated floats).\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str\n",
    "            Location of csv file containing GloVe vectors.\n",
    "        max_words: int, float\n",
    "            Set maximum number of words to read in from file. This can be used\n",
    "            during development to reduce wait times when loading data.\n",
    "        Returns\n",
    "        -------\n",
    "        Embeddings: Newly instantiated object.\n",
    "        \"\"\"\n",
    "        w2i = dict()\n",
    "        mat = []\n",
    "        with open(path, 'r') as f:\n",
    "            for i, line in enumerate(f):\n",
    "                # Faster testing\n",
    "                if i >= max_words: break\n",
    "                word, *nums = line.strip().split()\n",
    "                w2i[word] = i\n",
    "                mat.append(np.array(nums, dtype=float))\n",
    "                if i % print_freq == 0: print(i, word)\n",
    "        return cls(np.array(mat), w2i)\n",
    "\n",
    "    @classmethod\n",
    "    def from_pickle(cls, path):\n",
    "        \"\"\"If an Embeddings object previously saved its data in a pickle file,\n",
    "        loading it that way can avoid repeated computation.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str\n",
    "            Location of pickle file.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Embeddings: Newly instantiated object using the data that was stored in\n",
    "            the pickle file.\n",
    "        \"\"\"\n",
    "        return cls(**load(path))\n",
    "\n",
    "    def save(self, path, verbose=True):\n",
    "        \"\"\"Save data to a compressed pickle file. This reduces the amount of\n",
    "        space needed for storage (the csv is much larger) and can let us\n",
    "        avoid running PCA and building the embedding matrix again.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str\n",
    "            Path that object will be saved to.\n",
    "        verbose\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        data = dict(mat=self.mat,\n",
    "                    w2i=self.w2i,\n",
    "                    mat_2d=self.mat_2d)\n",
    "        save(data, path, verbose=verbose)\n",
    "\n",
    "    def vec(self, word):\n",
    "        \"\"\"Look up the embedding for a given word. Return None if not found.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        word: str\n",
    "            Input word to look up embedding for.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Embedding corresponding to the input word. If word not in\n",
    "            vocab, return None.\n",
    "        \"\"\"\n",
    "        idx = self[word]\n",
    "        if idx is not None:\n",
    "            return self.mat[idx]\n",
    "\n",
    "    def vec_2d(self, word):\n",
    "        \"\"\"Look up the compressed embedding for a word (PCA was used to shrink\n",
    "        dimensionality to 2). Return None if the word is not present in vocab.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        word: str\n",
    "            Input work to look up.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Compressed embedding of length 2. None if not found.\n",
    "        \"\"\"\n",
    "        idx = self[word]\n",
    "        if idx is not None:\n",
    "            return self.mat_2d[idx]\n",
    "\n",
    "    def _distances(self, vec, distance='cosine'):\n",
    "        \"\"\"Find distance from an input vector to every other vector in the\n",
    "        embedding matrix.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec: np.array\n",
    "            Vector for the input word.\n",
    "        distance: str\n",
    "            Specifies what distance metric to use for calculations.\n",
    "            One of ('euclidean', 'manhattan', 'cosine'). In a high dimensional\n",
    "            space, cosine is often a good choice.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: The i'th value corresponds to the distance to word i in the\n",
    "            vocabulary.\n",
    "        \"\"\"\n",
    "        if distance == 'euclidean':\n",
    "            dists = self.norm(self.mat - vec)\n",
    "        elif distance == 'cosine':\n",
    "            dists = self.cosine_distance(vec, self.mat)\n",
    "        elif distance == 'manhattan':\n",
    "            dists = self.manhattan_distance(vec, self.mat)\n",
    "        return dists\n",
    "\n",
    "    def nearest_neighbors(self, word, n=5, distance='cosine', digits=3):\n",
    "        \"\"\"Find the most similar words to a given word. This wrapper to\n",
    "        allows the user to pass in a word. To pass in a vector, use\n",
    "        _nearest_neighbors().\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        word: str\n",
    "            A word that must be in the vocabulary.\n",
    "        n: int\n",
    "            Number of neighbors to return.\n",
    "        distance: str\n",
    "            Distance method to use when computing nearest neighbors. One of\n",
    "            ('euclidean', 'manhattan', 'cosine').\n",
    "        digits: int\n",
    "            Digits to round output distances to.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict[str, float]: Dictionary mapping word to distance.\n",
    "        \"\"\"\n",
    "        # Error handling for words not in vocab.\n",
    "        if word not in self:\n",
    "            return None\n",
    "        return self._nearest_neighbors(self.vec(word), n, distance, digits)\n",
    "\n",
    "    def _nearest_neighbors(self, vec, n=5, distance='cosine', digits=3,\n",
    "                           skip_first=True):\n",
    "        \"\"\"Internal function behind nearest_neighbors(). This can be used if\n",
    "        we want to pass in a vector instead of a word. For more details, see\n",
    "        the wrapper method.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec: np.array\n",
    "        n: int\n",
    "        distance: str\n",
    "        digits: int\n",
    "        skip_first: bool\n",
    "            If True, the nearest result will be sliced off (this is desirable\n",
    "            when searching for a word's nearest neighbors, where we don't want\n",
    "            to return the word itself). When finding analogies or performing\n",
    "            embedding arithmetic, however, we likely don't want to slice off\n",
    "            the first result.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict[str, float]: Dictionary mapping word to distance.\n",
    "        \"\"\"\n",
    "        dists = self._distances(vec, distance)\n",
    "        idx = np.argsort(dists)[slice(skip_first, skip_first+n)]\n",
    "        return {self.i2w[i]: round(dists[i], digits) for i in idx}\n",
    "\n",
    "    def analogy(self, a, b, c, n=5, **kwargs):\n",
    "        \"\"\"Fill in the analogy: A is to B as C is to ___. Note that we always\n",
    "        treat A and B as valid candidates to fill in the blank. C is\n",
    "        only considered as a candidate in the trivial case where A=B, in which\n",
    "        case C should be the first choice.\n",
    "        Parameters\n",
    "        ----------\n",
    "        a: str\n",
    "            First word in analogy.\n",
    "        b: str\n",
    "            Second word in analogy.\n",
    "        c: str\n",
    "            Third word in analogy.\n",
    "        n: int\n",
    "            Number of candidates to return. Note that we specify this\n",
    "            separately fro kwargs since we need to alter its value before\n",
    "            passing it to _nearest_neighbors(). This will allow us to remove\n",
    "            the word c as a candidate if it is returned.\n",
    "        kwargs: distance (str), digits (int)\n",
    "            See _nearest_neighbors for details.\n",
    "        Returns\n",
    "        -------\n",
    "        list[str]: Best candidates to complete the analogy in descending order\n",
    "            of likelihood.\n",
    "        \"\"\"\n",
    "        # If any words missing from vocab, arithmetic w/ None will throw error.\n",
    "        try:\n",
    "            vec = self.vec(b) - self.vec(a) + self.vec(c)\n",
    "        except TypeError:\n",
    "            return None\n",
    "\n",
    "        # Except for trivial edge case, return 1 extra value in case neighbors\n",
    "        # includes c, which will be removed in these situations.\n",
    "        a, b, c = a.lower(), b.lower(), c.lower()\n",
    "        trivial = (a == b)\n",
    "        neighbors = self._nearest_neighbors(vec,\n",
    "                                            n=n+1-trivial,\n",
    "                                            skip_first=False,\n",
    "                                            **kwargs)\n",
    "        if not trivial and c in neighbors:\n",
    "            neighbors.pop(c)\n",
    "\n",
    "        # Relies on dicts being ordered in python >= 3.6.\n",
    "        return list(neighbors)[:n]\n",
    "    \n",
    "    def cbow(self, *args):\n",
    "        \"\"\"Wrapper to _cbow() that allows us to pass in strings instead of\n",
    "        vectors.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        args: str\n",
    "            Multiple words to average over.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Average of all input vectors. This will have the same\n",
    "            embedding dimension as each input.\n",
    "        \"\"\"\n",
    "        vecs = [arg for arg in map(self.vec, args) if arg is not None]\n",
    "        if vecs:\n",
    "            return self._cbow(*vecs)\n",
    "\n",
    "    def _cbow(self, *args):\n",
    "        \"\"\"Internal helper for cbow(). Can also use this directly if you want\n",
    "        to pass in vectors instead of words.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        args: np.array\n",
    "            Word vectors to average.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Average of all input vectors. This will have the same\n",
    "            embedding dimension as each input.\n",
    "        \"\"\"\n",
    "        return np.mean(args, axis=0)\n",
    "\n",
    "    def cbow_neighbors(self, *args, n=5, **kwargs):\n",
    "        \"\"\"Wrapper to cbow(). This lets us pass in words, compute their\n",
    "        average embedding, then return the words nearest this embedding. The\n",
    "        input words are not considered to be candidates for neighbors (e.g. if\n",
    "        you input the words 'happy' and 'cheerful', the neighbors returned will\n",
    "        not include those words even if they are the closest to the mean\n",
    "        embedding). The idea here is to find additional words that may be\n",
    "        similar to the group you've passed in.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        args: str\n",
    "            Input words to average over.\n",
    "        n: int\n",
    "            Number of neighbors to return.\n",
    "        kwargs: distance (str), digits (int)\n",
    "            See _nearest_neighbors() for details.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict[str, float]: Dictionary mapping word to distance from the average\n",
    "            of the input words' vectors.\n",
    "        \"\"\"\n",
    "        vec_avg = self.cbow(*args)\n",
    "        if vec_avg is None:\n",
    "            return\n",
    "        neighbors = self._nearest_neighbors(vec_avg, n=len(args)+n,\n",
    "                                            skip_first=False, **kwargs)\n",
    "\n",
    "        # Lowercase to help remove duplicates.\n",
    "        args = set(arg.lower() for arg in args)\n",
    "        return {word: neighbors[word] for word in\n",
    "                [n for n in neighbors if n not in args][:n]}\n",
    "\n",
    "    @staticmethod\n",
    "    def norm(vec):\n",
    "        \"\"\"Compute L2 norm of a vector. Euclidean distance between two vectors\n",
    "        can be found by the operation norm(vec1 - vec2).\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec: np.array\n",
    "            Input vector.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        float: L2 norm of input vector.\n",
    "        \"\"\"\n",
    "        return np.sqrt(np.sum(vec ** 2, axis=-1))\n",
    "\n",
    "    @staticmethod\n",
    "    def manhattan_distance(vec1, vec2):\n",
    "        \"\"\"Compute L1 distance between two vectors.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec1: np.array\n",
    "        vec2: np.array\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        float or np.array: Manhattan distance between vec1 and vec2. If two\n",
    "            vectors are passed in, the output will be a single number. When\n",
    "            computing distances between a vector and a matrix, the output\n",
    "            will be a vector (np.array).\n",
    "        \"\"\"\n",
    "        return np.sum(abs(vec1 - vec2), axis=-1)\n",
    "\n",
    "    def cosine_distance(self, vec1, vec2):\n",
    "        \"\"\"Compute cosine distance between two vectors.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec1: np.array\n",
    "        vec2: np.array\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        float or np.array: Cosine distance between vec1 and vec2. If two\n",
    "            vectors are passed in, the output will be a single number. When\n",
    "            computing distances between a vector and a matrix, the output\n",
    "            will be a vector (np.array).\n",
    "        \"\"\"\n",
    "        return 1 - (np.sum(vec1 * vec2, axis=-1) /\n",
    "                    (self.norm(vec1) * self.norm(vec2)))\n",
    "\n",
    "    def __getitem__(self, word):\n",
    "        return self.w2i.get(word.lower())\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n_embeddings\n",
    "\n",
    "    def __contains__(self, word):\n",
    "        return word.lower() in self.w2i\n",
    "\n",
    "    def __iter__(self):\n",
    "        \"\"\"Yields words in vocabulary.\"\"\"\n",
    "        yield from self.w2i.keys()\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'Embeddings(len={len(self)}, dim={self.dim})'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def back_translate(text, to, from_lang='en'):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    str: Same language and basically the same content as the original text,\n",
    "        but usually with slightly altered grammar, sentence structure, and/or\n",
    "        vocabulary.\n",
    "    \"\"\"\n",
    "    return str(\n",
    "        TextBlob(text)\\\n",
    "        .translate(to=to)\\\n",
    "        .translate(from_lang=to, to=from_lang)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Visit ESPN to get coverage of sports news, scores, highlights and comments from the NFL, MLB, NBA, college football, NCAA basketball and more.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"\"\"\n",
    "Visit ESPN to get up-to-the-minute sports news coverage, scores, highlights and commentary for NFL, MLB, NBA, College Football, NCAA Basketball and more.\n",
    "\"\"\"\n",
    "back_translate(text, 'es')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Visit ESPN for up-to-date sports information, scores, highlights and commentary for the NFL, MLB, NBA, college football, NCAA basketball and more.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"\"\"\n",
    "Visit ESPN to get up-to-the-minute sports news coverage, scores, highlights and commentary for NFL, MLB, NBA, College Football, NCAA Basketball and more.\n",
    "\"\"\"\n",
    "back_translate(text, 'fr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def postprocess_embeddings(emb, d=None):\n",
    "    \"\"\"Implements the algorithm from the paper:\n",
    "    \n",
    "    All-But-The-Top: Simple and Effective Post-Processing \n",
    "    for Word Representations (https://arxiv.org/pdf/1702.01417.pdf)\n",
    "    \n",
    "    There are three steps:\n",
    "    1. Compute the mean embedding and subtract this from the \n",
    "    original embedding matrix. \n",
    "    2. Perform PCA and extract the top d components.\n",
    "    3. Eliminate the principal components from the mean-adjusted\n",
    "    embeddings.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    emb: np.array\n",
    "        Embedding matrix of size (vocab_size, embedding_length).\n",
    "    d: int\n",
    "        Number of components to use in PCA. Defaults to \n",
    "        embedding_length/100 as recommended by the paper.\n",
    "    \"\"\"\n",
    "    d = d or emb.shape[1] // 100\n",
    "    emb_adj = emb - emb.mean(0)\n",
    "    u = PCA(d).fit(emb_adj).components_\n",
    "    return emb_adj - emb@u.T@u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def compress_embeddings(emb, new_dim, d=None):\n",
    "    \"\"\"Reduce embedding dimension as described in the paper:\n",
    "    \n",
    "    Simple and Effective Dimensionality Reduction for Word Embeddings\n",
    "    (https://lld-workshop.github.io/2017/papers/LLD_2017_paper_34.pdf)\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    emb: np.array\n",
    "        Embedding matrix of size (vocab_size, embedding_length).\n",
    "    d: int\n",
    "        Number of components to use in the post-processing\n",
    "        method described here: https://arxiv.org/pdf/1702.01417.pdf\n",
    "        Defaults to embedding_length/100 as recommended by the paper.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    np.array: Compressed embedding matrix of shape (vocab_size, new_dim).\n",
    "    \"\"\"\n",
    "    emb = postprocess_embeddings(emb, d)\n",
    "    emb = PCA(new_dim).fit_transform(emb)\n",
    "    return postprocess_embeddings(emb, d)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
