# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/05_nlp.ipynb (unless otherwise specified).

__all__ = ['tokenizer', 'tokenize', 'tokenize_many', 'Vocabulary', 'domain', 'domains_from_google_search', 'Embeddings',
           'back_translate', 'postprocess_embeddings', 'compress_embeddings', 'ParaphraseTransform',
           'GenerativeTransform', 'FillMaskTransform', 'NLP_TRANSFORMS', 'BacktranslateTransform', 'augment_text_df']


# Cell
from bs4 import BeautifulSoup
from collections import Counter
from collections.abc import Iterable
from functools import partial
from multipledispatch import dispatch
import multiprocessing
import numpy as np
import os
import pandas as pd
from pathlib import Path
import re
import requests
from sklearn.decomposition import PCA
from sklearn.utils.validation import check_is_fitted
import spacy
from textblob import TextBlob
import torch
from tqdm.auto import tqdm
from transformers import PegasusForConditionalGeneration, PegasusTokenizer, \
    PegasusTokenizerFast, Text2TextGenerationPipeline, \
    AutoModelForSeq2SeqLM, AutoTokenizer, TranslationPipeline, pipeline
from tldextract import extract
from transformers.modeling_utils import PreTrainedModel
import warnings
import wordninja as wn

from htools import save, load, add_docstring, tolist, auto_repr, listlike, \
    flatten, immutify_defaults, ifnone, item, lmap, func_name, valuecheck
from .utils import DEVICE


# Cell
tokenizer = partial(spacy.load, name='en_core_web_sm',
                    disable=('ner', 'parser', 'tagger'))


# Cell
def tokenize(text, nlp):
    """Word tokenize a single string.

    Parameters
    ----------
    x: str
        A piece of text to tokenize.
    nlp: spacy tokenizer, e.g. spacy.lang.en.English
        By default, a spacy tokenizer with a small English vocabulary
        is used. NER, parsing, and tagging are disabled. Any spacy
        tokenzer can be passed in, but keep in mind other configurations
        may slow down this function dramatically.

    Returns
    -------
    list[str]: List of word tokens from a single input string.
    """
    return [tok.text for tok in nlp(text)]


# Cell
def tokenize_many(rows, nlp=None, chunk=1_000):
    """Word tokenize a sequence of strings using multiprocessing. The max
    number of available processes are used.

    Parameters
    ----------
    rows: Iterable[str]
        A sequence of strings to tokenize. This could be a list, a column of
        a DataFrame, etc.
    nlp: spacy tokenizer, e.g. spacy.lang.en.English
        By default, a spacy tokenizer with a small English vocabulary
        is used. NER, parsing, and tagging are disabled. Any spacy
        tokenzer can be passed in, but keep in mind other configurations
        may slow down this function dramatically.
    chunk: int
        This determines how many items to send to multiprocessing at a time.
        The default of 1,000 is usually fine, but if you have extremely
        long pieces of text and memory is limited, you can always decrease it.
        Very small chunk sizes may increase processing time. Note that larger
        values will generally cause the progress bar to update more choppily.

    Returns
    -------
    list[list[str]]: Each nested list of word tokens corresponds to one
    of the input strings.
    """
    tokenize_ = partial(tokenize, nlp=nlp or tokenizer())
    length = len(rows)
    with multiprocessing.Pool() as p:
        res = list(tqdm(p.imap(tokenize_, rows, chunksize=chunk),
                        total=length))
    return res


