{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import gc\n",
    "import numpy as np\n",
    "import spacy\n",
    "import torch\n",
    "\n",
    "from htools import hdir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Counter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'a': 1, 'b': 1, ' ': 1, 'c': 1})\n",
      "Counter({' ': 4, 'c': 3, 'b': 2, 'e': 2, 'a': 1, ',': 1, 'd': 1})\n",
      "{'b', 'd', ' ', 'c', 'a', 'e', ','}\n"
     ]
    }
   ],
   "source": [
    "tmp1 = Counter('ab c')\n",
    "tmp2 = Counter('a, ccc bb d ee')\n",
    "print(tmp1)\n",
    "print(tmp2)\n",
    "print({k for k in set(list(tmp1.keys()) + list(tmp2.keys()))})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'b': 3, 'd': 1, ' ': 5, 'c': 4, 'a': 2, 'e': 2, ',': 1}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{k: tmp1.get(k, 0) + tmp2.get(k, 0) for k in {*tmp1.keys(), *tmp2.keys()}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scratch_add_dicts.ipynb              scratch_train.ipynb\r\n",
      "scratch_inheritance_and_mixins.ipynb scratch_vocabulary.ipynb\r\n"
     ]
    }
   ],
   "source": [
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 torch.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",
    "        if verbose:\n",
    "            print(f'Saving vocabulary to {path}.')\n",
    "        torch.save(self, path)\n",
    "    \n",
    "    def filter_tokens(self, tokens, max_words=None, min_freq=0, inplace=False, \n",
    "                     recompute=False):\n",
    "        \"\"\"Filter your vocabulary by specifying a max number of words or a min\n",
    "        frequency in the corpus. When done in place, this also sorts vocab by\n",
    "        frequency with more common words coming first (after idx_misc).\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        tokens: list[str]\n",
    "            A tokenized list of words in the corpus (must be all lowercase \n",
    "            when self.all_lower=True, such as when using GloVe vectors). There \n",
    "            is no need to hold out test data here since we are not using \n",
    "            labels.\n",
    "        max_words: int (optional)\n",
    "            Provides an upper threshold for the number of words in the\n",
    "            vocabulary. If no value is passed in, no maximum limit will be\n",
    "            enforced.\n",
    "        min_freq: int (optional)\n",
    "            Provides a lower threshold for the number of times a word must \n",
    "            appear in the corpus to remain in the vocabulary. If no value is\n",
    "            passed in, no minimum limit will be enforced.\n",
    "            \n",
    "            Note that we can specify values for both max_words and min_freq\n",
    "            if desired. If no values are passed in for either, no pruning of\n",
    "            the vocabulary will be performed.\n",
    "        inplace: bool\n",
    "            If True, will change the object's attributes \n",
    "            (w2idx, w2vec, and i2w) to reflect the newly filtered vocabulary.\n",
    "            If False, will not change the object, but will simply compute word\n",
    "            counts and return what the new w2idx would be. This can be helpful\n",
    "            for experimentation, as we may want to try out multiple values of\n",
    "            min_freq to decide how many words to keep. After the first call,\n",
    "            the attribute corpus_counts can also be examined to help determine\n",
    "            the desired vocab size.\n",
    "        recompute: bool\n",
    "            If True, will calculate word counts from the given tokens. If \n",
    "            False (the default), this will use existing counts if there are \n",
    "            any. \n",
    "            \n",
    "            The idea is that if we call this method, then realize we want\n",
    "            to change the corpus, we should calculate new word counts. \n",
    "            However, if we are simply calling this method multiple times on \n",
    "            the same corpus while deciding on the exact vocab size we want,\n",
    "            we should not recompute the word counts.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        dict or None: When called inplace, nothing is returned. When not \n",
    "        inplace, \n",
    "        \"\"\"\n",
    "        misc_len = len(self.idx_misc)\n",
    "        if recompute or not self.corpus_counts:\n",
    "            self.corpus_counts = Counter(tokens)\n",
    "        filtered = {word: i for i, (word, freq) \n",
    "                    in enumerate(self.corpus_counts.most_common(max_words),\n",
    "                                 misc_len) \n",
    "                    if freq >= min_freq}\n",
    "        filtered = {**self.idx_misc, **filtered}\n",
    "        \n",
    "        if inplace:\n",
    "            # Relies on python3.7 dicts retaining insertion order.\n",
    "            self.i2w = list(filtered.keys())\n",
    "            self.w2idx = filtered\n",
    "            self.w2vec = {word: self.vector(word) for word in filtered}\n",
    "        else:\n",
    "            return filtered\n",
    "        \n",
    "    def build_embedding_matrix(self, inplace=False):\n",
    "        \"\"\"Create a 2D numpy array of embedding vectors where row[i] \n",
    "        corresponds to word i in the vocabulary. This can be used to \n",
    "        initialize weights in the model's embedding layer.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        inplace: bool\n",
    "            If True, will store the output in the object's embedding_matrix\n",
    "            attribute. If False (default behavior), will simply return the \n",
    "            matrix without storing it as part of the object. In the \n",
    "            recommended case where inplace==False, we can store the output\n",
    "            in another variable which we can use to initialize the weights in\n",
    "            Torch, then delete the object and free up memory using \n",
    "            gc.collect().\n",
    "        \"\"\"\n",
    "        emb = np.zeros((len(self), self.dim))\n",
    "        for i, word in enumerate(self):\n",
    "            emb[i] = self.vector(word)\n",
    "            \n",
    "        if inplace:\n",
    "            self.embedding_matrix = emb\n",
    "        else:\n",
    "            return emb\n",
    "    \n",
    "    def idx(self, word):\n",
    "        \"\"\"This will map a word (str) to its index (int) in the vocabulary. \n",
    "        If a string is passed in and the word is not present, the index\n",
    "        corresponding to the <UNK> token is returned.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        word: str\n",
    "            A word that needs to be mapped to an integer index.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        int: The index of the given word in the vocabulary.\n",
    "            \n",
    "        Examples\n",
    "        ---------\n",
    "        >>> vocab.idx('the')\n",
    "        2\n",
    "        \"\"\"\n",
    "        if self.all_lower and word not in self.idx_misc:\n",
    "            word = word.lower()\n",
    "        return self.w2idx.get(word, self.w2idx['<UNK>'])\n",
    "    \n",
    "    def vector(self, word):\n",
    "        \"\"\"This maps a word to its corresponding embedding vector. If not\n",
    "        contained in the vocab, a vector of zeros will be returned.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        word: str\n",
    "            A word that needs to be mapped to a vector.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        np.array\n",
    "        \"\"\"\n",
    "        if self.all_lower and word not in self.idx_misc:\n",
    "            word = word.lower()\n",
    "        return self.w2vec.get(word, self.w2vec['<UNK>'])\n",
    "        \n",
    "    def encode(self, text, nlp, max_len, pad_end=True, trim_start=True):\n",
    "        \"\"\"Encode text so that each token is replaced by its integer index in \n",
    "        the vocab.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        text: str\n",
    "            Raw text to be encoded.\n",
    "        nlp: spacy.lang.en.English\n",
    "            Spacy tokenizer. Typically want to disable 'parser', 'tagger', and\n",
    "            'ner' as they aren't used here and slow down the encoding process.\n",
    "        max_len: int\n",
    "            Length of output encoding. If text is shorter, it will be padded \n",
    "            to fit the specified length. If text is longer, it will be \n",
    "            trimmed.\n",
    "        pad_end: bool\n",
    "            If True, add padding to the end of short sentences. If False, pad\n",
    "            the start of these sentences.\n",
    "        trim_start: bool\n",
    "            If True, trim off the start of sentences that are too long. If \n",
    "            False, trim off the end.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        np.array[int]: Array of length max_len containing integer indices\n",
    "            corresponding to the words passed in.\n",
    "        \"\"\"\n",
    "        output = np.ones(max_len) * self.idx('<PAD>')\n",
    "        encoded = [self.idx(tok.text) for tok in nlp(text)]\n",
    "        \n",
    "        # Trim sentence in case it's longer than max_len.\n",
    "        if len(encoded) > max_len:\n",
    "            if trim_start:\n",
    "                encoded = encoded[len(encoded) - max_len:]\n",
    "            else:\n",
    "                encoded = encoded[:max_len]\n",
    "\n",
    "        # Replace zeros at start or end, depending on choice of pad_end.\n",
    "        if pad_end:\n",
    "            output[:len(encoded)] = encoded\n",
    "        else:\n",
    "            output[max_len-len(encoded):] = encoded\n",
    "        return output.astype(int)\n",
    "    \n",
    "    def decode(self, idx):\n",
    "        \"\"\"Convert a list of indices to a string of words/tokens.\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        idx: list[int]\n",
    "            A list of integers indexing into the vocabulary. This will often \n",
    "            be the output of the encode() method.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        list[str]: A list of words/tokens reconstructed by indexing into the \n",
    "            vocabulary.\n",
    "        \"\"\"\n",
    "        return [self[i] for i in idx]\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        \"\"\"This will map an index (int) to a word (str).\n",
    "        \n",
    "        Parameters\n",
    "        -----------\n",
    "        i: int\n",
    "            Integer index for a word.\n",
    "            \n",
    "        Returns\n",
    "        --------\n",
    "        str: Word corresponding to the given index.\n",
    "            \n",
    "        Examples\n",
    "        ---------\n",
    "        >>> vocab = Vocabulary(w2idx, w2vec)\n",
    "        >>> vocab[1]\n",
    "        '<UNK>'\n",
    "        \"\"\"\n",
    "        return self.i2w[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        \"\"\"Number of words in vocabulary.\"\"\"\n",
    "        return len(self.w2idx)\n",
    "    \n",
    "    def __iter__(self):\n",
    "        for word in self.w2idx.keys():\n",
    "            yield word\n",
    "            \n",
    "    def __contains__(self, word):\n",
    "        return word in self.w2idx.keys()\n",
    "    \n",
    "    def __eq__(self, obj):\n",
    "        if not isinstance(obj, Vocabulary):\n",
    "            return False\n",
    "        \n",
    "        ignore = {'w2vec', 'embedding_matrix'}\n",
    "        attrs = [k for k, v in hdir(vocab).items() \n",
    "                 if v == 'attribute' and k not in ignore]\n",
    "        return all([getattr(self, attr) == getattr(obj, attr) \n",
    "                    for attr in attrs])\n",
    "    \n",
    "    def __repr__(self):\n",
    "        msg = f'Vocabulary({len(self)} words'\n",
    "        if self.dim > 1:\n",
    "            msg += f', {self.dim}-D embeddings'\n",
    "        return msg + ')'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load('en_core_web_sm')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocabulary(5002 words, 100-D embeddings)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "glove_path = '/Users/hmamin/data/glove/glove.6B.100d.txt'\n",
    "\n",
    "vocab = Vocabulary.from_glove_file(glove_path, 5_000)\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5002"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving vocabulary to tmp.pkl.\n"
     ]
    }
   ],
   "source": [
    "vocab.save('tmp.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v2 = Vocabulary.from_pickle('tmp.pkl')\n",
    "vocab == v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert v2.w2idx == vocab.w2idx\n",
    "assert v2.i2w == vocab.i2w\n",
    "assert v2.corpus_counts == vocab.corpus_counts\n",
    "assert v2.idx_misc == vocab.idx_misc\n",
    "assert v2.dim == vocab.dim\n",
    "assert len(v2.w2vec) == len(vocab.w2vec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "85"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del v2\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<PAD>', '<UNK>', 'the', ',', '.']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 <PAD> 1 [0. 0. 0.]\n",
      "1 <UNK> 1 [0. 0. 0.]\n",
      "2 the 2 [-0.038194 -0.24487   0.72812 ]\n",
      "3 , 3 [-0.10767  0.11053  0.59812]\n",
      "4 . 4 [-0.33979  0.20941  0.46348]\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    word = vocab[i]\n",
    "    print(i, word, vocab.idx(word), vocab.vector(word)[:3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = vocab.build_embedding_matrix()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.      ,  0.      ,  0.      ,  0.      ,  0.      ],\n",
       "       [ 0.      ,  0.      ,  0.      ,  0.      ,  0.      ],\n",
       "       [-0.038194, -0.24487 ,  0.72812 , -0.39961 ,  0.083172],\n",
       "       [-0.10767 ,  0.11053 ,  0.59812 , -0.54361 ,  0.67396 ],\n",
       "       [-0.33979 ,  0.20941 ,  0.46348 , -0.64792 , -0.38377 ],\n",
       "       [-0.1529  , -0.24279 ,  0.89837 ,  0.16996 ,  0.53516 ],\n",
       "       [-0.1897  ,  0.050024,  0.19084 , -0.049184, -0.089737],\n",
       "       [-0.071953,  0.23127 ,  0.023731, -0.50638 ,  0.33923 ],\n",
       "       [ 0.085703, -0.22201 ,  0.16569 ,  0.13373 ,  0.38239 ],\n",
       "       [-0.27086 ,  0.044006, -0.02026 , -0.17395 ,  0.6444  ]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(vocab.embedding_matrix)\n",
    "emb[:10, :5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1525"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del emb\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vocab.build_embedding_matrix(True)\n",
    "# vocab.embedding_matrix.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "54\n",
      "17\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'<PAD>': 0,\n",
       " '<UNK>': 1,\n",
       " 'blue': 2,\n",
       " '.': 3,\n",
       " 'sunny': 4,\n",
       " 'atmosphere': 5,\n",
       " 'the': 6,\n",
       " 'beach': 7,\n",
       " '\\n': 8,\n",
       " 'was': 9,\n",
       " 'and': 10,\n",
       " 'i': 11,\n",
       " 'went': 12,\n",
       " 'to': 13,\n",
       " 'today': 14,\n",
       " 'it': 15,\n",
       " 'warm': 16}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "raw = \"\"\"I went to the beach today. It was sunny and warm. The water was\n",
    "very blue and cold. The relaxing atmosphere provided a welcome break from the\n",
    "daily routine. beach sunny sunny sunny sunny atmosphere atmosphere atmosphere\n",
    "blue blue blue blue blue. beach next beach atmosphere\"\"\"\n",
    "corpus = [t.text for t in nlp(raw.lower(), disable=['parser', 'tagger', 'ner'])]\n",
    "print(len(corpus))\n",
    "\n",
    "filtered = vocab.filter_tokens(corpus, max_words=15)\n",
    "print(len(filtered))\n",
    "filtered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({'i': 1,\n",
       "         'went': 1,\n",
       "         'to': 1,\n",
       "         'the': 4,\n",
       "         'beach': 4,\n",
       "         'today': 1,\n",
       "         '.': 5,\n",
       "         'it': 1,\n",
       "         'was': 2,\n",
       "         'sunny': 5,\n",
       "         'and': 2,\n",
       "         'warm': 1,\n",
       "         'water': 1,\n",
       "         '\\n': 3,\n",
       "         'very': 1,\n",
       "         'blue': 6,\n",
       "         'cold': 1,\n",
       "         'relaxing': 1,\n",
       "         'atmosphere': 5,\n",
       "         'provided': 1,\n",
       "         'a': 1,\n",
       "         'welcome': 1,\n",
       "         'break': 1,\n",
       "         'from': 1,\n",
       "         'daily': 1,\n",
       "         'routine': 1,\n",
       "         'next': 1})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.corpus_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5001"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.filter_tokens(corpus, min_freq=2, inplace=True)\n",
    "len(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<PAD>': 0,\n",
       " '<UNK>': 1,\n",
       " 'blue': 2,\n",
       " '.': 3,\n",
       " 'sunny': 4,\n",
       " 'atmosphere': 5,\n",
       " 'the': 6,\n",
       " 'beach': 7,\n",
       " '\\n': 8,\n",
       " 'was': 9,\n",
       " 'and': 10}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.w2idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<PAD>',\n",
       " '<UNK>',\n",
       " 'blue',\n",
       " '.',\n",
       " 'sunny',\n",
       " 'atmosphere',\n",
       " 'the',\n",
       " 'beach',\n",
       " '\\n',\n",
       " 'was',\n",
       " 'and']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.i2w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.idx('BLUE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.all_lower = False\n",
    "vocab.idx('BLUE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.all_lower = True\n",
    "(vocab.vector('blue') == vocab.vector('BLUE')).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vocab.save('../data/vocab.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vocab2 = torch.load('../data/vocab.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(vocab.w2idx == vocab2.w2idx)\n",
    "# del vocab2\n",
    "# gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## custom idx_misc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocabulary(5004 words, 50-D embeddings)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_misc = {'<UNK>': 0, '<PAD>': 0, '<XXrep>': 2, '<XXcaps>': 3}\n",
    "vocab = Vocabulary.from_glove_file(glove_path.replace('100', '50'), 5_000, custom_misc)\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[4, 2381, 170, 6]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[vocab.idx(word) for word in ('the', 'boat', 'house', '.')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Without Glove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(text):\n",
    "    return [t.text for t in nlp(text, disable=['parser', 'tagger', 'ner'])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(text, vocab):\n",
    "    return [vocab.idx(t) for t in tokenize(text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocabulary(988 words)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_path = '/Users/hmamin/data/bbc/tech/401.txt'\n",
    "with open(text_path, 'r') as f:\n",
    "    tokens = tokenize(f.read())\n",
    "\n",
    "vocab = Vocabulary.from_tokens(tokens, all_lower=False)\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.vector('the')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.idx('The')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.idx('the')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 <PAD> 0 [0.]\n",
      "1 <UNK> 1 [0.]\n",
      "2 . 2 [0.]\n",
      "3 , 3 [0.]\n",
      "4 the 4 [0.]\n",
      "5 I 5 [0.]\n",
      "6 to 6 [0.]\n",
      "7 a 7 [0.]\n",
      "8 of 8 [0.]\n",
      "9 and 9 [0.]\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    word = vocab[i]\n",
    "    print(i, word, vocab.idx(word), vocab.vector(word))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocabulary(17 words)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.filter_tokens(tokens, max_words=15, min_freq=5, inplace=True)\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<PAD>', '<UNK>', '.', ',', 'the', 'I', 'to', 'a', 'of', 'and']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.],\n",
       "       [0.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.build_embedding_matrix()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocabulary(5002 words, 50-D embeddings)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab = Vocabulary.from_glove_file('/Users/hmamin/data/glove/glove.6B.50d.txt', 5_000)\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"It's kind of a really nice day.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.ones(5) * 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      "[22, 11, 923, 5, 9, 590, 3084, 124, 4]\n",
      "0 0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([  22,   11,  923,    5,    9,  590, 3084,  124,    4,    0])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded = vocab.encode(text, nlp, 10)\n",
    "encoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"it 's kind of a really nice day . <PAD>\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.decode(encoded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<UNK>'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<PAD>', '<UNK>', 'the', ',', '.', 'of']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[:6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load('en_core_web_sm', disable=['parser', 'tagger', 'ner'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder:\n",
    "    \n",
    "    def __init__(self, nlp, vocab, max_len, pad_end=True, trim_start=True):\n",
    "        self.nlp = nlp\n",
    "        self.vocab = vocab\n",
    "        self.max_len = max_len\n",
    "        self.pad_end = pad_end\n",
    "        self.trim_start = trim_start\n",
    "        self.pad_idx = self.vocab.idx('<PAD>')\n",
    "        \n",
    "    def encode(self, text):\n",
    "        output = np.zeros(self.max_len)\n",
    "        encoded = [self.vocab.idx(tok.text) \n",
    "                   for tok in self.nlp(text.lower())]\n",
    "        \n",
    "        # Trim sentence in case it's longer than max_len.\n",
    "        if len(encoded) > self.max_len:\n",
    "            if self.trim_start:\n",
    "                encoded = encoded[len(encoded) - self.max_len:]\n",
    "            else:\n",
    "                encoded = encoded[:self.max_len]\n",
    "\n",
    "        # Replace zeros at start or end, depending on choice of pad_end.\n",
    "        if self.pad_end:\n",
    "            output[:len(encoded)] = encoded\n",
    "        else:\n",
    "            output[self.max_len-len(encoded):] = encoded\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "PAD END True TRIM START True\n",
      "[2.200e+01 1.100e+01 9.230e+02 5.000e+00 9.000e+00 5.900e+02 3.084e+03\n",
      " 1.240e+02 4.000e+00 1.000e+00]\n",
      "[  16. 2239.    3.    7.   22.   16.    9. 1868.  366.    4.]\n",
      "\n",
      "PAD END True TRIM START False\n",
      "[2.200e+01 1.100e+01 9.230e+02 5.000e+00 9.000e+00 5.900e+02 3.084e+03\n",
      " 1.240e+02 4.000e+00 1.000e+00]\n",
      "[2.000e+00 1.571e+03 3.400e+01 4.131e+03 3.000e+00 2.000e+00 3.507e+03\n",
      " 1.600e+01 2.239e+03 3.000e+00]\n",
      "\n",
      "PAD END False TRIM START True\n",
      "[1.000e+00 2.200e+01 1.100e+01 9.230e+02 5.000e+00 9.000e+00 5.900e+02\n",
      " 3.084e+03 1.240e+02 4.000e+00]\n",
      "[  16. 2239.    3.    7.   22.   16.    9. 1868.  366.    4.]\n",
      "\n",
      "PAD END False TRIM START False\n",
      "[1.000e+00 2.200e+01 1.100e+01 9.230e+02 5.000e+00 9.000e+00 5.900e+02\n",
      " 3.084e+03 1.240e+02 4.000e+00]\n",
      "[2.000e+00 1.571e+03 3.400e+01 4.131e+03 3.000e+00 2.000e+00 3.507e+03\n",
      " 1.600e+01 2.239e+03 3.000e+00]\n"
     ]
    }
   ],
   "source": [
    "for pad_end_bool in (True, False):\n",
    "    for trim_start_bool in (True, False):\n",
    "        print('\\nPAD END', pad_end_bool, 'TRIM START', trim_start_bool)\n",
    "        print(vocab.encode(text, nlp, 10, pad_end_bool, trim_start_bool))\n",
    "        print(vocab.encode(text_long, nlp, 10, pad_end_bool, trim_start_bool))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  22.,   11.,  923.,    5.,    9.,  590., 3084.,  124.,    4.,\n",
       "          0.])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = Encoder(nlp, vocab, 10)\n",
    "encoder.encode(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.000e+00, 1.571e+03, 3.400e+01, 4.131e+03, 3.000e+00, 2.000e+00,\n",
       "       3.507e+03, 1.600e+01, 2.239e+03, 3.000e+00])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_long = \"The stars are bright, the sky is dark, and it is a cold night.\"\n",
    "encoder.encode(text_long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.000e+00, 1.571e+03, 3.400e+01, 4.131e+03, 3.000e+00, 2.000e+00,\n",
       "       3.507e+03, 1.600e+01, 2.239e+03, 3.000e+00])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = Encoder(nlp, vocab, 10, pad_end=False, trim_start=True)\n",
    "encoder.encode(text_long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([   0.,   22.,   11.,  923.,    5.,    9.,  590., 3084.,  124.,\n",
       "          4.])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder.encode(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['The',\n",
       " 'stars',\n",
       " 'are',\n",
       " 'bright',\n",
       " ',',\n",
       " 'the',\n",
       " 'sky',\n",
       " 'is',\n",
       " 'dark',\n",
       " ',',\n",
       " 'and',\n",
       " 'it',\n",
       " 'is',\n",
       " 'a',\n",
       " 'cold',\n",
       " 'night',\n",
       " '.']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[t.text for t in nlp(text_long)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "','"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab[3]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
