{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:14:46.891248Z",
     "start_time": "2021-05-09T21:14:46.888545Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp nlp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NLP\n",
    "\n",
    "> API details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:14:48.327158Z",
     "start_time": "2021-05-09T21:14:47.923218Z"
    }
   },
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:14:49.007764Z",
     "start_time": "2021-05-09T21:14:48.329406Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:14:53.713282Z",
     "start_time": "2021-05-09T21:14:49.011560Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "from bs4 import BeautifulSoup\n",
    "from collections import Counter\n",
    "from collections.abc import Iterable\n",
    "from functools import partial\n",
    "from multipledispatch import dispatch\n",
    "import multiprocessing\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "import re\n",
    "import requests\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.utils.validation import check_is_fitted\n",
    "import spacy\n",
    "from textblob import TextBlob\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import PegasusForConditionalGeneration, PegasusTokenizer, \\\n",
    "    PegasusTokenizerFast, Text2TextGenerationPipeline, \\\n",
    "    AutoModelForSeq2SeqLM, AutoTokenizer, TranslationPipeline, pipeline\n",
    "from tldextract import extract\n",
    "from transformers.modeling_utils import PreTrainedModel\n",
    "import warnings\n",
    "import wordninja as wn\n",
    "\n",
    "from htools import save, load, add_docstring, tolist, auto_repr, listlike, \\\n",
    "    flatten, immutify_defaults, ifnone, item, lmap, func_name, valuecheck\n",
    "from incendio.utils import DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:14:53.793092Z",
     "start_time": "2021-05-09T21:14:53.715519Z"
    }
   },
   "outputs": [],
   "source": [
    "# Only needed for testing.\n",
    "from string import ascii_lowercase\n",
    "\n",
    "from htools import assert_raises"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "tokenizer = partial(spacy.load, name='en_core_web_sm',\n",
    "                    disable=('ner', 'parser', 'tagger'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-16T23:56:29.859607Z",
     "start_time": "2021-01-16T23:56:29.834777Z"
    }
   },
   "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, join=True, sep=' '):\n",
    "        \"\"\"Convert a list of indices to a string or list 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",
    "        join: bool\n",
    "            If True, return a single string. If False, return a list of \n",
    "            strings.\n",
    "        sep: str\n",
    "            If join is True, this determines what character is used to join\n",
    "            tokens. Word tokens will usually be joined by a space, but some\n",
    "            tokenization schemes include spaces and can be joined with an \n",
    "            empty string ('').\n",
    "\n",
    "        Returns\n",
    "        --------\n",
    "        str or list[str]: A string or list of strings (words/tokens)\n",
    "        reconstructed by indexing into the vocabulary.\n",
    "        \"\"\"\n",
    "        tokens = [self[i] for i in idx]\n",
    "        if join: return sep.join(tokens)\n",
    "        return tokens\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": {
    "ExecuteTime": {
     "end_time": "2021-02-13T04:07:10.828799Z",
     "start_time": "2021-02-13T04:07:10.816191Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def domain(url, strip_www=True, mode='registered_domain'):\n",
    "    domain_ = getattr(extract(url), mode)\n",
    "    return domain_.replace('www.', '') if strip_www else domain_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-13T04:07:10.828799Z",
     "start_time": "2021-02-13T04:07:10.816191Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def domains_from_google_search(term, drop_probable_defaults=True):\n",
    "    \"\"\"Get domains from first page of google sarch for a given term.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    term: str\n",
    "        Term to search for.\n",
    "    drop_probable_defaults: bool\n",
    "        If True, remove two domains which seem to be present on all google\n",
    "        search results pages. I didn't see a simple identifier for which links\n",
    "        on the page are search results so we're left with this clumsy \n",
    "        filtering method (setting this to True means if 'accounts.google.com' \n",
    "        is a legitimate search result, we'd still remove it).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    set[str]: Domain names (no www) found on the first page of search results.\n",
    "    \"\"\"\n",
    "    url = f'http://www.google.com/search?q={\"+\".join(term.split())}'\n",
    "    r = requests.get(url)\n",
    "    links = BeautifulSoup(r.content, 'lxml').find_all('a')\n",
    "    # Use dict to remove duplicates while maintaining order.\n",
    "    domains = dict.fromkeys(domain(link['href'].replace('/url?q=', '')) \n",
    "                            for link in links \n",
    "                            if link['href'].startswith('/url?q='))\n",
    "    if drop_probable_defaults:\n",
    "        for d in ('accounts.google.com', 'support.google.com'):\n",
    "            domains.pop(d, None)\n",
    "    return list(domains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-13T04:07:16.553788Z",
     "start_time": "2021-02-13T04:07:13.434132Z"
    },
    "run_control": {
     "marked": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['healthline.com',\n",
       " 'webmd.com',\n",
       " 'health.harvard.edu',\n",
       " 'amazon.com',\n",
       " 'prevention.com',\n",
       " 'health.clevelandclinic.org',\n",
       " 'eatthis.com',\n",
       " 'medicalnewstoday.com']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "domains_from_google_search('protein powder')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-09T21:23:13.407041Z",
     "start_time": "2021-05-09T21:23:12.117731Z"
    }
   },
   "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 plotting. Generally, \n",
    "    user-facing methods let us pass in strings, while internal versions (same \n",
    "    name except prefixed with an underscore) allow us to pass in vectors.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, mat, w2i, pca=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",
    "        pca: sklearn.decomposition.PCA or None\n",
    "            If provided, this should be a PCA object with 2 components that\n",
    "            was previously fit on `mat`. If None, a new object will be created\n",
    "            and fit. This will let us plot embeddings in a way humans can\n",
    "            visually parse.\n",
    "        \"\"\"\n",
    "        self.mat = mat\n",
    "        max_id = max(w2i.values())\n",
    "        expected_ids = list(range(max_id + 1))\n",
    "        if list(w2i.values()) != expected_ids:\n",
    "            if sorted(w2i.values()) == expected_ids:\n",
    "                warnings.warn(\n",
    "                    'Your w2i dict is out of order (where ordered would mean '\n",
    "                    'the first key has id 0, the second has id 1, etc.). '\n",
    "                    'This should technically be fine but we recommend fixing '\n",
    "                    'it to be safe.'\n",
    "                )\n",
    "            else:\n",
    "                raise ValueError('Your w2i dict has missing indices. '\n",
    "                                 'We do not currently support gaps.')\n",
    "            \n",
    "        self.w2i = {k.lower(): v for k, v in w2i.items()}\n",
    "        if self.w2i != w2i:\n",
    "            if len(self.w2i) == len(w2i):\n",
    "                warnings.warn(\n",
    "                    'Your w2i dict contains 1 or more uppercase characters. '\n",
    "                    'Our current implementation force-lowercases everything. '\n",
    "                    'You don\\'t have any collisions (e.g. \"Dog\" and \"dog\") so '\n",
    "                    'this should be okay, just keep this behavior in mind.'\n",
    "                )\n",
    "            else:\n",
    "                raise ValueError(\n",
    "                    'Our current implementation force-lowercases your w2i '\n",
    "                    'dict and yours appears to contain a collision (e.g. '\n",
    "                    '\"Dog\" and \"dog\"). We tentatively plan to allow cased '\n",
    "                    'keys in the future.'\n",
    "                )\n",
    "                \n",
    "        self.i2w = [w for w, i in \n",
    "                    sorted(self.w2i.items(), key=lambda x: x[1])]\n",
    "        if len(self.w2i) != len(self.i2w):\n",
    "            warnings.warn(\n",
    "                'Some keys in your w2i share an index. Mapping keys to IDs '\n",
    "                'should still work (if this is your intent) but reversing the '\n",
    "                'operation may produce unexpected results (e.g. if \"dog\" and '\n",
    "                '\"bulldog\" both map to index 0, it\\'s unclear whether index 0 '\n",
    "                'should be decoded as \"dog\" or \"bulldog\").'\n",
    "            )\n",
    "            \n",
    "        self.n_embeddings, self.dim = self.mat.shape\n",
    "        # Sets \"pca\" and \"mat_2d\" attributes.\n",
    "        self._validate_or_fit_pca(pca)\n",
    "        \n",
    "    def _validate_or_fit_pca(self, pca=None):\n",
    "        \"\"\"Compresses embedding matrix using PCA. If an sklearn pca object is\n",
    "        passed in, we'll check that it's been fit already. We make this its\n",
    "        own method for cases where we perform inplace operations on the \n",
    "        embedding matrix (e.g. self.normed) because these require PCA to be\n",
    "        refit).\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        pca: None or sklearn.decomposition.PCA\n",
    "            If not None, the object should already be fitted.\n",
    "        \"\"\"\n",
    "        if pca is None:\n",
    "            pca = PCA(n_components=2).fit(self.mat)\n",
    "        else:\n",
    "            check_is_fitted(pca)\n",
    "        self.pca = pca\n",
    "        self.mat_2d = self.pca.transform(self.mat)\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",
    "            \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_word2vec(cls, w2vec, w2i=None):\n",
    "        if w2i is None:\n",
    "            mat = np.vstack(list(w2vec.values()))\n",
    "            w2i = {k: i for i, k in enumerate(w2vec)}\n",
    "        else:\n",
    "            mat = np.vstack([w2vec[word] for word, i\n",
    "                             in sorted(w2i.items(), key=lambda x: x[1])])\n",
    "        return cls(mat=mat, w2i=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\n",
    "        in 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",
    "        # No need to save mat_2d since pca can quickly transform `mat`.\n",
    "        data = dict(mat=self.mat,\n",
    "                    w2i=self.w2i,\n",
    "                    pca=self.pca)\n",
    "        save(data, path, verbose=verbose)\n",
    "        \n",
    "    def normed(self, inplace=False):\n",
    "        \"\"\"Create new Embeddings object where all vectors have unit norm.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        inplace: bool\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        Embeddings or None: If inplace is False, return a new Embeddings\n",
    "        object with the same indices. If it's True, return None. In either \n",
    "        case, the vectors of the resulting Embeddings object will have unit\n",
    "        norm.\n",
    "        \"\"\"\n",
    "        normed_mat = self.mat / self.norm(self.mat)[:, None]\n",
    "        if inplace:\n",
    "            self.mat = normed_mat\n",
    "            self._validate_or_fit_pca()\n",
    "        else:\n",
    "            return type(self)(normed_mat, self.w2i)\n",
    "    \n",
    "    def subset(self, n, recompute_2d=True):\n",
    "        \"\"\"Subset Embeddings to top n words. Nice way to see how results of\n",
    "        all other methods (e.g. nearest_neighbors) would change if we used a\n",
    "        smaller vocabulary.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n: int\n",
    "            Top n embeddings (indices 0 through n-1) will be included in \n",
    "            subset.\n",
    "        recompute_2d: bool\n",
    "            If True, a new 2D matrix will be computed only using the subset,\n",
    "            meaning information about the excluded embeddings will be ignored.\n",
    "            If False, a subset of the existing embeddings will be used.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Embeddings: Same as self but now with length n.\n",
    "        \"\"\"\n",
    "        return type(self)(\n",
    "            self.mat[:n],\n",
    "            {k: v for k, v in sorted(self.w2i.items(), key=lambda x: x[1])\n",
    "             if v < n},\n",
    "            pca=None if recompute_2d else self.pca\n",
    "        )\n",
    "\n",
    "    @dispatch(str)\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.get(word)\n",
    "        if idx is not None:\n",
    "            return self.mat[idx]\n",
    "        \n",
    "    @dispatch(Iterable)\n",
    "    def vec(self, words):\n",
    "        \"\"\"Get embedding vectors for a list of words and return them as a \n",
    "        single numpy array. Note that all words must be present here: we want\n",
    "        to guarantee the output has the same number of rows as the input.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        words: list[str]\n",
    "            Input words to look up embeddings for.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Embeddings corresponding to the input words. \n",
    "        Shape (len(words), emb.dim).\n",
    "        \"\"\"\n",
    "        # Don't just delegate to the other `vec` method because we want to \n",
    "        # ensure all words are present.\n",
    "        return np.vstack([self.mat[self[word]] for word in words])\n",
    "\n",
    "    @dispatch(str)\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.get(word)\n",
    "        if idx is not None:\n",
    "            return self.mat_2d[idx]\n",
    "        \n",
    "    @dispatch(Iterable)\n",
    "    def vec_2d(self, words):\n",
    "        \"\"\"Look up the compressed embeddings for multiple words \n",
    "        (PCA was used to shrink dimensionality to 2). Note that all words must \n",
    "        be present here: we want to guarantee the output has the same number \n",
    "        of rows as the input.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        words: list[str]\n",
    "            Input words to look up embeddings for.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.array: Shape (len(words), emb.dim). Row i corresponds to words[i].\n",
    "        \"\"\"\n",
    "        # Don't just delegate to the other `vec` method because we want to \n",
    "        # ensure all words are present.\n",
    "        return np.vstack([self.mat_2d[self[word]] for word in words])\n",
    "        \n",
    "    @staticmethod\n",
    "    def distance(vec1, vec2, distance='cosine'):\n",
    "        \"\"\"Find distance between two vectors.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        distance: str\n",
    "            One of ('cosine', 'euclidean', 'manhattan').\n",
    "        \"\"\"\n",
    "        if distance == 'euclidean':\n",
    "            dists = Embeddings.norm(vec1 - vec2)\n",
    "        elif distance == 'cosine':\n",
    "            dists = Embeddings.cosine_distance(vec1, vec2)\n",
    "        elif distance == 'manhattan':\n",
    "            dists = Embeddings.manhattan_distance(vec1, vec2)\n",
    "        # Let arrays have numpy dtypes, but scalars will just be floats.\n",
    "        return dists if isinstance(dists, Iterable) else float(dists)\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",
    "        return self.distance(self.mat, vec, distance=distance)\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\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",
    "        \"\"\"Find the most similar words to a given word's vector. \n",
    "        This is the internal function behind `nearest_neighbors`, so you pass\n",
    "        in a vector instead of a word.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        vec: np.array\n",
    "        n: int\n",
    "        distance: str\n",
    "            One of ('cosine', 'euclidean', 'manhattan').\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",
    "        # First convert to float, otherwise we get np.float32 or np.float64\n",
    "        # scalars which can cause annoying bugs in APIs or dash apps.\n",
    "        return {self.i2w[i]: round(float(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",
    "        \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 from 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",
    "            \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 throws 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+1-trivial,\n",
    "                                            skip_first=False, **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. Computes bag of words vector by averaging vectors for all \n",
    "        input words.\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` method that lets us pass in vectors \n",
    "        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, exclude_args=True, **kwargs):\n",
    "        \"\"\"Wrapper to `cbow` method. 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 \n",
    "        will not include those words even if they are the closest to the mean\n",
    "        embedding) unless you set exclude_args=False. The idea here is to\n",
    "        find additional words that may be similar to the group you've passed \n",
    "        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",
    "        w2dist = 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: w2dist[word] for word in \n",
    "                [w for w in w2dist if not exclude_args or w not in args][:n]}\n",
    "    \n",
    "    @valuecheck\n",
    "    def matching_keys(self, *terms,\n",
    "                      mode:('standard', 'regex', 'ninja')='standard'):\n",
    "        \"\"\"Find keys (usually URLs, but could be used on words) containing\n",
    "        a given term/prefix/regex. This helps us do things like create theme\n",
    "        vectors (e.g. use this to find all sites related to \"games\" or \n",
    "        \"gaming\", then proceed to average their embeddings and potentially\n",
    "        expand the group even more by finding nearest neighbors).\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        terms: str\n",
    "            One or more strings to search for. This can be a word ('gaming'),\n",
    "            a partial word ('gam'), a regex ('^gam.*').\n",
    "        mode: str\n",
    "            standard - Performs exact string matching.\n",
    "            regex - Allows passing in regular expressions like '^sports'.\n",
    "            ninja - Use wordninja to remove spurious matches that violate \n",
    "                likely word boundaries. Mostly useful for things like URLs.\n",
    "                E.g. if you search for 'math' with mode='standard', the URL\n",
    "                'mathiasmiller.com' would match. With mode='ninja', it would \n",
    "                not because this URL seems to refer to a person, not the word\n",
    "                'math'.\n",
    "                \n",
    "        Returns\n",
    "        -------\n",
    "        list[str]: Keys (words, URLs, etc.) matching any of the given terms.\n",
    "        \"\"\"\n",
    "        if mode == 'regex':\n",
    "            match_fn = re.search\n",
    "        else:\n",
    "            def match_fn(term, key):\n",
    "                return term in key\n",
    "\n",
    "            # This is a bit slow so perform it as a second step after first\n",
    "            # filtering for URLs that contain the term at all.\n",
    "            if mode == 'ninja':\n",
    "                def post_match_fn(term, key):\n",
    "                    return term in wn.split(key)\n",
    "\n",
    "        keys = [key for key in self if any(match_fn(t, key) for t in terms)]\n",
    "        if mode == 'ninja':\n",
    "            keys = [k for k in keys if any(post_match_fn(t, k) \n",
    "                                           for t in terms)]\n",
    "        return keys\n",
    "    \n",
    "    def compare_distances(self, key, distance='cosine', as_df=True,\n",
    "                          sort_df=True, **vectors):\n",
    "        \"\"\"Compare how far a word/domain is from one or more vectors. Intended\n",
    "        for use with cbow results: e.g. checking if a word is closer to a\n",
    "        'liberal' vector or a 'conservative' vector.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        key: str\n",
    "            Word/domain to analyze.\n",
    "        distance: str\n",
    "            One of ('cosine', 'euclidean', 'manhattan'). Determines distance\n",
    "            method to use.\n",
    "        as_df: bool\n",
    "            If True, output is a dataframe. If False, it's a dict mapping \n",
    "            keys (strings) to distances (floats).\n",
    "        sort_df: bool\n",
    "            If True and as_df is True, the output df will be sorted by \n",
    "            distance from closest to furthest (relatively speaking - all \n",
    "            results will be relatively close).\n",
    "        vectors: np.array\n",
    "            One or more vectors to compare the input key's vector to. These\n",
    "            are kwargs rather than args because we need names for the \n",
    "            resulting df to show which distance corresponds to which vector.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        pd.DataFrame or dict[str, float]: Type depends on value of `as_df`.\n",
    "        \"\"\"\n",
    "        return self._compare_distances(self.vec(key), distance, as_df, \n",
    "                                       sort_df, **vectors)\n",
    "            \n",
    "    def _compare_distances(self, src_vec, distance='cosine', as_df=True,\n",
    "                           sort_df=True, **vectors):\n",
    "        \"\"\"Internal version of `compare_distances` that accepts a vector\n",
    "        rather than a word.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        src_vec: np.array\n",
    "            Embedding for the word/domain to analyze.\n",
    "        distance: str\n",
    "            One of ('cosine', 'euclidean', 'manhattan'). Determines distance\n",
    "            method to use.\n",
    "        as_df: bool\n",
    "            If True, output is a dataframe. If False, it's a dict mapping \n",
    "            keys (strings) to distances (floats).\n",
    "        sort_df: bool\n",
    "            If True and as_df is True, the output df will be sorted by \n",
    "            distance from closest to furthest (relatively speaking - all \n",
    "            results will be relatively close).\n",
    "        vectors: np.array\n",
    "            One or more vectors to compare the input key's vector to. These\n",
    "            are kwargs rather than args because we need names for the \n",
    "            resulting df to show which distance corresponds to which vector.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        pd.DataFrame or dict[str, float]: Type depends on value of `as_df`.\n",
    "        \"\"\"\n",
    "        d2dist = {k: self.distance(src_vec, v, distance)\n",
    "                  for k, v in vectors.items()}\n",
    "        if as_df:\n",
    "            d2dist = pd.DataFrame.from_dict(d2dist, orient='index',\n",
    "                                            columns=['score'])\n",
    "            if sort_df: d2dist = d2dist.sort_values('score', ascending=True)\n",
    "        return d2dist\n",
    "    \n",
    "    @valuecheck\n",
    "    def semantic_vector(\n",
    "            self, *queries, n=25, \n",
    "            include_queries:('always', 'never', 'auto')='always',\n",
    "            mode='standard', google_missing=False\n",
    "    ):\n",
    "        \"\"\"Mostly for domains rather than words: create a vector matching the\n",
    "        \"semantic theme\" of 1 or more input queries (usually several). \n",
    "        For example, to create a \"movie\" theme, you could pass in 'imdb.com',\n",
    "        'rottentomatoes.com', and 'letterboxd.com'.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        queries: str\n",
    "            These will often be keys in the Embeddings object (i.e. domains).\n",
    "            However, we also support words (e.g. \"movie\"), regular expressions\n",
    "            (e.g. \"^movie*\"), or even phrases (e.g. \"scary movies\"). Some\n",
    "            combinations of the above are supported: the only limitation as of\n",
    "            2/12/21 is you must choose a str matching mode (see \n",
    "            `self.matching_keys`). So you can pass in a mix of domains, words,\n",
    "            and phrases, but you can't use string matching for some words and\n",
    "            regex matching for others. Might be supported in the future if I\n",
    "            encounter situations where it seems useful.\n",
    "        n: int\n",
    "            Number of neighbors to find in `cbow_neighbors` method. Note that\n",
    "            this won't necessarily be the final number of keys returned - \n",
    "            `include_queries` will affect that too.\n",
    "        include_queries: str\n",
    "            Determines whether URLs retrieved from the initial step of query\n",
    "            string matching should be included in results. 'never' is useful\n",
    "            if you specifically want URLs that DON'T contain the queries \n",
    "            (e.g. music-related sites without \"music\" in the URL), but \n",
    "            I suspect 'always' may give better quality results. 'auto' will\n",
    "            allow these matches to be retained but won't force them to if they\n",
    "            aren't close to the final theme vector.\n",
    "        mode: str\n",
    "            Determines type of string matching used to find initial \"seed\"\n",
    "            domains (see `self.matching_keys`). Options are \n",
    "            ('standard', 'regex', 'ninja').\n",
    "        google_missing: bool\n",
    "            If True, terms that don't yield any string match results will be\n",
    "            googled and any domains on the first page of results will be used\n",
    "            if they're in our vocabulary.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        tuple[np.array, list[str]] or None: First item is the \"bag of words\"\n",
    "        array created by averaging over n neighbors of the input queries.\n",
    "        Second item is a list of the neighbor names sorted by similarity in\n",
    "        decreasing order.\n",
    "        \"\"\"\n",
    "        # If a query isn't in our embeddings, we assume it's intended to be a\n",
    "        # str matching query, not a full domain. If it is a full domain, this\n",
    "        # is a bit inefficient but it won't affect the final list of queries.\n",
    "        #  Don't use a list comprehension because we need to warn when no\n",
    "        # matches are found since we remove missing queries before calling\n",
    "        # `cbow_neighbors` (i.e. no built-in warning). This means our final\n",
    "        # list won't have missing queries, which is nice.\n",
    "        all_queries = []\n",
    "        for q in queries:\n",
    "            if q in self:\n",
    "                matches = [q]\n",
    "            else:\n",
    "                matches = self.matching_keys(q, mode=mode)\n",
    "                if not matches and google_missing:\n",
    "                    matches = domains_from_google_search(q)\n",
    "                if not matches:\n",
    "                    warnings.warn(f'No matches found for {q}.')\n",
    "            all_queries.extend(matches)\n",
    "        if not all_queries:\n",
    "            warnings.warn('Queries yielded zero matches.')\n",
    "            # Return None (not vector) for consistency with other methods.\n",
    "            return None, []\n",
    "\n",
    "        # Keep as dict for now for O(1) lookup time in case we use huge n.\n",
    "        matches = self.cbow_neighbors(*all_queries, n=n,\n",
    "                                      exclude_args=include_queries=='never')\n",
    "\n",
    "        if include_queries == 'always':\n",
    "            # Use dict as ordered set: remove duplicates but maintain order.\n",
    "            matches = dict.fromkeys(\n",
    "                list(matches) + [q for q in all_queries if q not in matches]\n",
    "            )\n",
    "        vec = self.cbow(*matches)\n",
    "        return vec, list(matches)\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",
    "    @staticmethod\n",
    "    def cosine_distance(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",
    "                    (Embeddings.norm(vec1) * Embeddings.norm(vec2)))\n",
    "\n",
    "    @dispatch(str)\n",
    "    def __getitem__(self, key):\n",
    "        \"\"\"When indexing with a string, this acts as a word->index method.\n",
    "        \n",
    "        Examples\n",
    "        --------\n",
    "        >>> emb['the']\n",
    "        1\n",
    "        \"\"\"\n",
    "        return self.w2i[key.lower()]\n",
    "    \n",
    "    @dispatch((int, slice))\n",
    "    def __getitem__(self, i):\n",
    "        \"\"\"When indexing with an integer, this acts as an index->word method.\n",
    "        \n",
    "        Examples\n",
    "        --------\n",
    "        >>> emb[1]\n",
    "        'the'\n",
    "        \n",
    "        >>> emb[:3]\n",
    "        ['a', the', 'is']\n",
    "        \"\"\"\n",
    "        return self.i2w[i]\n",
    "    \n",
    "    @dispatch(Iterable)\n",
    "    def __getitem__(self, keys):\n",
    "        \"\"\"Allows indexing in with a list of keys/indices. You can pass in a\n",
    "        mix of strings and integers though I can't imagine why that would be\n",
    "        necessary.\n",
    "        \n",
    "        Examples\n",
    "        --------\n",
    "        >>> emb[[1, 100, 7]]\n",
    "        ['the', 'frog', 'dog']\n",
    "        \n",
    "        >>> emb[['the', 'dog', 'frog']]\n",
    "        [1, 7, 100]\n",
    "        \"\"\"\n",
    "        return [self[key] for key in keys]\n",
    "            \n",
    "    def get(self, key, default=None):\n",
    "        \"\"\"Returns None if word is not present just like dict.get.\"\"\"\n",
    "        try:\n",
    "            return self.w2i[key.lower()]\n",
    "        except KeyError:\n",
    "            warnings.warn(f'{key} not in Embeddings.')\n",
    "            return default\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 in insertion order (may differ from\n",
    "        index order).\n",
    "        \"\"\"\n",
    "        yield from self.w2i.keys()\n",
    "        \n",
    "    def __eq__(self, obj):\n",
    "        if not isinstance(obj, Embeddings):\n",
    "            return False\n",
    "\n",
    "        ignore = {'pca'}\n",
    "        for k, v in vars(obj).items():\n",
    "            if k in ignore: continue\n",
    "            v_self = getattr(self, k)\n",
    "            if isinstance(v, np.ndarray):\n",
    "                if not np.allclose(v, v_self): return False\n",
    "            elif v != v_self: return False\n",
    "        return True\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'Embeddings(len={len(self)}, dim={self.dim})'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T21:13:47.047805Z",
     "start_time": "2021-04-27T21:13:46.964410Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got ValueError(Our current implementation force-lowercases your w2i dict and yours appears to contain a collision (e.g. \"Dog\" and \"dog\"). We tentatively plan to allow cased keys in the future.).\n"
     ]
    }
   ],
   "source": [
    "w2i = {'a': 0, 'A': 1, 'b': 2}\n",
    "mat = np.arange(6).reshape(3, 2)\n",
    "\n",
    "with assert_raises(ValueError):\n",
    "    emb = Embeddings(mat, w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T21:13:47.257371Z",
     "start_time": "2021-04-27T21:13:47.168697Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got ValueError(Your w2i dict has missing indices. We do not currently support gaps.).\n"
     ]
    }
   ],
   "source": [
    "w2i = {'B': 0, 'c': 1, 'a': 3}\n",
    "with assert_raises(ValueError):\n",
    "    emb = Embeddings(mat, w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T21:13:47.468295Z",
     "start_time": "2021-04-27T21:13:47.379405Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/harrisonmamin/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/ipykernel_launcher.py:44: UserWarning: Your w2i dict contains 1 or more uppercase characters. Our current implementation force-lowercases everything. You don't have any collisions (e.g. \"Dog\" and \"dog\") so this should be okay, just keep this behavior in mind.\n"
     ]
    }
   ],
   "source": [
    "w2i = {'a': 0, 'B': 1, 'c': 2}\n",
    "emb = Embeddings(mat, w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T21:13:47.573375Z",
     "start_time": "2021-04-27T21:13:47.490191Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/harrisonmamin/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/ipykernel_launcher.py:31: UserWarning: Your w2i dict is out of order (where ordered would mean the first key has id 0, the second has id 1, etc.). This should technically be fine but we recommend fixing it to be safe.\n",
      "/Users/harrisonmamin/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/ipykernel_launcher.py:44: UserWarning: Your w2i dict contains 1 or more uppercase characters. Our current implementation force-lowercases everything. You don't have any collisions (e.g. \"Dog\" and \"dog\") so this should be okay, just keep this behavior in mind.\n"
     ]
    }
   ],
   "source": [
    "w2i = {'B': 1, 'c': 2, 'a': 0, }\n",
    "emb = Embeddings(mat, w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T21:13:50.032672Z",
     "start_time": "2021-04-27T21:13:49.884838Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/harrisonmamin/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/ipykernel_launcher.py:31: UserWarning: Your w2i dict is out of order (where ordered would mean the first key has id 0, the second has id 1, etc.). This should technically be fine but we recommend fixing it to be safe.\n",
      "/Users/harrisonmamin/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/ipykernel_launcher.py:44: UserWarning: Your w2i dict contains 1 or more uppercase characters. Our current implementation force-lowercases everything. You don't have any collisions (e.g. \"Dog\" and \"dog\") so this should be okay, just keep this behavior in mind.\n"
     ]
    }
   ],
   "source": [
    "emb2 = Embeddings(mat, w2i)\n",
    "assert emb == emb2, 'Should evaluate as equal when w2i and mat are equal.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-17T00:07:46.418601Z",
     "start_time": "2021-01-17T00:07:46.412698Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def back_translate(text, to, from_lang='en'):\n",
    "    \"\"\"Translate a piece of text into another language, then back to English\n",
    "    for data augmentation purposes. This is rate limited but we now have a\n",
    "    pure ML version in the form of BacktranslateTransform.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    text: str\n",
    "        Text to back translate.\n",
    "    to: str\n",
    "        Language to translate to before translating back to English.\n",
    "    from_lang: str\n",
    "        Language of input text (usually 'en' for English).\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-26T20:52:36.750546Z",
     "start_time": "2020-11-26T20:52:36.679127Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@auto_repr\n",
    "class ParaphraseTransform:\n",
    "    \"\"\"Text transform that paraphrases input text as a method of data\n",
    "    augmentation. This is rather slow so it's recommended to precompute \n",
    "    samples and save them, but you could generate samples on the fly if \n",
    "    desired. One further downside of that approach is you'll have a huge\n",
    "    paraphrasing model on the GPU while (presumably) training another model.\n",
    "    \n",
    "    Other paraphrasing models exist on Model Hub but as of 11/14/2020, none of\n",
    "    the results compared favorably to this pegasus model, at least based on\n",
    "    a rough \"eyeball check\". While smaller and presumably faster, many of \n",
    "    these appear to require processing a single example at a time which \n",
    "    diminishes these gains. If you do attempt to use them, you'll likely need \n",
    "    to write a new class with a preprocessing method that does something like\n",
    "    the following:\n",
    "    _preprocess(text) -> 'paraphrase: {text}</s>'\n",
    "    I'm recording this here because many are missing documentation and it took\n",
    "    me some time to discover this.\n",
    "    \"\"\"\n",
    "\n",
    "    name = 'tuner007/pegasus_paraphrase'\n",
    "\n",
    "    def __init__(self, n=1, pipe=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n: int\n",
    "            Default number of samples to generate. You can override this in\n",
    "            __call__.\n",
    "        pipe: transformers Text2TextGenerationPipeline or None\n",
    "        \"\"\"\n",
    "        if pipe:\n",
    "            self.pipe = pipe\n",
    "            self.name = pipe.model.config._name_or_path\n",
    "        else:\n",
    "            self.pipe = Text2TextGenerationPipeline(\n",
    "                PegasusForConditionalGeneration.from_pretrained(self.name),\n",
    "                PegasusTokenizer.from_pretrained(self.name),\n",
    "                device=0 if torch.cuda.is_available() else -1\n",
    "            )\n",
    "        self.n = n\n",
    "            \n",
    "        assert type(self.pipe).__name__ == 'Text2TextGenerationPipeline'\n",
    "        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():\n",
    "            warnings.warn('The pipeline passed in is not using cuda. '\n",
    "                          'Did you mean to use the available GPU?')\n",
    "                \n",
    "    def _preprocess(self, text):\n",
    "        \"\"\"Does nothing (just want shared interface with other transforms).\"\"\"\n",
    "        return text\n",
    "    \n",
    "    @add_docstring(PreTrainedModel.generate)\n",
    "    def __call__(self, text, n=None, flat=True, **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "            Raw text to transform.\n",
    "        n: int or None\n",
    "            If None, use the default self.n.\n",
    "        flat: bool\n",
    "            If True, return flat list of strings. If False, return list of \n",
    "            nested lists where list i contains n augmentations of input i.\n",
    "        kwargs: any\n",
    "            Additional kwargs are passed to the model's text generation \n",
    "            method. Its docstring is included below for convenience.\n",
    "            \n",
    "        Returns\n",
    "        -------        \n",
    "        list: either a list with n strings per input string, or a list of\n",
    "        lists, each of length n, if flat=False.\n",
    "        \"\"\"\n",
    "        n = n or self.n\n",
    "        rows = [row['generated_text'] for row in \n",
    "                self.pipe(text, num_return_sequences=n, **kwargs)]\n",
    "        if listlike(text) and not flat: \n",
    "            rows = [rows[i*n:(i+1)*n] for i in range(len(text))]\n",
    "        return rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-11T04:36:40.431581Z",
     "start_time": "2021-02-11T04:36:40.428413Z"
    }
   },
   "outputs": [],
   "source": [
    "text = 'It was a beautiful sunny day and birds were chirping.'\n",
    "texts = ['Play fun games online for free! Watch your favorite movies and tv '\n",
    "         'shows here.', \n",
    "         'Bill hated school, especially math. His teacher was losing '\n",
    "         'patience with him.']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:26:43.488355Z",
     "start_time": "2020-11-22T05:26:08.661229Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/socket.py:660: ResourceWarning: unclosed <socket.socket fd=66, family=AddressFamily.AF_INET, type=SocketKind.SOCK_STREAM, proto=6, laddr=('192.168.1.119', 63347), raddr=('52.217.101.22', 443)>\n",
      "  self._sock = None\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/Users/hmamin/anaconda3/lib/python3.7/socket.py:660: ResourceWarning: unclosed <socket.socket fd=66, family=AddressFamily.AF_INET, type=SocketKind.SOCK_STREAM, proto=6, laddr=('192.168.1.119', 63354), raddr=('52.217.101.22', 443)>\n",
      "  self._sock = None\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n"
     ]
    }
   ],
   "source": [
    "p_tfm = ParaphraseTransform()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:26:54.281221Z",
     "start_time": "2020-11-22T05:26:43.493353Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Birds were singing in the sun.',\n",
       " 'Birds were singing on a nice sunny day.',\n",
       " 'Birds were singing on a sunny day.']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_tfm(text, n=3, temperature=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-26T20:52:39.667079Z",
     "start_time": "2020-11-26T20:52:39.602617Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@auto_repr\n",
    "class GenerativeTransform:\n",
    "    \"\"\"Text transform that truncates a piece of text and completes it using\n",
    "    a text generation model for the purposes of data augmentation. We\n",
    "    recommend precomputing samples and saving them for later use, but you\n",
    "    could generate samples on the fly if desired. Aside from speed, this\n",
    "    approach also has the drawback of having a text generation model on the \n",
    "    GPU while (presumably) training another model.\n",
    "    \"\"\"\n",
    "    \n",
    "    name = 'text-generation'\n",
    "    \n",
    "    def __init__(self, n=1, pipe=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n: int\n",
    "            Default number of samples to generate. You can override this in\n",
    "            __call__.\n",
    "        pipe: Transformers TextGenerationPipeline or None\n",
    "        \"\"\"\n",
    "        if pipe:\n",
    "            self.pipe = pipe\n",
    "            self.name = pipe.model.config._name_or_path\n",
    "        else:\n",
    "            self.pipe = pipeline(self.name, \n",
    "                                 device=0 if torch.cuda.is_available() \n",
    "                                 else -1)\n",
    "        self.n = n\n",
    "        \n",
    "        assert type(self.pipe).__name__ == 'TextGenerationPipeline'\n",
    "        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():\n",
    "            warnings.warn('The pipeline passed in is not using cuda. '\n",
    "                          'Did you mean to use the available GPU?')\n",
    "    \n",
    "    def _preprocess(self, text, drop=None, drop_pct=None, rand_low=None, \n",
    "                    rand_high=None, min_keep=3, return_tuple=False):\n",
    "        \"\"\"Truncate text so we can generate the ending.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "        drop: None or int\n",
    "            If provided, specifies the number of words to drop. We use a\n",
    "            simple \"split on spaces\" strategy since it's fast and simple.\n",
    "            Drop strategies occur in the signature in order of priority, so\n",
    "            if this is non-None it will override any values passed in for\n",
    "            drop_pct or rand_low/rand_high.\n",
    "        drop_pct: float or None\n",
    "            If provided, this should be a value between 0.0 and 1.0 specifying\n",
    "            the proportion of words to drop. \n",
    "        rand_low: int or None\n",
    "            If provided, specifies the minimum number of words to drop. The\n",
    "            max will be set by rand_high (which must also be provided with\n",
    "            rand_low). A random integer will be selected for each row of text.\n",
    "        rand_high: int or None\n",
    "            See rand_low: helps define bounds when randomly truncating rows\n",
    "            of text.\n",
    "        min_keep: int\n",
    "            The minimum number of words to keep. Sequences of this length or\n",
    "            shorter will therefore remain un-transformed. You could set this\n",
    "            to zero to enforce no minimum.\n",
    "        return_tuple: bool\n",
    "            If True, return a tuple where the first item is the truncated\n",
    "            text and the second item is the number of words masked. This is\n",
    "            rarely needed but it might be helpful if you want to use this for\n",
    "            some sort of self-supervised pre-training task.\n",
    "        \"\"\"\n",
    "        if listlike(text):\n",
    "            return [self._preprocess(row, drop, drop_pct, rand_low, rand_high,\n",
    "                                     min_keep, return_tuple) for row in text]\n",
    "        \n",
    "        tokens = text.split()\n",
    "        if len(tokens) <= min_keep:\n",
    "            n_drop = 0\n",
    "        else:\n",
    "            # Default is to truncate the last 20% of the sequence.\n",
    "            if drop:\n",
    "                n_drop = drop\n",
    "            elif drop_pct:\n",
    "                n_drop = int(drop_pct * len(tokens))\n",
    "            elif rand_low is not None and rand_high is not None:\n",
    "                n_drop = np.random.randint(rand_low, rand_high)\n",
    "            else:\n",
    "                n_drop = int(np.ceil(.2 * len(tokens)))\n",
    "            n_drop = np.clip(n_drop, 0, len(tokens) - min_keep)\n",
    "            tokens = tokens[:-n_drop]\n",
    "        truncated = ' '.join(tokens)\n",
    "        return (truncated, n_drop) if return_tuple else truncated\n",
    "    \n",
    "    @add_docstring(PreTrainedModel.generate)\n",
    "    def __call__(self, text, n=None, flat=True, min_length=2, max_length=7, \n",
    "                 drop=None, drop_pct=None, rand_low=None, rand_high=None, \n",
    "                 min_keep=3, **generate_kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "        n: int or None\n",
    "            Number of samples to generate for each input. Defaults to self.n\n",
    "            if None.\n",
    "        flat: bool\n",
    "            If True, return flat list of strings. If False, return list of \n",
    "            nested lists where list i contains n augmentations of input i.\n",
    "        min_length: int\n",
    "            Min number of tokens to generate.\n",
    "        max_length: int\n",
    "            Max number of tokens to generate. You could set this equal to \n",
    "            min_length to enforce a constant number.\n",
    "        drop: None or int\n",
    "            If provided, specifies the number of words to drop. We use a\n",
    "            simple \"split on spaces\" strategy since it's fast and simple.\n",
    "            Drop strategies occur in the signature in order of priority, so\n",
    "            if this is non-None it will override any values passed in for\n",
    "            drop_pct or rand_low/rand_high.\n",
    "        drop_pct: float or None\n",
    "            If provided, this should be a value between 0.0 and 1.0 specifying\n",
    "            the proportion of words to drop. \n",
    "        rand_low: int or None\n",
    "            If provided, specifies the minimum number of words to drop. The\n",
    "            max will be set by rand_high (which must also be provided with\n",
    "            rand_low). A random integer will be selected for each row of text.\n",
    "        rand_high: int or None\n",
    "            See rand_low: helps define bounds when randomly truncating rows\n",
    "            of text.\n",
    "        min_keep: int\n",
    "            The minimum number of words to keep. Sequences of this length or\n",
    "            shorter will therefore remain un-transformed. You could set this\n",
    "            to zero to enforce no minimum.\n",
    "        generate_kwargs: any\n",
    "            Forwarded to model's `generate` method. For convenience, its\n",
    "            docstring is provided below.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        list: either a list with n strings per input string, or a list of\n",
    "        lists, each of length n, if flat=False.\n",
    "        \"\"\"\n",
    "        n = n or self.n\n",
    "        if listlike(text):\n",
    "            res = [self(row, n, flat=flat, min_length=min_length, \n",
    "                        max_length=max_length, drop=drop, drop_pct=drop_pct, \n",
    "                        rand_low=rand_low, rand_high=rand_high, \n",
    "                        min_keep=min_keep, **generate_kwargs) for row in text]\n",
    "            return flatten(res) if flat else res\n",
    "    \n",
    "        # `generate` counts current length as part of min_length. \n",
    "        text = self._preprocess(text, drop, drop_pct, rand_low=rand_low, \n",
    "                                rand_high=rand_high, min_keep=min_keep)\n",
    "        n_curr = len(self.pipe.tokenizer.tokenize(text))\n",
    "        res = self.pipe(text, min_length=n_curr + min_length,\n",
    "                        max_length=n_curr + max_length,\n",
    "                        num_return_sequences=n, **generate_kwargs)\n",
    "        return [row['generated_text'] for row in res]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:10.480361Z",
     "start_time": "2020-11-22T05:34:10.429597Z"
    }
   },
   "outputs": [],
   "source": [
    "g_tfm = GenerativeTransform(n=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:27.870159Z",
     "start_time": "2020-11-22T05:34:27.800346Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Play fun games online for free! Watch your favorite movies and tv shows here.', 'Bill hated school, especially math. His teacher was losing patience with him.']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['Play fun games online', 'Bill hated school,']"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(texts)\n",
    "g_tfm._preprocess(texts, drop_pct=.75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:30.338284Z",
     "start_time": "2020-11-22T05:34:28.920514Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['It was a disaster. It was a disaster in',\n",
       " 'It was a really rough thing to say,\" the',\n",
       " 'It was a different kind of world. The city']"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_tfm(text, drop_pct=.75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:35.065519Z",
     "start_time": "2020-11-22T05:34:31.706349Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n",
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['Play fun games online for free! Watch your favorite movies and tv series from Disney in this game',\n",
       " 'Play fun games online for free! Watch your favorite movies and TV shows on Netflix, Amazon Video',\n",
       " 'Play fun games online for free! Watch your favorite movies and TV shows online at your mobile device',\n",
       " 'Bill hated school, especially math. His teacher was losing her eyesight due to a small',\n",
       " 'Bill hated school, especially math. His teacher was losing sight of the fact that his work',\n",
       " 'Bill hated school, especially math. His teacher was losing touch with him. Even though she']"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_tfm(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:39.115370Z",
     "start_time": "2020-11-22T05:34:36.943232Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n",
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[['Play fun games online for free! Watch your favorite movies and streams live for free! Read,',\n",
       "  'Play fun games online for free! Watch your favorite movies and TV shows play out! Explore The',\n",
       "  'Play fun games online for free! Watch your favorite movies and play them yourself.\\n\\nFree'],\n",
       " [\"Bill hated school, especially math. His teacher was losing sleep over his student's grades,\",\n",
       "  'Bill hated school, especially math. His teacher was losing her job at the time.\\n',\n",
       "  'Bill hated school, especially math. His teacher was losing control of his mind. He wouldn']]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = g_tfm(texts, flat=False)\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:34:56.989216Z",
     "start_time": "2020-11-22T05:34:54.969703Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n",
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[\"Play fun games online for free! Watch your favorite movies and sports videos and you'll\",\n",
       " 'Play fun games online for free! Watch your favorite movies and television shows in real time',\n",
       " 'Bill hated school, especially math. His teacher was losing her temper, so he',\n",
       " 'Bill hated school, especially math. His teacher was losing her job and he just']"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_tfm(texts, n=2, min_length=3, max_length=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:35:05.308625Z",
     "start_time": "2020-11-22T05:35:03.855093Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[\"It was a beautiful sunny day. I'm thankful. My\",\n",
       " 'It was a beautiful sunny day in August but we had to',\n",
       " 'It was a beautiful sunny day, and I went down to',\n",
       " 'It was a beautiful sunny spot in our neighbourhood. We had',\n",
       " 'It was a beautiful sunny Friday evening.\\n\\n\"It']"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_tfm(text, n=5, drop_pct=.5, min_keep=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-11T04:34:25.769674Z",
     "start_time": "2021-02-11T04:34:25.639124Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@auto_repr\n",
    "class FillMaskTransform:    \n",
    "    \"\"\"Text transform that masks one or more words in a piece of text and \n",
    "    fills them using RoBERTa for the purposes of data augmentation. We\n",
    "    recommend precomputing samples and saving them for later use, but you\n",
    "    could generate samples on the fly if desired. In addition to being slow,\n",
    "    that approach also entails having a mask filling model on the GPU while \n",
    "    (presumably) training another model.\n",
    "    \"\"\"\n",
    "\n",
    "    name = 'fill-mask'\n",
    "    MASK = '<mask>'\n",
    "    \n",
    "    def __init__(self, n=1, max_n=None, pipe=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n: int\n",
    "            n is intentionally bigger than the default n in __call__. This is\n",
    "            the number of candidates generated, so if we use strategy='random'\n",
    "            it makes sense for this to be larger.\n",
    "        max_n: int or None\n",
    "            Used as topk attribute when sampling generated candidates. This \n",
    "            must be >=n. Increasing this should not slow down generation. When\n",
    "            using strategy='random', you probably want this to be strictly >n.\n",
    "            Defaults to n+2.\n",
    "        pipe: transformers FillMaskPipeline\n",
    "            We let users pass in an existing pipeline since instantiation can\n",
    "            be slow.\n",
    "        \"\"\"\n",
    "        if pipe:\n",
    "            self.pipe = pipe\n",
    "            self.name = pipe.model.config._name_or_path\n",
    "        else:\n",
    "            self.pipe = pipeline(self.name, \n",
    "                                 device=0 if torch.cuda.is_available()\n",
    "                                 else -1)\n",
    "        # Set n before max_n.\n",
    "        self.n = n\n",
    "        self.max_n = max_n or n+2\n",
    "        \n",
    "        assert type(self.pipe).__name__ == 'FillMaskPipeline'\n",
    "        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():\n",
    "            warnings.warn('The pipeline passed in is not using cuda. '\n",
    "                          'Did you mean to use the available GPU?')\n",
    "    \n",
    "    def _preprocess(self, text, min_keep=3, errors='raise'):\n",
    "        \"\"\"Randomly mask one word from an input piece of text to prepare it \n",
    "        for RoBERTa to fill. Notice that even if the user chooses to mask\n",
    "        multiple words, each call to `_preprocess` only masks one since the\n",
    "        model can only fill one at a time.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "            One or more pieces of text to process.\n",
    "        min_keep: int\n",
    "            Minimum number of words to keep after truncating each piece of \n",
    "            text.\n",
    "        errors: str\n",
    "            If 'warn', we show a warning when min_keep is violated but allow\n",
    "            masking to take place. Any other value will result in an error\n",
    "            being raised.\n",
    "        \"\"\"\n",
    "        if listlike(text):\n",
    "            return [self._preprocess(row, min_keep, errors) for row in text]\n",
    "        \n",
    "        tokens = text.split()\n",
    "        if len(tokens) < min_keep + 1:\n",
    "            msg = (f'Text \"{text[:25]}...\" is too short to mask while '\n",
    "                   f'enforcing min_keep={min_keep}.')\n",
    "            # Err on side of caution: typos raise error too.\n",
    "            if errors == 'warn':\n",
    "                warnings.warn(msg)\n",
    "            else:\n",
    "                raise ValueError(msg)\n",
    "        \n",
    "        idx = np.random.choice(range(len(tokens)))\n",
    "        return ' '.join(self.MASK if i == idx else t \n",
    "                        for i, t in enumerate(tokens))\n",
    "    \n",
    "    @add_docstring(PreTrainedModel.generate)\n",
    "    def __call__(self, text, n=None, flat=True, n_mask=1, min_keep=3, \n",
    "                 return_all=False, errors='raise', strategy='best', **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "        n: int or None\n",
    "            Number of variations to return per piece of input text. If -1,\n",
    "            return all generated examples for the given mask count.\n",
    "            This can become very large when n_mask is large.\n",
    "            Example: if self.max_n=3, n=-1, and n_mask=4, we first mask once \n",
    "            and generate 3 samples. Then we mask each of those 3 and generate\n",
    "            a total of 9 samples, then 27, then finally 81 which is what will \n",
    "            be returned. The intermediate samples can be returned by\n",
    "            specifying `return_all=True`. \n",
    "        flat: bool\n",
    "            If True, return flat list of strings. If False, return list of\n",
    "            nested lists where list i contains n augmented versions of input \n",
    "            i.\n",
    "        n_mask: int\n",
    "            Number of words to mask. Because the model can only fill 1 masked\n",
    "            word at a time, `n_mask` forward passes will be performed.\n",
    "        min_keep: int\n",
    "            Minimum number of words to keep (presumably, you wouldn't want to\n",
    "            mask every word in an input sentence since that would strip it of\n",
    "            all existing meaning). This can be strictly enforced or not,\n",
    "            depending on your choice of `errors`.\n",
    "        return_all: bool\n",
    "            If True, return all intermediate generated samples rather than \n",
    "            just the final samples (e.g. if n_mask is 3, we first have to \n",
    "            generate samples with 1 masked word, then samples with 2 masked \n",
    "            words, and finally samples with 3 masked words. This is because\n",
    "            RoBERTa can only fill one masked word at a time.) See the \n",
    "            explanation of the `n` parameter for an example.\n",
    "        errors: str\n",
    "            One of ('warn', 'raise'). 'raise' will raise a ValueError if we're \n",
    "            about to violate `min_keep`. 'warn' will only show a warning if \n",
    "            this happens but will not strictly prevent it from occurring.\n",
    "        strategy: str\n",
    "            One of ('random', 'best'). The model will generate self.max_n\n",
    "            samples and if n < self.max_n, this means we need some way of \n",
    "            selecting which samples to keep. 'random' selects randomly without\n",
    "            replacement, while 'best' chooses the n most likely generations.\n",
    "            Note: when n_mask > 1, you should probably use strategy='random'\n",
    "            if you want relatively diverse results. If 'best', the benefit of\n",
    "            additional iterations is diminished because we are likely to end\n",
    "            up with very similar (or even identical) results.\n",
    "        kwargs: any\n",
    "            Forwarded to model's `generate` method. Its docstring is provided\n",
    "            below for convenience.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        list: either a list with n strings per input string, or a list of\n",
    "        lists, each of length n, if flat=False. This is slightly different if\n",
    "        return_all=True (see its description for details).\n",
    "        \"\"\"\n",
    "        # Make sure we generate adequate number of sequences. Model topk must\n",
    "        # be >= our desired n.\n",
    "        n = n or self.n\n",
    "        if n > self.max_n:\n",
    "            self.max_n = n\n",
    "            \n",
    "        # Each item will be a list of strings. Each string in res[i]\n",
    "        # will have i words changed. If text is a sequence of strings, we must\n",
    "        # handle each one separately because each is passed through pipeline\n",
    "        # repeatedly.\n",
    "        if listlike(text):\n",
    "            res = [self(row, n=n, flat=flat, n_mask=n_mask, min_keep=min_keep,\n",
    "                        return_all=return_all, errors=errors, \n",
    "                        strategy=strategy, **kwargs) for row in text]\n",
    "            return flatten(res) if flat else res\n",
    "\n",
    "        res = [[text]]\n",
    "        for i in range(n_mask):\n",
    "            seqs = self.pipe(self._preprocess(res[-1], min_keep=min_keep,\n",
    "                                              errors=errors))\n",
    "            # Transformers returns either list of dicts or list of list of \n",
    "            # dicts depending on whether input list has 1 item or multiple.\n",
    "            if isinstance(seqs[0], list): \n",
    "                seqs = [seq for group in seqs for seq in group]\n",
    "            text = [seq['sequence'].replace('<s>', '').replace('</s>', '') \n",
    "                    for seq in seqs]\n",
    "            \n",
    "            # Keep all generated samples when n is -1.\n",
    "            if n != -1:\n",
    "                if strategy == 'random':\n",
    "                    text = np.random.choice(text, n, replace=False)\n",
    "                elif strategy == 'best':\n",
    "                    text = text[:n]\n",
    "                else:\n",
    "                    raise ValueError('strategy should be \"random\" or \"best\".')\n",
    "            res.append(text)\n",
    "        if not return_all: res = res[n_mask]\n",
    "        return flatten(res) if flat else res\n",
    "    \n",
    "    @property\n",
    "    def max_n(self):\n",
    "        return self.pipe.topk\n",
    "    \n",
    "    @max_n.setter\n",
    "    def max_n(self, max_n):\n",
    "        \"\"\"Need to ensure the model generates enough options to return the\n",
    "        desired number of samples.\"\"\"\n",
    "        if not isinstance(max_n, int):\n",
    "            raise TypeError('max_n must be an integer.')\n",
    "        if max_n < self.n:\n",
    "            raise ValueError(f'max_n must be >= self.n (currently {self.n}.')\n",
    "        self.pipe.topk = max_n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-11T04:35:30.661049Z",
     "start_time": "2021-02-11T04:34:26.440412Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7824e28cfcd24df0b9908539b1b86e81",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/480 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8425da883cbb434d89c4a5f9e3c67848",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/331M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4641ac3b9f10422fa706b256a27b3a7b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c40b94fbb26c4799bd65ac18aa398ac7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bce36c36ec23435ebfba1ad1c0b6d92e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaForMaskedLM were not initialized from the model checkpoint at distilroberta-base and are newly initialized: ['lm_head.decoder.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "fm_tfm = FillMaskTransform(n=4, max_n=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-11T04:36:49.699263Z",
     "start_time": "2021-02-11T04:36:49.470594Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Play fun games online for free! Watch your favorite movies and TV shows here.',\n",
       " 'Play fun games online for free! Watch your favorite movies and tv shows here.',\n",
       " 'Play fun games online for free! Watch your favorite movies and television shows here.',\n",
       " 'Play fun games online for free! Watch your favorite movies and Broadway shows here.',\n",
       " 'Bill hated school, especially math. His teacher was losing patience with him.',\n",
       " 'Bill hates school, especially math. His teacher was losing patience with him.',\n",
       " 'Bill loved school, especially math. His teacher was losing patience with him.',\n",
       " 'Bill loves school, especially math. His teacher was losing patience with him.']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:05.303620Z",
     "start_time": "2020-11-22T05:40:04.659046Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['Play fun games online for free! Watch your favorite movies and tv shows here.',\n",
       "  'Play fun games online for free! Watch your favourite movies and tv shows here.',\n",
       "  'Play fun games online for free! Watch your own movies and tv shows here.',\n",
       "  'Play fun games online for free! Watch your favorites movies and tv shows here.'],\n",
       " ['Bill hated school, especially math. His teacher was losing patience with him.',\n",
       "  'Bill hated school, especially math. His teacher was having patience with him.',\n",
       "  'Bill hated school, especially math. His teacher was no patience with him.',\n",
       "  'Bill hated school, especially math. His teacher was lacking patience with him.']]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(texts, flat=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:06.560578Z",
     "start_time": "2020-11-22T05:40:06.133384Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['It was a beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and birds kept chirping.']"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(text, n=2, strategy='best')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:07.285970Z",
     "start_time": "2020-11-22T05:40:07.001621Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['It was a beautiful sunny day and they were chirping.',\n",
       " 'It was a beautiful sunny day and monkeys were chirping.']"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(text, n=2, strategy='random')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:08.155349Z",
     "start_time": "2020-11-22T05:40:07.750281Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['It was a beautiful sunny day and birds were chirping.'],\n",
       " ['It was a beautiful sunny day and birds were chirping.',\n",
       "  \"It's a beautiful sunny day and birds were chirping.\",\n",
       "  'It is a beautiful sunny day and birds were chirping.',\n",
       "  'It seemed a beautiful sunny day and birds were chirping.'],\n",
       " ['It was a beautiful sunny day and birds were chirping.',\n",
       "  'It was a beautiful sunny day and birds kept chirping.',\n",
       "  'It was a beautiful sunny day and birds started chirping.',\n",
       "  'It was a beautiful sunny day and birds are chirping.']]"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(text, n_mask=2, return_all=True, flat=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how quickly samples pile up when using n=-1, n_mask>1, and return_all=True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:10.481303Z",
     "start_time": "2020-11-22T05:40:09.708966Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['It was a beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and birds were chirping.',\n",
       " 'This was a beautiful sunny day and birds were chirping.',\n",
       " 'Today was a beautiful sunny day and birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day and birds were chirping.',\n",
       " 'Sunday was a beautiful sunny day and birds were chirping.',\n",
       " ' It was a beautiful sunny day and birds were chirping.',\n",
       " ' it was a beautiful sunny day and birds were chirping.',\n",
       " 'Saturday was a beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and we were chirping.',\n",
       " 'It was a beautiful sunny day and kids were chirping.',\n",
       " 'It was a beautiful sunny day and frogs were chirping.',\n",
       " 'It was a beautiful sunny day and chickens were chirping.',\n",
       " 'It was a beautiful sunny day and they were chirping.',\n",
       " 'It was a beautiful sunny day and children were chirping.',\n",
       " 'It was a beautiful sunny day and monkeys were chirping.',\n",
       " 'This was a beautiful sunny day and birds were chirping.',\n",
       " 'This was a beautiful sunny day and birds kept chirping.',\n",
       " 'This was a beautiful sunny day and birds are chirping.',\n",
       " 'This was a beautiful sunny day and birds started chirping.',\n",
       " 'This was a beautiful sunny day and birds began chirping.',\n",
       " 'This was a beautiful sunny day and birds happily chirping.',\n",
       " 'This was a beautiful sunny day and birds everywhere chirping.',\n",
       " 'This was a beautiful sunny day and birds stopped chirping.',\n",
       " 'Today was a beautiful sunny day and birds were chirping.',\n",
       " 'Today was a beautiful sunny afternoon and birds were chirping.',\n",
       " 'Today was a beautiful sunny morning and birds were chirping.',\n",
       " 'Today was a beautiful sunny spring and birds were chirping.',\n",
       " 'Today was a beautiful sunny summer and birds were chirping.',\n",
       " 'Today was a beautiful sunny weather and birds were chirping.',\n",
       " 'Today was a beautiful sunny sky and birds were chirping.',\n",
       " 'Today was a beautiful sunny evening and birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day and birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day where birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day when birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day as birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day; birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day while birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day, birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day - birds were chirping.',\n",
       " 'Sunday was a beautiful sunny day and birds were singing',\n",
       " 'Sunday was a beautiful sunny day and birds were flying',\n",
       " 'Sunday was a beautiful sunny day and birds were plentiful',\n",
       " 'Sunday was a beautiful sunny day and birds were nesting',\n",
       " 'Sunday was a beautiful sunny day and birds were soaring',\n",
       " 'Sunday was a beautiful sunny day and birds were abundant',\n",
       " 'Sunday was a beautiful sunny day and birds were buzzing',\n",
       " 'Sunday was a beautiful sunny day and birds were watching',\n",
       " 'It was a beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and we were chirping.',\n",
       " 'It was a beautiful sunny day and kids were chirping.',\n",
       " 'It was a beautiful sunny day and frogs were chirping.',\n",
       " 'It was a beautiful sunny day and chickens were chirping.',\n",
       " 'It was a beautiful sunny day and they were chirping.',\n",
       " 'It was a beautiful sunny day and children were chirping.',\n",
       " 'It was a beautiful sunny day and monkeys were chirping.',\n",
       " 'it was a beautiful sunny day and birds were chirping.',\n",
       " 'it was another beautiful sunny day and birds were chirping.',\n",
       " 'it was this beautiful sunny day and birds were chirping.',\n",
       " 'it was the beautiful sunny day and birds were chirping.',\n",
       " 'it was one beautiful sunny day and birds were chirping.',\n",
       " 'it was that beautiful sunny day and birds were chirping.',\n",
       " 'it was an beautiful sunny day and birds were chirping.',\n",
       " 'it was very beautiful sunny day and birds were chirping.',\n",
       " 'It was a beautiful sunny day and birds were chirping.',\n",
       " 'This was a beautiful sunny day and birds were chirping.',\n",
       " 'Today was a beautiful sunny day and birds were chirping.',\n",
       " 'Yesterday was a beautiful sunny day and birds were chirping.',\n",
       " 'Sunday was a beautiful sunny day and birds were chirping.',\n",
       " ' It was a beautiful sunny day and birds were chirping.',\n",
       " ' it was a beautiful sunny day and birds were chirping.',\n",
       " 'Saturday was a beautiful sunny day and birds were chirping.']"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(text, n=-1, n_mask=2, return_all=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-22T05:40:23.594057Z",
     "start_time": "2020-11-22T05:40:22.370826Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[['Play fun games online for free! Watch your favorite movies and tv shows here.'],\n",
       "  ['Play fun games online for free! Watch your favorite movies and tv shows here.',\n",
       "   'Play fun games online for free! Watch your favourite movies and tv shows here.'],\n",
       "  ['Play fun games online for free! Watch your favorite movies and TV shows here.',\n",
       "   'Play fun games online for free! Watch your favorite movies and tv shows here.']],\n",
       " [['Bill hated school, especially math. His teacher was losing patience with him.'],\n",
       "  ['Bill hated school, especially math. His teacher was losing patience with him.',\n",
       "   'Bill hated school, especially math. His teacher was losing patience for him.'],\n",
       "  ['Bill hated school, especially math. His teacher was losing patience with him.',\n",
       "   'Bill hated school, especially math. His teacher was having patience with him.']]]"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fm_tfm(texts, n=2, n_mask=2, return_all=True, flat=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also provide a convenience function that allows us to easily generate new samples from a source dataframe or csv while preserving any desired metadata so we can map each generated row to the correct label, ID, raw text, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-26T20:53:01.427319Z",
     "start_time": "2020-11-26T20:53:01.357390Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "NLP_TRANSFORMS = {\n",
    "    'fillmask': FillMaskTransform,\n",
    "    'paraphrase': ParaphraseTransform,\n",
    "    'generative': GenerativeTransform\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-11T04:48:35.299877Z",
     "start_time": "2021-02-11T04:48:35.245602Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class BacktranslateTransform:\n",
    "    \"\"\"Augment/perturb text inputs by translating them to a different language\n",
    "    and then back to English (this process can be repeated as many times as \n",
    "    you want by specifying multiple target languages). As of 2/10/21, this is\n",
    "    excluded from incendio's NLP_TRANSFORMS variable since its interface is a\n",
    "    little different from the other transforms: it has 2 pipelines, not 1, so\n",
    "    has variables `names` and `pipes` (both lists) instead of `name` and \n",
    "    `pipe`. It also has no `_preprocess` method.\n",
    "    \"\"\"\n",
    "\n",
    "    names = ['Helsinki-NLP/opus-mt-en-ROMANCE',\n",
    "             'Helsinki-NLP/opus-mt-ROMANCE-en']\n",
    "\n",
    "    language_codes = {\n",
    "        'es': 'spanish',\n",
    "        'fr': 'french',\n",
    "        'it': 'italian',\n",
    "        'pt': 'portuguese',\n",
    "        'pt_br': 'portuguese (brazil)',\n",
    "        'ro': 'romanian',\n",
    "        'ca': 'catalan',\n",
    "        'gl': 'galician',\n",
    "        'pt_BR': 'portuguese (brazil?)',\n",
    "        'la': 'latin',\n",
    "        'wa': 'walloon',\n",
    "        'fur': 'friulian (?)',\n",
    "        'oc': 'occitan',\n",
    "        'fr_CA': 'french (canada)',\n",
    "        'sc': 'sardianian',\n",
    "        'es_ES': 'spanish',\n",
    "        'es_MX': 'spanish (mexico)',\n",
    "        'es_AR': 'spanish (argentina)',\n",
    "        'es_PR': 'spanish (puerto rico)',\n",
    "        'es_UY': 'spanish (uruguay)',\n",
    "        'es_CL': 'spanish (chile)',\n",
    "        'es_CO': 'spanish (colombia)',\n",
    "        'es_CR': 'spanish (croatia)',\n",
    "        'es_GT': 'spanish (guatemala)',\n",
    "        'es_HN': 'spanish (honduras)',\n",
    "        'es_NI': 'spanish (nicaragua)',\n",
    "        'es_PA': 'spanish (panama)',\n",
    "        'es_PE': 'spanish (peru)',\n",
    "        'es_VE': 'spanish (venezuela)',\n",
    "        'es_DO': 'spanish (dominican republic)',\n",
    "        'es_EC': 'spanish (ecuador)',\n",
    "        'es_SV': 'spanish (el salvador)',\n",
    "        'an': 'aragonese',\n",
    "        'pt_PT': 'portuguese (portugal)',\n",
    "        'frp': 'franco provencal',\n",
    "        'lad': 'ladino',\n",
    "        'vec': 'venetian',\n",
    "        'fr_FR': 'france (france)',\n",
    "        'co': 'corsican',\n",
    "        'it_IT': 'italian (italy)',\n",
    "        'lld': 'ladin',\n",
    "        'lij': 'ligurian',\n",
    "        'lmo': 'lombard',\n",
    "        'nap': 'neapolitan',\n",
    "        'rm': 'rhaetian (?)',\n",
    "        'scn': 'sicilian',\n",
    "        'mwl': 'mirandese'\n",
    "    }\n",
    "\n",
    "    def __init__(self, to_langs, pipes=()):\n",
    "        \"\"\" \n",
    "        Parameters\n",
    "        ----------\n",
    "        to_langs: Iterable[str]\n",
    "            One or more language codes to use for backtranslation (see \n",
    "            self.language_codes for all options). They will be applied in\n",
    "            order: for instance, passing in ['es', 'fr'] will translate input\n",
    "            from \n",
    "            english -> spanish -> english -> french -> english. You can \n",
    "            override these later in specific calls but the value(s) you \n",
    "            provide here will be defaults.\n",
    "        pipes: Iterable[Pipeline]\n",
    "            Huggingface pipelines, the first of which translates English to\n",
    "            Romance languages and the second of which does the reverse. It's\n",
    "            usually easiest to let the Transform create these for you, but if\n",
    "            you already have them passing them in will be faster.\n",
    "        \"\"\"\n",
    "        if not pipes:\n",
    "            pipes = [TranslationPipeline(\n",
    "                        model=AutoModelForSeq2SeqLM.from_pretrained(name),\n",
    "                        tokenizer=AutoTokenizer.from_pretrained(name),\n",
    "                        device=1 - torch.cuda.is_available()\n",
    "                     ) for name in names]\n",
    "        self.pipes = pipes\n",
    "        self.to_langs = tolist(to_langs)\n",
    "\n",
    "    def __call__(self, text, intermediate=False, flat=True, to_langs=(), \n",
    "                 **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        text: str or Iterable[str]\n",
    "            The input pieces of text to translate.\n",
    "        intermediate: bool\n",
    "            If True, return all intermediate backtranslations if more than one\n",
    "            target language is provided. Otherwise, only return the final\n",
    "            backtranslation.\n",
    "        flat: bool\n",
    "            If True, return a flat list of strings. If False, return a list of\n",
    "            lists where item i contains all the intermediate backtranslations\n",
    "            (n languages in `to_langs` will generate n backtranslations). Note\n",
    "            that if intermediate=False, results will always be flat.\n",
    "        to_langs: Iterable[str]\n",
    "            One or more language codes to use for backtranslation. They will \n",
    "            be applied in order: for instance, passing in ['es', 'fr'] will\n",
    "            translate input from \n",
    "            english -> spanish -> english -> french -> english. If not\n",
    "            specified, this defaults to self.to_langs.\n",
    "        kwargs: any\n",
    "            Ignored. Just provided for consistency with other transforms.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        list[str] or list[list[str]]: Default is list of strings where item\n",
    "        i of output corresponds to item i of input. If intermediate=True and\n",
    "        flat=False, we get a list of lists where each nested list contains\n",
    "        n backtranslations of input i.\n",
    "        \"\"\"\n",
    "        text = tolist(text)\n",
    "        to_langs = tolist(to_langs) or self.to_langs\n",
    "        assert not set(to_langs) - set(self.language_codes), \\\n",
    "            'to_langs codes should all be present in self.language_codes.'\n",
    "\n",
    "        steps = []\n",
    "        for lang in to_langs:\n",
    "            text = [f'>>{lang}<< {t}' for t in text]\n",
    "            text = [row['translation_text'] for row in self.pipes[0](text)]\n",
    "            text = [row['translation_text'] for row in self.pipes[1](text)]\n",
    "            steps.append(text)\n",
    "        if intermediate:\n",
    "            steps = zip(*steps)\n",
    "            return flatten(steps) if flat else lmap(list, *steps)\n",
    "        else:\n",
    "            return text\n",
    "\n",
    "    def __repr__(self):\n",
    "        lang_str = \", \".join(repr(lang) for lang in self.to_langs)\n",
    "        return f'{func_name(self)}(to_langs=[{lang_str}])'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-26T21:17:56.919410Z",
     "start_time": "2020-11-26T21:17:56.823572Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@immutify_defaults\n",
    "def augment_text_df(source, transform='fillmask', dest=None, n=5, \n",
    "                    text_col='text', id_cols=(), nrows=None, tfm_kwargs={},\n",
    "                    call_kwargs={}):\n",
    "    \"\"\"Create augmented versions of a dataframe of text, optionally preserving\n",
    "    other columns for identification purposes. We recommend precomputing and\n",
    "    saving variations of your data rather than doing this on the fly in a \n",
    "    torch dataset since they can be rather space- and time-intensive. \n",
    "    Augmented versions of an input row should generally be kept in the same \n",
    "    training split: in order to keep the label the same, we usually want to \n",
    "    make relatively limited changes to the raw text (just enough to provide a\n",
    "    regularizing effect).\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    source: str, Path, or pd.DataFrame\n",
    "        If str or Path, this is a csv containing our text data. Alternatively,\n",
    "        you can pass in a dataframe itself.\n",
    "    transform: str or callable\n",
    "        If str, this must be one of the keys in `NLP_TRANSFORMS` from this \n",
    "        same module - this will be used to create a new transform object.\n",
    "        Alternatively, you can pass in a previously created object (NOT the\n",
    "        class). The default is the mask filling transform as it's relatively\n",
    "        quick and effective. 'paraphrase' may give better (but slower) \n",
    "        results. Anecdotally, 'generative' seems to provide lower quality\n",
    "        results, but perhaps by experimenting with hyperparameters it could\n",
    "        be more useful.\n",
    "    dest: str, Path, or None\n",
    "        If str or Path, this is where the output file will be saved to \n",
    "        (directories will be created as needed). If None, nothing will be \n",
    "        saved and the function will merely return the output DF for you to do \n",
    "        with as you wish.\n",
    "    n: int\n",
    "        Number of samples to generate for each raw row.\n",
    "    text_col: str\n",
    "        Name of column in DF containing the text to augment.\n",
    "    id_cols: Iterable[str]\n",
    "        Columns containing identifying information such as labels, row_ids, \n",
    "        etc. These also help us map the augmented text rows to their\n",
    "        corresponding raw rows.\n",
    "    nrows: int or None\n",
    "        Max number of rows from the source DF to generate text for. Useful for\n",
    "        testing (equivalently, you could pass in df.head(nrows) and leave this \n",
    "        as None).\n",
    "    tfm_kwargs: dict\n",
    "        Arguments to pass to `transform`'s constructor. These are ignored when\n",
    "        passing in a transform object rather than a string.\n",
    "    call_kwargs: dict\n",
    "        Arguments to pass to the __call__ method of `transform` to affect\n",
    "        the augmentation process.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    pd.DataFrame: DF of generated text with columns `text_col` and `id_cols`.\n",
    "    By default, this will have 5x the rows as your source DF, but this can\n",
    "    easily be adjusted through the `nrows` parameter.\n",
    "    \"\"\"\n",
    "    # Load data.\n",
    "    if isinstance(source, (str, Path)):\n",
    "        df = pd.read_csv(Path(source), usecols=[text_col] + list(id_cols), \n",
    "                         nrows=nrows)\n",
    "    elif isinstance(source, pd.DataFrame):\n",
    "        df = source.head(nrows)\n",
    "    else:\n",
    "        raise TypeError('`source` must be a str/Path or pd.DataFrame.')\n",
    "        \n",
    "    # Prepare for output file if necessary.\n",
    "    if isinstance(dest, (str, Path)):\n",
    "        dest = Path(dest)\n",
    "        os.makedirs(dest.parent, exist_ok=True)\n",
    "    elif dest is not None:\n",
    "        raise ValueError('`dest` must be a str/Path containing the output '\n",
    "                         'file name to create, or None if you just want to '\n",
    "                         'return a df.')\n",
    "\n",
    "    # For simplicity, we stick to one transform at a time. Slow to load so at\n",
    "    # least for now, let user pass in the transform itself.\n",
    "    transform = NLP_TRANSFORMS[transform](n=n, **tfm_kwargs) \\\n",
    "        if isinstance(transform, str) else transform\n",
    "    \n",
    "    # Generate new variations of input text.\n",
    "    res = transform(df[text_col].tolist(), **{**call_kwargs, 'flat': True})\n",
    "    res = pd.DataFrame(res, columns=[text_col])\n",
    "    \n",
    "    # Attach identifier columns to output (e.g. we usually want to store \n",
    "    # labels and or sample IDs. Most of our augmentation methods make \n",
    "    # relatively minor changes to the input so all variations of 1 input \n",
    "    # should remain in the same set, usually training.).\n",
    "    if id_cols:\n",
    "        df_id = pd.concat([df[col].repeat(res.shape[0] // df.shape[0])\n",
    "                           for col in id_cols], axis=1).reset_index(drop=True)\n",
    "        res = pd.concat([df_id, res], axis=1)\n",
    "        \n",
    "    # Optionally save output.\n",
    "    if dest: res.to_csv(dest, index=False)\n",
    "    return res"
   ]
  }
 ],
 "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.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