# Cell
class Vocabulary:

    def __init__(self, w2idx, w2vec=None, idx_misc=None, corpus_counts=None,
                 all_lower=True):
        """Defines a vocabulary object for NLP problems, allowing users to
        encode text with indices or embeddings.

        Parameters
        -----------
        w2idx: dict[str, int]
            Dictionary mapping words to their integer index in a vocabulary.
            The indices must allow for idx_misc to be added to the dictionary,
            so in the default case this should have a minimum index of 2. If
            a longer idx_misc is passed in, the minimum index would be larger.
        w2vec: dict[str, np.array]
            Dictionary mapping words to their embedding vectors stored as
            numpy arrays (optional).
        idx_misc: dict
            A dictionary mapping non-word tokens to indices. If none is passed
            in, a default version will be used with keys for unknown tokens
            and padding. A customized version might pass in additional tokens
            for repeated characters or all caps, for example.
        corpus_counts: collections.Counter
            Counter dict mapping words to their number of occurrences in a
            corpus (optional).
        all_lower: bool
            Specifies whether the data you've passed in (w2idx, w2vec, i2w) is
            all lowercase. Note that this will NOT change any of this data. If
            True, it simply lowercases user-input words when looking up their
            index or vector.
        """
        if not idx_misc:
            idx_misc = {'<PAD>': 0,
                        '<UNK>': 1}
        self.idx_misc = idx_misc
        # Check that space has been left for misc keys.
        assert len(idx_misc) == min(w2idx.values())

        # Core data structures.
        self.w2idx = {**self.idx_misc, **w2idx}
        self.i2w = [word for word, idx in sorted(self.w2idx.items(),
                                                 key=lambda x: x[1])]
        self.w2vec = w2vec or dict()

        # Miscellaneous other attributes.
        if w2vec:
            self.dim = len(w2vec[self[-1]])
        else:
            self.dim = 1
        self.corpus_counts = corpus_counts
        self.embedding_matrix = None
        self.w2vec['<UNK>'] = np.zeros(self.dim)
        self.all_lower = all_lower

    @classmethod
    def from_glove_file(cls, path, max_lines=float('inf'), idx_misc=None):
        """Create a new Vocabulary object by loading GloVe vectors from a text
        file. The embeddings are all lowercase so the user does not have the
        option to set the all_lower parameter.

        Parameters
        -----------
        path: str
            Path to file containing glove vectors.
        max_lines: int, float (optional)
            Loading the GloVe vectors can be slow, so for testing purposes
            it can be helpful to read in a subset. If no value is provided,
            all 400,000 lines in the file will be read in.
        idx_misc: dict
            Map non-standard tokens to indices. See constructor docstring.
        """
        w2idx = dict()
        w2vec = dict()
        misc_len = 2 if not idx_misc else len(idx_misc)

        with open(path, 'r') as f:
            for i, line in enumerate(f):
                if i >= max_lines:
                    break
                word, *values = line.strip().split(' ')
                w2idx[word] = i + misc_len
                w2vec[word] = np.array(values, dtype=np.float)

        return cls(w2idx, w2vec, idx_misc)

    @classmethod
    def from_tokens(cls, tokens, idx_misc=None, all_lower=True):
        """Construct a Vocabulary object from a list or array of tokens.

        Parameters
        -----------
        tokens: list[str]
            The word-tokenized corpus.
        idx_misc: dict
            Map non-standard tokens to indices. See constructor docstring.
        all_lower: bool
            Specifies whether your tokens are all lowercase.

        Returns
        --------
        Vocabulary
        """
        misc_len = 2 if not idx_misc else len(idx_misc)
        counts = Counter(tokens)
        w2idx = {word: i for i, (word, freq)
                 in enumerate(counts.most_common(), misc_len)}
        return cls(w2idx, idx_misc=idx_misc, corpus_counts=counts,
                   all_lower=all_lower)

    @staticmethod
    def from_pickle(path):
        """Load a previously saved Vocabulary object.

        Parameters
        -----------
        path: str
            Location of pickled Vocabulary file.

        Returns
        --------
        Vocabulary
        """
        return load(path)

    def save(self, path, verbose=True):
        """Pickle Vocabulary object for later use. We can then quickly load
        the object using torch.load(path), which can be much faster than
        re-computing everything when the vocab size becomes large.

        Parameters
        -----------
        path: str
            Where to save the output file.
        verbose: bool
            If True, print message showing where the object was saved to.
        """
        save(self, path, verbose)

    def filter_tokens(self, tokens, max_words=None, min_freq=0, inplace=False,
                      recompute=False):
        """Filter your vocabulary by specifying a max number of words or a min
        frequency in the corpus. When done in place, this also sorts vocab by
        frequency with more common words coming first (after idx_misc).

        Parameters
        -----------
        tokens: list[str]
            A tokenized list of words in the corpus (must be all lowercase
            when self.all_lower=True, such as when using GloVe vectors). There
            is no need to hold out test data here since we are not using
            labels.
        max_words: int (optional)
            Provides an upper threshold for the number of words in the
            vocabulary. If no value is passed in, no maximum limit will be
            enforced.
        min_freq: int (optional)
            Provides a lower threshold for the number of times a word must
            appear in the corpus to remain in the vocabulary. If no value is
            passed in, no minimum limit will be enforced.

            Note that we can specify values for both max_words and min_freq
            if desired. If no values are passed in for either, no pruning of
            the vocabulary will be performed.
        inplace: bool
            If True, will change the object's attributes
            (w2idx, w2vec, and i2w) to reflect the newly filtered vocabulary.
            If False, will not change the object, but will simply compute word
            counts and return what the new w2idx would be. This can be helpful
            for experimentation, as we may want to try out multiple values of
            min_freq to decide how many words to keep. After the first call,
            the attribute corpus_counts can also be examined to help determine
            the desired vocab size.
        recompute: bool
            If True, will calculate word counts from the given tokens. If
            False (the default), this will use existing counts if there are
            any.

            The idea is that if we call this method, then realize we want
            to change the corpus, we should calculate new word counts.
            However, if we are simply calling this method multiple times on
            the same corpus while deciding on the exact vocab size we want,
            we should not recompute the word counts.

        Returns
        --------
        dict or None: When called inplace, nothing is returned. When not
        inplace,
        """
        misc_len = len(self.idx_misc)
        if recompute or not self.corpus_counts:
            self.corpus_counts = Counter(tokens)
        filtered = {word: i for i, (word, freq)
                    in enumerate(self.corpus_counts.most_common(max_words),
                                 misc_len)
                    if freq >= min_freq}
        filtered = {**self.idx_misc, **filtered}

        if inplace:
            # Relies on python3.7 dicts retaining insertion order.
            self.i2w = list(filtered.keys())
            self.w2idx = filtered
            self.w2vec = {word: self.vector(word) for word in filtered}
        else:
            return filtered

    def build_embedding_matrix(self, inplace=False):
        """Create a 2D numpy array of embedding vectors where row[i]
        corresponds to word i in the vocabulary. This can be used to
        initialize weights in the model's embedding layer.

        Parameters
        -----------
        inplace: bool
            If True, will store the output in the object's embedding_matrix
            attribute. If False (default behavior), will simply return the
            matrix without storing it as part of the object. In the
            recommended case where inplace==False, we can store the output
            in another variable which we can use to initialize the weights in
            Torch, then delete the object and free up memory using
            gc.collect().
        """
        emb = np.zeros((len(self), self.dim))
        for i, word in enumerate(self):
            emb[i] = self.vector(word)

        if inplace:
            self.embedding_matrix = emb
        else:
            return emb

    def idx(self, word):
        """This will map a word (str) to its index (int) in the vocabulary.
        If a string is passed in and the word is not present, the index
        corresponding to the <UNK> token is returned.

        Parameters
        -----------
        word: str
            A word that needs to be mapped to an integer index.

        Returns
        --------
        int: The index of the given word in the vocabulary.

        Examples
        ---------
        >>> vocab.idx('the')
        2
        """
        if self.all_lower and word not in self.idx_misc:
            word = word.lower()
        return self.w2idx.get(word, self.w2idx['<UNK>'])

    def vector(self, word):
        """This maps a word to its corresponding embedding vector. If not
        contained in the vocab, a vector of zeros will be returned.

        Parameters
        -----------
        word: str
            A word that needs to be mapped to a vector.

        Returns
        --------
        np.array
        """
        if self.all_lower and word not in self.idx_misc:
            word = word.lower()
        return self.w2vec.get(word, self.w2vec['<UNK>'])

    def encode(self, text, nlp, max_len, pad_end=True, trim_start=True):
        """Encode text so that each token is replaced by its integer index in
        the vocab.

        Parameters
        -----------
        text: str
            Raw text to be encoded.
        nlp: spacy.lang.en.English
            Spacy tokenizer. Typically want to disable 'parser', 'tagger', and
            'ner' as they aren't used here and slow down the encoding process.
        max_len: int
            Length of output encoding. If text is shorter, it will be padded
            to fit the specified length. If text is longer, it will be
            trimmed.
        pad_end: bool
            If True, add padding to the end of short sentences. If False, pad
            the start of these sentences.
        trim_start: bool
            If True, trim off the start of sentences that are too long. If
            False, trim off the end.

        Returns
        --------
        np.array[int]: Array of length max_len containing integer indices
            corresponding to the words passed in.
        """
        output = np.ones(max_len) * self.idx('<PAD>')
        encoded = [self.idx(tok.text) for tok in nlp(text)]

        # Trim sentence in case it's longer than max_len.
        if len(encoded) > max_len:
            if trim_start:
                encoded = encoded[len(encoded) - max_len:]
            else:
                encoded = encoded[:max_len]

        # Replace zeros at start or end, depending on choice of pad_end.
        if pad_end:
            output[:len(encoded)] = encoded
        else:
            output[max_len-len(encoded):] = encoded
        return output.astype(int)

    def decode(self, idx, join=True, sep=' '):
        """Convert a list of indices to a string or list of words/tokens.

        Parameters
        -----------
        idx: list[int]
            A list of integers indexing into the vocabulary. This will often
            be the output of the encode() method.
        join: bool
            If True, return a single string. If False, return a list of
            strings.
        sep: str
            If join is True, this determines what character is used to join
            tokens. Word tokens will usually be joined by a space, but some
            tokenization schemes include spaces and can be joined with an
            empty string ('').

        Returns
        --------
        str or list[str]: A string or list of strings (words/tokens)
        reconstructed by indexing into the vocabulary.
        """
        tokens = [self[i] for i in idx]
        if join: return sep.join(tokens)
        return tokens

    def __getitem__(self, i):
        """This will map an index (int) to a word (str).

        Parameters
        -----------
        i: int
            Integer index for a word.

        Returns
        --------
        str: Word corresponding to the given index.

        Examples
        ---------
        >>> vocab = Vocabulary(w2idx, w2vec)
        >>> vocab[1]
        '<UNK>'
        """
        return self.i2w[i]

    def __len__(self):
        """Number of words in vocabulary."""
        return len(self.w2idx)

    def __iter__(self):
        for word in self.w2idx.keys():
            yield word

    def __contains__(self, word):
        return word in self.w2idx.keys()

    def __eq__(self, obj):
        if not isinstance(obj, Vocabulary):
            return False

        ignore = {'w2vec', 'embedding_matrix'}
        attrs = [k for k, v in hdir(vocab).items()
                 if v == 'attribute' and k not in ignore]
        return all([getattr(self, attr) == getattr(obj, attr)
                    for attr in attrs])

    def __repr__(self):
        msg = f'Vocabulary({len(self)} words'
        if self.dim > 1:
            msg += f', {self.dim}-D embeddings'
        return msg + ')'


# Cell
def domain(url, strip_www=True, mode='registered_domain'):
    domain_ = getattr(extract(url), mode)
    return domain_.replace('www.', '') if strip_www else domain_


# Cell
def domains_from_google_search(term, drop_probable_defaults=True):
    """Get domains from first page of google sarch for a given term.

    Parameters
    ----------
    term: str
        Term to search for.
    drop_probable_defaults: bool
        If True, remove two domains which seem to be present on all google
        search results pages. I didn't see a simple identifier for which links
        on the page are search results so we're left with this clumsy
        filtering method (setting this to True means if 'accounts.google.com'
        is a legitimate search result, we'd still remove it).

    Returns
    -------
    set[str]: Domain names (no www) found on the first page of search results.
    """
    url = f'http://www.google.com/search?q={"+".join(term.split())}'
    r = requests.get(url)
    links = BeautifulSoup(r.content, 'lxml').find_all('a')
    # Use dict to remove duplicates while maintaining order.
    domains = dict.fromkeys(domain(link['href'].replace('/url?q=', ''))
                            for link in links
                            if link['href'].startswith('/url?q='))
    if drop_probable_defaults:
        for d in ('accounts.google.com', 'support.google.com'):
            domains.pop(d, None)
    return list(domains)


# Cell
class Embeddings:
    """Embeddings object. Lets us easily map word to index, index to
    word, and word to vector. We can use this to find similar words,
    build analogies, or get 2D representations for plotting. Generally,
    user-facing methods let us pass in strings, while internal versions (same
    name except prefixed with an underscore) allow us to pass in vectors.
    """

    def __init__(self, mat, w2i, pca=None):
        """
        Parameters
        ----------
        mat: str
            Numpy array of embeddings where row i corresponds to ID i
            in w2i.
        w2i: dict[str, int]
            Dictionary mapping word to its index in the vocabulary.
        pca: sklearn.decomposition.PCA or None
            If provided, this should be a PCA object with 2 components that
            was previously fit on `mat`. If None, a new object will be created
            and fit. This will let us plot embeddings in a way humans can
            visually parse.
        """
        self.mat = mat
        max_id = max(w2i.values())
        expected_ids = list(range(max_id + 1))
        if list(w2i.values()) != expected_ids:
            if sorted(w2i.values()) == expected_ids:
                warnings.warn(
                    '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.'
                )
            else:
                raise ValueError('Your w2i dict has missing indices. '
                                 'We do not currently support gaps.')

        self.w2i = {k.lower(): v for k, v in w2i.items()}
        if self.w2i != w2i:
            if len(self.w2i) == len(w2i):
                warnings.warn(
                    '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.'
                )
            else:
                raise 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.'
                )

        self.i2w = [w for w, i in
                    sorted(self.w2i.items(), key=lambda x: x[1])]
        if len(self.w2i) != len(self.i2w):
            warnings.warn(
                'Some keys in your w2i share an index. Mapping keys to IDs '
                'should still work (if this is your intent) but reversing the '
                'operation may produce unexpected results (e.g. if "dog" and '
                '"bulldog" both map to index 0, it\'s unclear whether index 0 '
                'should be decoded as "dog" or "bulldog").'
            )

        self.n_embeddings, self.dim = self.mat.shape
        # Sets "pca" and "mat_2d" attributes.
        self._validate_or_fit_pca(pca)

    def _validate_or_fit_pca(self, pca=None):
        """Compresses embedding matrix using PCA. If an sklearn pca object is
        passed in, we'll check that it's been fit already. We make this its
        own method for cases where we perform inplace operations on the
        embedding matrix (e.g. self.normed) because these require PCA to be
        refit).

        Parameters
        ----------
        pca: None or sklearn.decomposition.PCA
            If not None, the object should already be fitted.
        """
        if pca is None:
            pca = PCA(n_components=2).fit(self.mat)
        else:
            check_is_fitted(pca)
        self.pca = pca
        self.mat_2d = self.pca.transform(self.mat)

    @classmethod
    def from_text_file(cls, path, max_words=float('inf'), print_freq=10_000):
        """Create a new Embeddings object from a raw text file using the
        GloVe format (each row contains a word and its embedding as
        space-separated floats).

        Parameters
        ----------
        path: str
            Location of csv file containing GloVe vectors.
        max_words: int, float
            Set maximum number of words to read in from file. This can be used
            during development to reduce wait times when loading data.

        Returns
        -------
        Embeddings: Newly instantiated object.
        """
        w2i = dict()
        mat = []
        with open(path, 'r') as f:
            for i, line in enumerate(f):
                # Faster testing
                if i >= max_words: break
                word, *nums = line.strip().split()
                w2i[word] = i
                mat.append(np.array(nums, dtype=float))
                if i % print_freq == 0: print(i, word)
        return cls(np.array(mat), w2i)

    @classmethod
    def from_word2vec(cls, w2vec, w2i=None):
        if w2i is None:
            mat = np.vstack(list(w2vec.values()))
            w2i = {k: i for i, k in enumerate(w2vec)}
        else:
            mat = np.vstack([w2vec[word] for word, i
                             in sorted(w2i.items(), key=lambda x: x[1])])
        return cls(mat=mat, w2i=w2i)

    @classmethod
    def from_pickle(cls, path):
        """If an Embeddings object previously saved its data in a pickle file,
        loading it that way can avoid repeated computation.

        Parameters
        ----------
        path: str
            Location of pickle file.

        Returns
        -------
        Embeddings: Newly instantiated object using the data that was stored
        in the pickle file.
        """
        return cls(**load(path))

    def save(self, path, verbose=True):
        """Save data to a compressed pickle file. This reduces the amount of
        space needed for storage (the csv is much larger) and can let us
        avoid running PCA and building the embedding matrix again.

        Parameters
        ----------
        path: str
            Path that object will be saved to.
        verbose

        Returns
        -------
        None
        """
        # No need to save mat_2d since pca can quickly transform `mat`.
        data = dict(mat=self.mat,
                    w2i=self.w2i,
                    pca=self.pca)
        save(data, path, verbose=verbose)

    def normed(self, inplace=False):
        """Create new Embeddings object where all vectors have unit norm.

        Parameters
        ----------
        inplace: bool

        Returns
        -------
        Embeddings or None: If inplace is False, return a new Embeddings
        object with the same indices. If it's True, return None. In either
        case, the vectors of the resulting Embeddings object will have unit
        norm.
        """
        normed_mat = self.mat / self.norm(self.mat)[:, None]
        if inplace:
            self.mat = normed_mat
            self._validate_or_fit_pca()
        else:
            return type(self)(normed_mat, self.w2i)

    def subset(self, n, recompute_2d=True):
        """Subset Embeddings to top n words. Nice way to see how results of
        all other methods (e.g. nearest_neighbors) would change if we used a
        smaller vocabulary.

        Parameters
        ----------
        n: int
            Top n embeddings (indices 0 through n-1) will be included in
            subset.
        recompute_2d: bool
            If True, a new 2D matrix will be computed only using the subset,
            meaning information about the excluded embeddings will be ignored.
            If False, a subset of the existing embeddings will be used.

        Returns
        -------
        Embeddings: Same as self but now with length n.
        """
        return type(self)(
            self.mat[:n],
            {k: v for k, v in sorted(self.w2i.items(), key=lambda x: x[1])
             if v < n},
            pca=None if recompute_2d else self.pca
        )

    @dispatch(str)
    def vec(self, word):
        """Look up the embedding for a given word. Return None if not found.

        Parameters
        ----------
        word: str
            Input word to look up embedding for.

        Returns
        -------
        np.array: Embedding corresponding to the input word. If word not in
            vocab, return None.
        """
        idx = self.get(word)
        if idx is not None:
            return self.mat[idx]

    @dispatch(Iterable)
    def vec(self, words):
        """Get embedding vectors for a list of words and return them as a
        single numpy array. Note that all words must be present here: we want
        to guarantee the output has the same number of rows as the input.

        Parameters
        ----------
        words: list[str]
            Input words to look up embeddings for.

        Returns
        -------
        np.array: Embeddings corresponding to the input words.
        Shape (len(words), emb.dim).
        """
        # Don't just delegate to the other `vec` method because we want to
        # ensure all words are present.
        return np.vstack([self.mat[self[word]] for word in words])

    @dispatch(str)
    def vec_2d(self, word):
        """Look up the compressed embedding for a word (PCA was used to shrink
        dimensionality to 2). Return None if the word is not present in vocab.

        Parameters
        ----------
        word: str
            Input work to look up.

        Returns
        -------
        np.array: Compressed embedding of length 2. None if not found.
        """
        idx = self.get(word)
        if idx is not None:
            return self.mat_2d[idx]

    @dispatch(Iterable)
    def vec_2d(self, words):
        """Look up the compressed embeddings for multiple words
        (PCA was used to shrink dimensionality to 2). Note that all words must
        be present here: we want to guarantee the output has the same number
        of rows as the input.

        Parameters
        ----------
        words: list[str]
            Input words to look up embeddings for.

        Returns
        -------
        np.array: Shape (len(words), emb.dim). Row i corresponds to words[i].
        """
        # Don't just delegate to the other `vec` method because we want to
        # ensure all words are present.
        return np.vstack([self.mat_2d[self[word]] for word in words])

    @staticmethod
    def distance(vec1, vec2, distance='cosine'):
        """Find distance between two vectors.

        Parameters
        ----------
        distance: str
            One of ('cosine', 'euclidean', 'manhattan').
        """
        if distance == 'euclidean':
            dists = Embeddings.norm(vec1 - vec2)
        elif distance == 'cosine':
            dists = Embeddings.cosine_distance(vec1, vec2)
        elif distance == 'manhattan':
            dists = Embeddings.manhattan_distance(vec1, vec2)
        # Let arrays have numpy dtypes, but scalars will just be floats.
        return dists if isinstance(dists, Iterable) else float(dists)

    def _distances(self, vec, distance='cosine'):
        """Find distance from an input vector to every other vector in the
        embedding matrix.

        Parameters
        ----------
        vec: np.array
            Vector for the input word.
        distance: str
            Specifies what distance metric to use for calculations.
            One of ('euclidean', 'manhattan', 'cosine'). In a high dimensional
            space, cosine is often a good choice.

        Returns
        -------
        np.array: The i'th value corresponds to the distance to word i in the
            vocabulary.
        """
        return self.distance(self.mat, vec, distance=distance)

    def nearest_neighbors(self, word, n=5, distance='cosine', digits=3):
        """Find the most similar words to a given word. This wrapper
        allows the user to pass in a word. To pass in a vector, use
        `_nearest_neighbors`.

        Parameters
        ----------
        word: str
            A word that must be in the vocabulary.
        n: int
            Number of neighbors to return.
        distance: str
            Distance method to use when computing nearest neighbors. One of
            ('euclidean', 'manhattan', 'cosine').
        digits: int
            Digits to round output distances to.

        Returns
        -------
        dict[str, float]: Dictionary mapping word to distance.
        """
        # Error handling for words not in vocab.
        if word not in self:
            return None
        return self._nearest_neighbors(self.vec(word), n, distance, digits)

    def _nearest_neighbors(self, vec, n=5, distance='cosine', digits=3,
                           skip_first=True):
        """Find the most similar words to a given word's vector.
        This is the internal function behind `nearest_neighbors`, so you pass
        in a vector instead of a word.

        Parameters
        ----------
        vec: np.array
        n: int
        distance: str
            One of ('cosine', 'euclidean', 'manhattan').
        digits: int
        skip_first: bool
            If True, the nearest result will be sliced off (this is desirable
            when searching for a word's nearest neighbors, where we don't want
            to return the word itself). When finding analogies or performing
            embedding arithmetic, however, we likely don't want to slice off
            the first result.

        Returns
        -------
        dict[str, float]: Dictionary mapping word to distance.
        """
        dists = self._distances(vec, distance)
        idx = np.argsort(dists)[slice(skip_first, skip_first+n)]
        # First convert to float, otherwise we get np.float32 or np.float64
        # scalars which can cause annoying bugs in APIs or dash apps.
        return {self.i2w[i]: round(float(dists[i]), digits) for i in idx}

    def analogy(self, a, b, c, n=5, **kwargs):
        """Fill in the analogy: A is to B as C is to ___. Note that we always
        treat A and B as valid candidates to fill in the blank. C is
        only considered as a candidate in the trivial case where A=B, in which
        case C should be the first choice.

        Parameters
        ----------
        a: str
            First word in analogy.
        b: str
            Second word in analogy.
        c: str
            Third word in analogy.
        n: int
            Number of candidates to return. Note that we specify this
            separately from kwargs since we need to alter its value before
            passing it to `_nearest_neighbors`. This will allow us to remove
            the word c as a candidate if it is returned.
        kwargs: distance (str), digits (int)
            See _nearest_neighbors for details.

        Returns
        -------
        list[str]: Best candidates to complete the analogy in descending order
            of likelihood.
        """
        # If any words missing from vocab, arithmetic w/ None throws error.
        try:
            vec = self.vec(b) - self.vec(a) + self.vec(c)
        except TypeError:
            return None

        # Except for trivial edge case, return 1 extra value in case neighbors
        # includes c, which will be removed in these situations.
        a, b, c = a.lower(), b.lower(), c.lower()
        trivial = (a == b)
        neighbors = self._nearest_neighbors(vec, n=n+1-trivial,
                                            skip_first=False, **kwargs)
        if not trivial and c in neighbors:
            neighbors.pop(c)

        # Relies on dicts being ordered in python >= 3.6.
        return list(neighbors)[:n]

    def cbow(self, *args):
        """Wrapper to `_cbow` that allows us to pass in strings instead of
        vectors. Computes bag of words vector by averaging vectors for all
        input words.

        Parameters
        ----------
        args: str
            Multiple words to average over.

        Returns
        -------
        np.array: Average of all input vectors. This will have the same
            embedding dimension as each input.
        """
        vecs = [arg for arg in map(self.vec, args) if arg is not None]
        if vecs:
            return self._cbow(*vecs)

    def _cbow(self, *args):
        """Internal helper for `cbow` method that lets us pass in vectors
        instead of words.

        Parameters
        ----------
        args: np.array
            Word vectors to average.

        Returns
        -------
        np.array: Average of all input vectors. This will have the same
            embedding dimension as each input.
        """
        return np.mean(args, axis=0)

    def cbow_neighbors(self, *args, n=5, exclude_args=True, **kwargs):
        """Wrapper to `cbow` method. This lets us pass in words, compute their
        average embedding, then return the words nearest this embedding. The
        input words are not considered to be candidates for neighbors (e.g. if
        you input the words 'happy' and 'cheerful', the neighbors returned
        will not include those words even if they are the closest to the mean
        embedding) unless you set exclude_args=False. The idea here is to
        find additional words that may be similar to the group you've passed
        in.

        Parameters
        ----------
        args: str
            Input words to average over.
        n: int
            Number of neighbors to return.
        kwargs: distance (str), digits (int)
            See _nearest_neighbors() for details.

        Returns
        -------
        dict[str, float]: Dictionary mapping word to distance from the average
            of the input words' vectors.
        """
        vec_avg = self.cbow(*args)
        if vec_avg is None:
            return
        w2dist = self._nearest_neighbors(vec_avg, n=len(args)+n,
                                         skip_first=False, **kwargs)

        # Lowercase to help remove duplicates.
        args = set(arg.lower() for arg in args)
        return {word: w2dist[word] for word in
                [w for w in w2dist if not exclude_args or w not in args][:n]}

    @valuecheck
    def matching_keys(self, *terms,
                      mode:('standard', 'regex', 'ninja')='standard'):
        """Find keys (usually URLs, but could be used on words) containing
        a given term/prefix/regex. This helps us do things like create theme
        vectors (e.g. use this to find all sites related to "games" or
        "gaming", then proceed to average their embeddings and potentially
        expand the group even more by finding nearest neighbors).

        Parameters
        ----------
        terms: str
            One or more strings to search for. This can be a word ('gaming'),
            a partial word ('gam'), a regex ('^gam.*').
        mode: str
            standard - Performs exact string matching.
            regex - Allows passing in regular expressions like '^sports'.
            ninja - Use wordninja to remove spurious matches that violate
                likely word boundaries. Mostly useful for things like URLs.
                E.g. if you search for 'math' with mode='standard', the URL
                'mathiasmiller.com' would match. With mode='ninja', it would
                not because this URL seems to refer to a person, not the word
                'math'.

        Returns
        -------
        list[str]: Keys (words, URLs, etc.) matching any of the given terms.
        """
        if mode == 'regex':
            match_fn = re.search
        else:
            def match_fn(term, key):
                return term in key

            # This is a bit slow so perform it as a second step after first
            # filtering for URLs that contain the term at all.
            if mode == 'ninja':
                def post_match_fn(term, key):
                    return term in wn.split(key)

        keys = [key for key in self if any(match_fn(t, key) for t in terms)]
        if mode == 'ninja':
            keys = [k for k in keys if any(post_match_fn(t, k)
                                           for t in terms)]
        return keys

    def compare_distances(self, key, distance='cosine', as_df=True,
                          sort_df=True, **vectors):
        """Compare how far a word/domain is from one or more vectors. Intended
        for use with cbow results: e.g. checking if a word is closer to a
        'liberal' vector or a 'conservative' vector.

        Parameters
        ----------
        key: str
            Word/domain to analyze.
        distance: str
            One of ('cosine', 'euclidean', 'manhattan'). Determines distance
            method to use.
        as_df: bool
            If True, output is a dataframe. If False, it's a dict mapping
            keys (strings) to distances (floats).
        sort_df: bool
            If True and as_df is True, the output df will be sorted by
            distance from closest to furthest (relatively speaking - all
            results will be relatively close).
        vectors: np.array
            One or more vectors to compare the input key's vector to. These
            are kwargs rather than args because we need names for the
            resulting df to show which distance corresponds to which vector.

        Returns
        -------
        pd.DataFrame or dict[str, float]: Type depends on value of `as_df`.
        """
        return self._compare_distances(self.vec(key), distance, as_df,
                                       sort_df, **vectors)

    def _compare_distances(self, src_vec, distance='cosine', as_df=True,
                           sort_df=True, **vectors):
        """Internal version of `compare_distances` that accepts a vector
        rather than a word.

        Parameters
        ----------
        src_vec: np.array
            Embedding for the word/domain to analyze.
        distance: str
            One of ('cosine', 'euclidean', 'manhattan'). Determines distance
            method to use.
        as_df: bool
            If True, output is a dataframe. If False, it's a dict mapping
            keys (strings) to distances (floats).
        sort_df: bool
            If True and as_df is True, the output df will be sorted by
            distance from closest to furthest (relatively speaking - all
            results will be relatively close).
        vectors: np.array
            One or more vectors to compare the input key's vector to. These
            are kwargs rather than args because we need names for the
            resulting df to show which distance corresponds to which vector.

        Returns
        -------
        pd.DataFrame or dict[str, float]: Type depends on value of `as_df`.
        """
        d2dist = {k: self.distance(src_vec, v, distance)
                  for k, v in vectors.items()}
        if as_df:
            d2dist = pd.DataFrame.from_dict(d2dist, orient='index',
                                            columns=['score'])
            if sort_df: d2dist = d2dist.sort_values('score', ascending=True)
        return d2dist

    @valuecheck
    def semantic_vector(
            self, *queries, n=25,
            include_queries:('always', 'never', 'auto')='always',
            mode='standard', google_missing=False
    ):
        """Mostly for domains rather than words: create a vector matching the
        "semantic theme" of 1 or more input queries (usually several).
        For example, to create a "movie" theme, you could pass in 'imdb.com',
        'rottentomatoes.com', and 'letterboxd.com'.

        Parameters
        ----------
        queries: str
            These will often be keys in the Embeddings object (i.e. domains).
            However, we also support words (e.g. "movie"), regular expressions
            (e.g. "^movie*"), or even phrases (e.g. "scary movies"). Some
            combinations of the above are supported: the only limitation as of
            2/12/21 is you must choose a str matching mode (see
            `self.matching_keys`). So you can pass in a mix of domains, words,
            and phrases, but you can't use string matching for some words and
            regex matching for others. Might be supported in the future if I
            encounter situations where it seems useful.
        n: int
            Number of neighbors to find in `cbow_neighbors` method. Note that
            this won't necessarily be the final number of keys returned -
            `include_queries` will affect that too.
        include_queries: str
            Determines whether URLs retrieved from the initial step of query
            string matching should be included in results. 'never' is useful
            if you specifically want URLs that DON'T contain the queries
            (e.g. music-related sites without "music" in the URL), but
            I suspect 'always' may give better quality results. 'auto' will
            allow these matches to be retained but won't force them to if they
            aren't close to the final theme vector.
        mode: str
            Determines type of string matching used to find initial "seed"
            domains (see `self.matching_keys`). Options are
            ('standard', 'regex', 'ninja').
        google_missing: bool
            If True, terms that don't yield any string match results will be
            googled and any domains on the first page of results will be used
            if they're in our vocabulary.

        Returns
        -------
        tuple[np.array, list[str]] or None: First item is the "bag of words"
        array created by averaging over n neighbors of the input queries.
        Second item is a list of the neighbor names sorted by similarity in
        decreasing order.
        """
        # If a query isn't in our embeddings, we assume it's intended to be a
        # str matching query, not a full domain. If it is a full domain, this
        # is a bit inefficient but it won't affect the final list of queries.
        #  Don't use a list comprehension because we need to warn when no
        # matches are found since we remove missing queries before calling
        # `cbow_neighbors` (i.e. no built-in warning). This means our final
        # list won't have missing queries, which is nice.
        all_queries = []
        for q in queries:
            if q in self:
                matches = [q]
            else:
                matches = self.matching_keys(q, mode=mode)
                if not matches and google_missing:
                    matches = domains_from_google_search(q)
                if not matches:
                    warnings.warn(f'No matches found for {q}.')
            all_queries.extend(matches)
        if not all_queries:
            warnings.warn('Queries yielded zero matches.')
            # Return None (not vector) for consistency with other methods.
            return None, []

        # Keep as dict for now for O(1) lookup time in case we use huge n.
        matches = self.cbow_neighbors(*all_queries, n=n,
                                      exclude_args=include_queries=='never')

        if include_queries == 'always':
            # Use dict as ordered set: remove duplicates but maintain order.
            matches = dict.fromkeys(
                list(matches) + [q for q in all_queries if q not in matches]
            )
        vec = self.cbow(*matches)
        return vec, list(matches)

    @staticmethod
    def norm(vec):
        """Compute L2 norm of a vector. Euclidean distance between two vectors
        can be found by the operation norm(vec1 - vec2).

        Parameters
        ----------
        vec: np.array
            Input vector.

        Returns
        -------
        float: L2 norm of input vector.
        """
        return np.sqrt(np.sum(vec ** 2, axis=-1))

    @staticmethod
    def manhattan_distance(vec1, vec2):
        """Compute L1 distance between two vectors.

        Parameters
        ----------
        vec1: np.array
        vec2: np.array

        Returns
        -------
        float or np.array: Manhattan distance between vec1 and vec2. If two
            vectors are passed in, the output will be a single number. When
            computing distances between a vector and a matrix, the output
            will be a vector (np.array).
        """
        return np.sum(abs(vec1 - vec2), axis=-1)

    @staticmethod
    def cosine_distance(vec1, vec2):
        """Compute cosine distance between two vectors.

        Parameters
        ----------
        vec1: np.array
        vec2: np.array

        Returns
        -------
        float or np.array: Cosine distance between vec1 and vec2. If two
            vectors are passed in, the output will be a single number. When
            computing distances between a vector and a matrix, the output
            will be a vector (np.array).
        """
        return 1 - (np.sum(vec1 * vec2, axis=-1) /
                    (Embeddings.norm(vec1) * Embeddings.norm(vec2)))

    @dispatch(str)
    def __getitem__(self, key):
        """When indexing with a string, this acts as a word->index method.

        Examples
        --------
        >>> emb['the']
        1
        """
        return self.w2i[key.lower()]

    @dispatch((int, slice))
    def __getitem__(self, i):
        """When indexing with an integer, this acts as an index->word method.

        Examples
        --------
        >>> emb[1]
        'the'

        >>> emb[:3]
        ['a', the', 'is']
        """
        return self.i2w[i]

    @dispatch(Iterable)
    def __getitem__(self, keys):
        """Allows indexing in with a list of keys/indices. You can pass in a
        mix of strings and integers though I can't imagine why that would be
        necessary.

        Examples
        --------
        >>> emb[[1, 100, 7]]
        ['the', 'frog', 'dog']

        >>> emb[['the', 'dog', 'frog']]
        [1, 7, 100]
        """
        return [self[key] for key in keys]

    def get(self, key, default=None):
        """Returns None if word is not present just like dict.get."""
        try:
            return self.w2i[key.lower()]
        except KeyError:
            warnings.warn(f'{key} not in Embeddings.')
            return default

    def __len__(self):
        return self.n_embeddings

    def __contains__(self, word):
        return word.lower() in self.w2i

    def __iter__(self):
        """Yields words in vocabulary in insertion order (may differ from
        index order).
        """
        yield from self.w2i.keys()

    def __eq__(self, obj):
        if not isinstance(obj, Embeddings):
            return False

        ignore = {'pca'}
        for k, v in vars(obj).items():
            if k in ignore: continue
            v_self = getattr(self, k)
            if isinstance(v, np.ndarray):
                if not np.allclose(v, v_self): return False
            elif v != v_self: return False
        return True

    def __repr__(self):
        return f'Embeddings(len={len(self)}, dim={self.dim})'


# Cell
def back_translate(text, to, from_lang='en'):
    """Translate a piece of text into another language, then back to English
    for data augmentation purposes. This is rate limited but we now have a
    pure ML version in the form of BacktranslateTransform.

    Parameters
    ----------
    text: str
        Text to back translate.
    to: str
        Language to translate to before translating back to English.
    from_lang: str
        Language of input text (usually 'en' for English).

    Returns
    -------
    str: Same language and basically the same content as the original text,
        but usually with slightly altered grammar, sentence structure, and/or
        vocabulary.
    """
    return str(
        TextBlob(text)\
        .translate(to=to)\
        .translate(from_lang=to, to=from_lang)
    )


# Cell
def postprocess_embeddings(emb, d=None):
    """Implements the algorithm from the paper:

    All-But-The-Top: Simple and Effective Post-Processing
    for Word Representations (https://arxiv.org/pdf/1702.01417.pdf)

    There are three steps:
    1. Compute the mean embedding and subtract this from the
    original embedding matrix.
    2. Perform PCA and extract the top d components.
    3. Eliminate the principal components from the mean-adjusted
    embeddings.

    Parameters
    ----------
    emb: np.array
        Embedding matrix of size (vocab_size, embedding_length).
    d: int
        Number of components to use in PCA. Defaults to
        embedding_length/100 as recommended by the paper.
    """
    d = d or emb.shape[1] // 100
    emb_adj = emb - emb.mean(0)
    u = PCA(d).fit(emb_adj).components_
    return emb_adj - emb@u.T@u


# Cell
def compress_embeddings(emb, new_dim, d=None):
    """Reduce embedding dimension as described in the paper:

    Simple and Effective Dimensionality Reduction for Word Embeddings
    (https://lld-workshop.github.io/2017/papers/LLD_2017_paper_34.pdf)

    Parameters
    ----------
    emb: np.array
        Embedding matrix of size (vocab_size, embedding_length).
    d: int
        Number of components to use in the post-processing
        method described here: https://arxiv.org/pdf/1702.01417.pdf
        Defaults to embedding_length/100 as recommended by the paper.

    Returns
    -------
    np.array: Compressed embedding matrix of shape (vocab_size, new_dim).
    """
    emb = postprocess_embeddings(emb, d)
    emb = PCA(new_dim).fit_transform(emb)
    return postprocess_embeddings(emb, d)


# Cell
@auto_repr
class ParaphraseTransform:
    """Text transform that paraphrases input text as a method of data
    augmentation. This is rather slow so it's recommended to precompute
    samples and save them, but you could generate samples on the fly if
    desired. One further downside of that approach is you'll have a huge
    paraphrasing model on the GPU while (presumably) training another model.

    Other paraphrasing models exist on Model Hub but as of 11/14/2020, none of
    the results compared favorably to this pegasus model, at least based on
    a rough "eyeball check". While smaller and presumably faster, many of
    these appear to require processing a single example at a time which
    diminishes these gains. If you do attempt to use them, you'll likely need
    to write a new class with a preprocessing method that does something like
    the following:
    _preprocess(text) -> 'paraphrase: {text}</s>'
    I'm recording this here because many are missing documentation and it took
    me some time to discover this.
    """

    name = 'tuner007/pegasus_paraphrase'

    def __init__(self, n=1, pipe=None):
        """
        Parameters
        ----------
        n: int
            Default number of samples to generate. You can override this in
            __call__.
        pipe: transformers Text2TextGenerationPipeline or None
        """
        if pipe:
            self.pipe = pipe
            self.name = pipe.model.config._name_or_path
        else:
            self.pipe = Text2TextGenerationPipeline(
                PegasusForConditionalGeneration.from_pretrained(self.name),
                PegasusTokenizer.from_pretrained(self.name),
                device=0 if torch.cuda.is_available() else -1
            )
        self.n = n

        assert type(self.pipe).__name__ == 'Text2TextGenerationPipeline'
        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():
            warnings.warn('The pipeline passed in is not using cuda. '
                          'Did you mean to use the available GPU?')

    def _preprocess(self, text):
        """Does nothing (just want shared interface with other transforms)."""
        return text

    @add_docstring(PreTrainedModel.generate)
    def __call__(self, text, n=None, flat=True, **kwargs):
        """
        Parameters
        ----------
        text: str or Iterable[str]
            Raw text to transform.
        n: int or None
            If None, use the default self.n.
        flat: bool
            If True, return flat list of strings. If False, return list of
            nested lists where list i contains n augmentations of input i.
        kwargs: any
            Additional kwargs are passed to the model's text generation
            method. Its docstring is included below for convenience.

        Returns
        -------
        list: either a list with n strings per input string, or a list of
        lists, each of length n, if flat=False.
        """
        n = n or self.n
        rows = [row['generated_text'] for row in
                self.pipe(text, num_return_sequences=n, **kwargs)]
        if listlike(text) and not flat:
            rows = [rows[i*n:(i+1)*n] for i in range(len(text))]
        return rows


# Cell
@auto_repr
class GenerativeTransform:
    """Text transform that truncates a piece of text and completes it using
    a text generation model for the purposes of data augmentation. We
    recommend precomputing samples and saving them for later use, but you
    could generate samples on the fly if desired. Aside from speed, this
    approach also has the drawback of having a text generation model on the
    GPU while (presumably) training another model.
    """

    name = 'text-generation'

    def __init__(self, n=1, pipe=None):
        """
        Parameters
        ----------
        n: int
            Default number of samples to generate. You can override this in
            __call__.
        pipe: Transformers TextGenerationPipeline or None
        """
        if pipe:
            self.pipe = pipe
            self.name = pipe.model.config._name_or_path
        else:
            self.pipe = pipeline(self.name,
                                 device=0 if torch.cuda.is_available()
                                 else -1)
        self.n = n

        assert type(self.pipe).__name__ == 'TextGenerationPipeline'
        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():
            warnings.warn('The pipeline passed in is not using cuda. '
                          'Did you mean to use the available GPU?')

    def _preprocess(self, text, drop=None, drop_pct=None, rand_low=None,
                    rand_high=None, min_keep=3, return_tuple=False):
        """Truncate text so we can generate the ending.

        Parameters
        ----------
        text: str or Iterable[str]
        drop: None or int
            If provided, specifies the number of words to drop. We use a
            simple "split on spaces" strategy since it's fast and simple.
            Drop strategies occur in the signature in order of priority, so
            if this is non-None it will override any values passed in for
            drop_pct or rand_low/rand_high.
        drop_pct: float or None
            If provided, this should be a value between 0.0 and 1.0 specifying
            the proportion of words to drop.
        rand_low: int or None
            If provided, specifies the minimum number of words to drop. The
            max will be set by rand_high (which must also be provided with
            rand_low). A random integer will be selected for each row of text.
        rand_high: int or None
            See rand_low: helps define bounds when randomly truncating rows
            of text.
        min_keep: int
            The minimum number of words to keep. Sequences of this length or
            shorter will therefore remain un-transformed. You could set this
            to zero to enforce no minimum.
        return_tuple: bool
            If True, return a tuple where the first item is the truncated
            text and the second item is the number of words masked. This is
            rarely needed but it might be helpful if you want to use this for
            some sort of self-supervised pre-training task.
        """
        if listlike(text):
            return [self._preprocess(row, drop, drop_pct, rand_low, rand_high,
                                     min_keep, return_tuple) for row in text]

        tokens = text.split()
        if len(tokens) <= min_keep:
            n_drop = 0
        else:
            # Default is to truncate the last 20% of the sequence.
            if drop:
                n_drop = drop
            elif drop_pct:
                n_drop = int(drop_pct * len(tokens))
            elif rand_low is not None and rand_high is not None:
                n_drop = np.random.randint(rand_low, rand_high)
            else:
                n_drop = int(np.ceil(.2 * len(tokens)))
            n_drop = np.clip(n_drop, 0, len(tokens) - min_keep)
            tokens = tokens[:-n_drop]
        truncated = ' '.join(tokens)
        return (truncated, n_drop) if return_tuple else truncated

    @add_docstring(PreTrainedModel.generate)
    def __call__(self, text, n=None, flat=True, min_length=2, max_length=7,
                 drop=None, drop_pct=None, rand_low=None, rand_high=None,
                 min_keep=3, **generate_kwargs):
        """
        Parameters
        ----------
        text: str or Iterable[str]
        n: int or None
            Number of samples to generate for each input. Defaults to self.n
            if None.
        flat: bool
            If True, return flat list of strings. If False, return list of
            nested lists where list i contains n augmentations of input i.
        min_length: int
            Min number of tokens to generate.
        max_length: int
            Max number of tokens to generate. You could set this equal to
            min_length to enforce a constant number.
        drop: None or int
            If provided, specifies the number of words to drop. We use a
            simple "split on spaces" strategy since it's fast and simple.
            Drop strategies occur in the signature in order of priority, so
            if this is non-None it will override any values passed in for
            drop_pct or rand_low/rand_high.
        drop_pct: float or None
            If provided, this should be a value between 0.0 and 1.0 specifying
            the proportion of words to drop.
        rand_low: int or None
            If provided, specifies the minimum number of words to drop. The
            max will be set by rand_high (which must also be provided with
            rand_low). A random integer will be selected for each row of text.
        rand_high: int or None
            See rand_low: helps define bounds when randomly truncating rows
            of text.
        min_keep: int
            The minimum number of words to keep. Sequences of this length or
            shorter will therefore remain un-transformed. You could set this
            to zero to enforce no minimum.
        generate_kwargs: any
            Forwarded to model's `generate` method. For convenience, its
            docstring is provided below.

        Returns
        -------
        list: either a list with n strings per input string, or a list of
        lists, each of length n, if flat=False.
        """
        n = n or self.n
        if listlike(text):
            res = [self(row, n, flat=flat, min_length=min_length,
                        max_length=max_length, drop=drop, drop_pct=drop_pct,
                        rand_low=rand_low, rand_high=rand_high,
                        min_keep=min_keep, **generate_kwargs) for row in text]
            return flatten(res) if flat else res

        # `generate` counts current length as part of min_length.
        text = self._preprocess(text, drop, drop_pct, rand_low=rand_low,
                                rand_high=rand_high, min_keep=min_keep)
        n_curr = len(self.pipe.tokenizer.tokenize(text))
        res = self.pipe(text, min_length=n_curr + min_length,
                        max_length=n_curr + max_length,
                        num_return_sequences=n, **generate_kwargs)
        return [row['generated_text'] for row in res]


# Cell
@auto_repr
class FillMaskTransform:
    """Text transform that masks one or more words in a piece of text and
    fills them using RoBERTa for the purposes of data augmentation. We
    recommend precomputing samples and saving them for later use, but you
    could generate samples on the fly if desired. In addition to being slow,
    that approach also entails having a mask filling model on the GPU while
    (presumably) training another model.
    """

    name = 'fill-mask'
    MASK = '<mask>'

    def __init__(self, n=1, max_n=None, pipe=None):
        """
        Parameters
        ----------
        n: int
            n is intentionally bigger than the default n in __call__. This is
            the number of candidates generated, so if we use strategy='random'
            it makes sense for this to be larger.
        max_n: int or None
            Used as topk attribute when sampling generated candidates. This
            must be >=n. Increasing this should not slow down generation. When
            using strategy='random', you probably want this to be strictly >n.
            Defaults to n+2.
        pipe: transformers FillMaskPipeline
            We let users pass in an existing pipeline since instantiation can
            be slow.
        """
        if pipe:
            self.pipe = pipe
            self.name = pipe.model.config._name_or_path
        else:
            self.pipe = pipeline(self.name,
                                 device=0 if torch.cuda.is_available()
                                 else -1)
        # Set n before max_n.
        self.n = n
        self.max_n = max_n or n+2

        assert type(self.pipe).__name__ == 'FillMaskPipeline'
        if 'cuda' not in str(self.pipe.device) and torch.cuda.is_available():
            warnings.warn('The pipeline passed in is not using cuda. '
                          'Did you mean to use the available GPU?')

    def _preprocess(self, text, min_keep=3, errors='raise'):
        """Randomly mask one word from an input piece of text to prepare it
        for RoBERTa to fill. Notice that even if the user chooses to mask
        multiple words, each call to `_preprocess` only masks one since the
        model can only fill one at a time.

        Parameters
        ----------
        text: str or Iterable[str]
            One or more pieces of text to process.
        min_keep: int
            Minimum number of words to keep after truncating each piece of
            text.
        errors: str
            If 'warn', we show a warning when min_keep is violated but allow
            masking to take place. Any other value will result in an error
            being raised.
        """
        if listlike(text):
            return [self._preprocess(row, min_keep, errors) for row in text]

        tokens = text.split()
        if len(tokens) < min_keep + 1:
            msg = (f'Text "{text[:25]}..." is too short to mask while '
                   f'enforcing min_keep={min_keep}.')
            # Err on side of caution: typos raise error too.
            if errors == 'warn':
                warnings.warn(msg)
            else:
                raise ValueError(msg)

        idx = np.random.choice(range(len(tokens)))
        return ' '.join(self.MASK if i == idx else t
                        for i, t in enumerate(tokens))

    @add_docstring(PreTrainedModel.generate)
    def __call__(self, text, n=None, flat=True, n_mask=1, min_keep=3,
                 return_all=False, errors='raise', strategy='best', **kwargs):
        """
        Parameters
        ----------
        text: str or Iterable[str]
        n: int or None
            Number of variations to return per piece of input text. If -1,
            return all generated examples for the given mask count.
            This can become very large when n_mask is large.
            Example: if self.max_n=3, n=-1, and n_mask=4, we first mask once
            and generate 3 samples. Then we mask each of those 3 and generate
            a total of 9 samples, then 27, then finally 81 which is what will
            be returned. The intermediate samples can be returned by
            specifying `return_all=True`.
        flat: bool
            If True, return flat list of strings. If False, return list of
            nested lists where list i contains n augmented versions of input
            i.
        n_mask: int
            Number of words to mask. Because the model can only fill 1 masked
            word at a time, `n_mask` forward passes will be performed.
        min_keep: int
            Minimum number of words to keep (presumably, you wouldn't want to
            mask every word in an input sentence since that would strip it of
            all existing meaning). This can be strictly enforced or not,
            depending on your choice of `errors`.
        return_all: bool
            If True, return all intermediate generated samples rather than
            just the final samples (e.g. if n_mask is 3, we first have to
            generate samples with 1 masked word, then samples with 2 masked
            words, and finally samples with 3 masked words. This is because
            RoBERTa can only fill one masked word at a time.) See the
            explanation of the `n` parameter for an example.
        errors: str
            One of ('warn', 'raise'). 'raise' will raise a ValueError if we're
            about to violate `min_keep`. 'warn' will only show a warning if
            this happens but will not strictly prevent it from occurring.
        strategy: str
            One of ('random', 'best'). The model will generate self.max_n
            samples and if n < self.max_n, this means we need some way of
            selecting which samples to keep. 'random' selects randomly without
            replacement, while 'best' chooses the n most likely generations.
            Note: when n_mask > 1, you should probably use strategy='random'
            if you want relatively diverse results. If 'best', the benefit of
            additional iterations is diminished because we are likely to end
            up with very similar (or even identical) results.
        kwargs: any
            Forwarded to model's `generate` method. Its docstring is provided
            below for convenience.

        Returns
        -------
        list: either a list with n strings per input string, or a list of
        lists, each of length n, if flat=False. This is slightly different if
        return_all=True (see its description for details).
        """
        # Make sure we generate adequate number of sequences. Model topk must
        # be >= our desired n.
        n = n or self.n
        if n > self.max_n:
            self.max_n = n

        # Each item will be a list of strings. Each string in res[i]
        # will have i words changed. If text is a sequence of strings, we must
        # handle each one separately because each is passed through pipeline
        # repeatedly.
        if listlike(text):
            res = [self(row, n=n, flat=flat, n_mask=n_mask, min_keep=min_keep,
                        return_all=return_all, errors=errors,
                        strategy=strategy, **kwargs) for row in text]
            return flatten(res) if flat else res

        res = [[text]]
        for i in range(n_mask):
            seqs = self.pipe(self._preprocess(res[-1], min_keep=min_keep,
                                              errors=errors))
            # Transformers returns either list of dicts or list of list of
            # dicts depending on whether input list has 1 item or multiple.
            if isinstance(seqs[0], list):
                seqs = [seq for group in seqs for seq in group]
            text = [seq['sequence'].replace('<s>', '').replace('</s>', '')
                    for seq in seqs]

            # Keep all generated samples when n is -1.
            if n != -1:
                if strategy == 'random':
                    text = np.random.choice(text, n, replace=False)
                elif strategy == 'best':
                    text = text[:n]
                else:
                    raise ValueError('strategy should be "random" or "best".')
            res.append(text)
        if not return_all: res = res[n_mask]
        return flatten(res) if flat else res

    @property
    def max_n(self):
        return self.pipe.topk

    @max_n.setter
    def max_n(self, max_n):
        """Need to ensure the model generates enough options to return the
        desired number of samples."""
        if not isinstance(max_n, int):
            raise TypeError('max_n must be an integer.')
        if max_n < self.n:
            raise ValueError(f'max_n must be >= self.n (currently {self.n}.')
        self.pipe.topk = max_n


# Cell
NLP_TRANSFORMS = {
    'fillmask': FillMaskTransform,
    'paraphrase': ParaphraseTransform,
    'generative': GenerativeTransform
}


# Cell
class BacktranslateTransform:
    """Augment/perturb text inputs by translating them to a different language
    and then back to English (this process can be repeated as many times as
    you want by specifying multiple target languages). As of 2/10/21, this is
    excluded from incendio's NLP_TRANSFORMS variable since its interface is a
    little different from the other transforms: it has 2 pipelines, not 1, so
    has variables `names` and `pipes` (both lists) instead of `name` and
    `pipe`. It also has no `_preprocess` method.
    """

    names = ['Helsinki-NLP/opus-mt-en-ROMANCE',
             'Helsinki-NLP/opus-mt-ROMANCE-en']

    language_codes = {
        'es': 'spanish',
        'fr': 'french',
        'it': 'italian',
        'pt': 'portuguese',
        'pt_br': 'portuguese (brazil)',
        'ro': 'romanian',
        'ca': 'catalan',
        'gl': 'galician',
        'pt_BR': 'portuguese (brazil?)',
        'la': 'latin',
        'wa': 'walloon',
        'fur': 'friulian (?)',
        'oc': 'occitan',
        'fr_CA': 'french (canada)',
        'sc': 'sardianian',
        'es_ES': 'spanish',
        'es_MX': 'spanish (mexico)',
        'es_AR': 'spanish (argentina)',
        'es_PR': 'spanish (puerto rico)',
        'es_UY': 'spanish (uruguay)',
        'es_CL': 'spanish (chile)',
        'es_CO': 'spanish (colombia)',
        'es_CR': 'spanish (croatia)',
        'es_GT': 'spanish (guatemala)',
        'es_HN': 'spanish (honduras)',
        'es_NI': 'spanish (nicaragua)',
        'es_PA': 'spanish (panama)',
        'es_PE': 'spanish (peru)',
        'es_VE': 'spanish (venezuela)',
        'es_DO': 'spanish (dominican republic)',
        'es_EC': 'spanish (ecuador)',
        'es_SV': 'spanish (el salvador)',
        'an': 'aragonese',
        'pt_PT': 'portuguese (portugal)',
        'frp': 'franco provencal',
        'lad': 'ladino',
        'vec': 'venetian',
        'fr_FR': 'france (france)',
        'co': 'corsican',
        'it_IT': 'italian (italy)',
        'lld': 'ladin',
        'lij': 'ligurian',
        'lmo': 'lombard',
        'nap': 'neapolitan',
        'rm': 'rhaetian (?)',
        'scn': 'sicilian',
        'mwl': 'mirandese'
    }

    def __init__(self, to_langs, pipes=()):
        """
        Parameters
        ----------
        to_langs: Iterable[str]
            One or more language codes to use for backtranslation (see
            self.language_codes for all options). They will be applied in
            order: for instance, passing in ['es', 'fr'] will translate input
            from
            english -> spanish -> english -> french -> english. You can
            override these later in specific calls but the value(s) you
            provide here will be defaults.
        pipes: Iterable[Pipeline]
            Huggingface pipelines, the first of which translates English to
            Romance languages and the second of which does the reverse. It's
            usually easiest to let the Transform create these for you, but if
            you already have them passing them in will be faster.
        """
        if not pipes:
            pipes = [TranslationPipeline(
                        model=AutoModelForSeq2SeqLM.from_pretrained(name),
                        tokenizer=AutoTokenizer.from_pretrained(name),
                        device=1 - torch.cuda.is_available()
                     ) for name in names]
        self.pipes = pipes
        self.to_langs = tolist(to_langs)

    def __call__(self, text, intermediate=False, flat=True, to_langs=(),
                 **kwargs):
        """
        Parameters
        ----------
        text: str or Iterable[str]
            The input pieces of text to translate.
        intermediate: bool
            If True, return all intermediate backtranslations if more than one
            target language is provided. Otherwise, only return the final
            backtranslation.
        flat: bool
            If True, return a flat list of strings. If False, return a list of
            lists where item i contains all the intermediate backtranslations
            (n languages in `to_langs` will generate n backtranslations). Note
            that if intermediate=False, results will always be flat.
        to_langs: Iterable[str]
            One or more language codes to use for backtranslation. They will
            be applied in order: for instance, passing in ['es', 'fr'] will
            translate input from
            english -> spanish -> english -> french -> english. If not
            specified, this defaults to self.to_langs.
        kwargs: any
            Ignored. Just provided for consistency with other transforms.

        Returns
        -------
        list[str] or list[list[str]]: Default is list of strings where item
        i of output corresponds to item i of input. If intermediate=True and
        flat=False, we get a list of lists where each nested list contains
        n backtranslations of input i.
        """
        text = tolist(text)
        to_langs = tolist(to_langs) or self.to_langs
        assert not set(to_langs) - set(self.language_codes), \
            'to_langs codes should all be present in self.language_codes.'

        steps = []
        for lang in to_langs:
            text = [f'>>{lang}<< {t}' for t in text]
            text = [row['translation_text'] for row in self.pipes[0](text)]
            text = [row['translation_text'] for row in self.pipes[1](text)]
            steps.append(text)
        if intermediate:
            steps = zip(*steps)
            return flatten(steps) if flat else lmap(list, *steps)
        else:
            return text

    def __repr__(self):
        lang_str = ", ".join(repr(lang) for lang in self.to_langs)
        return f'{func_name(self)}(to_langs=[{lang_str}])'


# Cell
@immutify_defaults
def augment_text_df(source, transform='fillmask', dest=None, n=5,
                    text_col='text', id_cols=(), nrows=None, tfm_kwargs={},
                    call_kwargs={}):
    """Create augmented versions of a dataframe of text, optionally preserving
    other columns for identification purposes. We recommend precomputing and
    saving variations of your data rather than doing this on the fly in a
    torch dataset since they can be rather space- and time-intensive.
    Augmented versions of an input row should generally be kept in the same
    training split: in order to keep the label the same, we usually want to
    make relatively limited changes to the raw text (just enough to provide a
    regularizing effect).

    Parameters
    ----------
    source: str, Path, or pd.DataFrame
        If str or Path, this is a csv containing our text data. Alternatively,
        you can pass in a dataframe itself.
    transform: str or callable
        If str, this must be one of the keys in `NLP_TRANSFORMS` from this
        same module - this will be used to create a new transform object.
        Alternatively, you can pass in a previously created object (NOT the
        class). The default is the mask filling transform as it's relatively
        quick and effective. 'paraphrase' may give better (but slower)
        results. Anecdotally, 'generative' seems to provide lower quality
        results, but perhaps by experimenting with hyperparameters it could
        be more useful.
    dest: str, Path, or None
        If str or Path, this is where the output file will be saved to
        (directories will be created as needed). If None, nothing will be
        saved and the function will merely return the output DF for you to do
        with as you wish.
    n: int
        Number of samples to generate for each raw row.
    text_col: str
        Name of column in DF containing the text to augment.
    id_cols: Iterable[str]
        Columns containing identifying information such as labels, row_ids,
        etc. These also help us map the augmented text rows to their
        corresponding raw rows.
    nrows: int or None
        Max number of rows from the source DF to generate text for. Useful for
        testing (equivalently, you could pass in df.head(nrows) and leave this
        as None).
    tfm_kwargs: dict
        Arguments to pass to `transform`'s constructor. These are ignored when
        passing in a transform object rather than a string.
    call_kwargs: dict
        Arguments to pass to the __call__ method of `transform` to affect
        the augmentation process.

    Returns
    -------
    pd.DataFrame: DF of generated text with columns `text_col` and `id_cols`.
    By default, this will have 5x the rows as your source DF, but this can
    easily be adjusted through the `nrows` parameter.
    """
    # Load data.
    if isinstance(source, (str, Path)):
        df = pd.read_csv(Path(source), usecols=[text_col] + list(id_cols),
                         nrows=nrows)
    elif isinstance(source, pd.DataFrame):
        df = source.head(nrows)
    else:
        raise TypeError('`source` must be a str/Path or pd.DataFrame.')

    # Prepare for output file if necessary.
    if isinstance(dest, (str, Path)):
        dest = Path(dest)
        os.makedirs(dest.parent, exist_ok=True)
    elif dest is not None:
        raise ValueError('`dest` must be a str/Path containing the output '
                         'file name to create, or None if you just want to '
                         'return a df.')

    # For simplicity, we stick to one transform at a time. Slow to load so at
    # least for now, let user pass in the transform itself.
    transform = NLP_TRANSFORMS[transform](n=n, **tfm_kwargs) \
        if isinstance(transform, str) else transform

    # Generate new variations of input text.
    res = transform(df[text_col].tolist(), **{**call_kwargs, 'flat': True})
    res = pd.DataFrame(res, columns=[text_col])

    # Attach identifier columns to output (e.g. we usually want to store
    # labels and or sample IDs. Most of our augmentation methods make
    # relatively minor changes to the input so all variations of 1 input
    # should remain in the same set, usually training.).
    if id_cols:
        df_id = pd.concat([df[col].repeat(res.shape[0] // df.shape[0])
                           for col in id_cols], axis=1).reset_index(drop=True)
        res = pd.concat([df_id, res], axis=1)

    # Optionally save output.
    if dest: res.to_csv(dest, index=False)
    return res