{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data\n",
    "\n",
    "> Tools to help construct datasets, which may be related to loading, processing, or encoding data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from collections import deque\n",
    "import mmh3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "import warnings\n",
    "\n",
    "from htools import auto_repr, save, load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only needed for testing.\n",
    "from collections import Counter\n",
    "from itertools import chain\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from htools import eprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def probabilistic_hash_item(x, n_buckets, mode=int, n_hashes=3):\n",
    "    \"\"\"Slightly hacky way to probabilistically hash an integer by\n",
    "    first converting it to a string.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: int\n",
    "        The integer or string to hash.\n",
    "    n_buckets: int\n",
    "        The number of buckets that items will be mapped to. Typically \n",
    "        this would occur outside the hashing function, but since \n",
    "        the intended use case is so narrow here it makes sense to me \n",
    "        to include it here.\n",
    "    mode: type\n",
    "        The type of input you want to hash. This is user-provided to prevent\n",
    "        accidents where we pass in a different item than intended and hash \n",
    "        the wrong thing. One of (int, str). When using this inside a\n",
    "        BloomEmbedding layer, this must be `int` because there are no\n",
    "        string tensors. When used inside a dataset or as a one-time\n",
    "        pre-processing step, you can choose either as long as you\n",
    "        pass in the appropriate inputs.\n",
    "    n_hashes: int\n",
    "        The number of times to hash x, each time with a different seed.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[int]: A list of integers with length `n_hashes`, where each integer\n",
    "        is in [0, n_buckets).\n",
    "    \"\"\"\n",
    "    # Check type to ensure we don't accidentally hash Tensor(5) instead of 5.\n",
    "    assert isinstance(x, mode), f'Input `x` must have type {mode}.'\n",
    "    return [mmh3.hash(str(x), i, signed=False) % n_buckets \n",
    "            for i in range(n_hashes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def probabilistic_hash_tensor(x_r2, n_buckets, n_hashes=3, pad_idx=0):\n",
    "    \"\"\"Hash a rank 2 LongTensor.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x_r2: torch.LongTensor\n",
    "        Rank 2 tensor of integers. Shape: (bs, seq_len)\n",
    "    n_buckets: int\n",
    "        Number of buckets to hash items into (i.e. the number of \n",
    "        rows in the embedding matrix). Typically a moderately large\n",
    "        prime number, like 251 or 997.\n",
    "    n_hashes: int\n",
    "        Number of hashes to take for each input index. This determines\n",
    "        the number of rows of the embedding matrix that will be summed\n",
    "        to get the representation for each word. Typically 2-5.\n",
    "    pad_idx: int or None\n",
    "        If you want to pad sequences with vectors of zeros, pass in an\n",
    "        integer (same as the `padding_idx` argument to nn.Embedding).\n",
    "        If None, no padding index will be used. The sequences must be\n",
    "        padded before passing them into this function.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.LongTensor: Tensor of indices where each row corresponds\n",
    "        to one of the input indices. Shape: (bs, seq_len, n_hashes)\n",
    "    \"\"\"\n",
    "    return torch.tensor(\n",
    "        [[probabilistic_hash_item(x.item(), n_buckets, int, n_hashes) \n",
    "          if x != pad_idx else [pad_idx]*n_hashes for x in row]\n",
    "         for row in x_r2]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sents = [\n",
    "    'I walked to the store so I hope it is not closed.',\n",
    "    'The theater is closed today and the sky is grey.',\n",
    "    'His dog is brown while hers is grey.'\n",
    "]\n",
    "labels = [0, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data(Dataset):\n",
    "    \n",
    "    def __init__(self, sentences, labels, seq_len):\n",
    "        x = [s.split(' ') for s in sentences]\n",
    "        self.w2i = self.make_w2i(x)\n",
    "        self.seq_len = seq_len\n",
    "        self.x = self.encode(x)\n",
    "        self.y = torch.tensor(labels)\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "    \n",
    "    def make_w2i(self, tok_rows):\n",
    "        return {k: i for i, (k, v) in \n",
    "                enumerate(Counter(chain(*tok_rows)).most_common(), 1)}\n",
    "    \n",
    "    def encode(self, tok_rows):\n",
    "        enc = np.zeros((len(tok_rows), self.seq_len), dtype=int)\n",
    "        for i, row in enumerate(tok_rows):\n",
    "            trunc = [self.w2i.get(w, 0) for w in row[:self.seq_len]]\n",
    "            enc[i, :len(trunc)] = trunc\n",
    "        return torch.tensor(enc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We construct a toy dataset with a vocabulary of size 23. In reality, you might wish to lowercase text or use a better tokenizer, but this is sufficient for the purposes of demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = Data(sents, labels, 10)\n",
    "len(ds.w2i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]), tensor([0, 1, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dl = DataLoader(ds, batch_size=3)\n",
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 10])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We hash each word index 4 times, as specified by the `n_hashes` parameter in `probabilistic_hash_tensor`. Notice that we only use 7 buckets, meaning the embedding matrix will have 7 rows rather than 23 (not counting a padding row)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x shape: torch.Size([3, 10])\n",
      "x_hashed shape: torch.Size([3, 10, 4])\n"
     ]
    }
   ],
   "source": [
    "x_hashed = probabilistic_hash_tensor(x, n_buckets=7, n_hashes=4)\n",
    "print('x shape:', x.shape)\n",
    "print('x_hashed shape:', x_hashed.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, each row of 4 numbers encodes a single word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2, 0, 2, 2],\n",
       "         [1, 6, 2, 4],\n",
       "         [2, 0, 1, 4],\n",
       "         [5, 2, 6, 5],\n",
       "         [1, 5, 1, 3],\n",
       "         [0, 4, 0, 0],\n",
       "         [2, 0, 2, 2],\n",
       "         [2, 4, 0, 2],\n",
       "         [3, 4, 4, 6],\n",
       "         [5, 0, 3, 6]],\n",
       "\n",
       "        [[5, 4, 4, 2],\n",
       "         [5, 3, 3, 1],\n",
       "         [5, 0, 3, 6],\n",
       "         [2, 1, 1, 1],\n",
       "         [2, 4, 4, 4],\n",
       "         [2, 4, 2, 6],\n",
       "         [5, 2, 6, 5],\n",
       "         [3, 5, 0, 0],\n",
       "         [5, 0, 3, 6],\n",
       "         [1, 5, 2, 6]],\n",
       "\n",
       "        [[5, 5, 3, 4],\n",
       "         [4, 5, 5, 1],\n",
       "         [5, 0, 3, 6],\n",
       "         [6, 2, 0, 6],\n",
       "         [4, 2, 6, 1],\n",
       "         [3, 6, 1, 6],\n",
       "         [5, 0, 3, 6],\n",
       "         [1, 5, 2, 6],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_hashed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See how each word is mapped to a list of 4 indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I [2, 0, 2, 2]\n",
      "walked [1, 6, 2, 4]\n",
      "to [2, 0, 1, 4]\n",
      "the [5, 2, 6, 5]\n",
      "store [1, 5, 1, 3]\n",
      "so [0, 4, 0, 0]\n",
      "I [2, 0, 2, 2]\n",
      "hope [2, 4, 0, 2]\n",
      "it [3, 4, 4, 6]\n",
      "is [5, 0, 3, 6]\n"
     ]
    }
   ],
   "source": [
    "for word, i in zip(sents[0].split(' '), x[0]):\n",
    "    print(word, probabilistic_hash_item(i.item(), 7, int, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that hashing the words directly is also possible, but the resulting hashes will be **different** than if hashing after encoding words as integers. This is fine as long as you are consistent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0: ('I', [0, 5, 5])\n",
      " 1: ('walked', [2, 5, 1])\n",
      " 2: ('to', [10, 4, 6])\n",
      " 3: ('the', [4, 1, 4])\n",
      " 4: ('store', [4, 6, 3])\n",
      " 5: ('so', [7, 8, 8])\n",
      " 6: ('I', [0, 5, 5])\n",
      " 7: ('hope', [1, 2, 7])\n",
      " 8: ('it', [3, 9, 0])\n",
      " 9: ('is', [3, 1, 3])\n",
      "10: ('not', [6, 10, 4])\n",
      "11: ('closed.', [3, 6, 10])\n",
      "\n",
      " 0: ('The', [1, 6, 9])\n",
      " 1: ('theater', [8, 10, 2])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('closed', [5, 5, 0])\n",
      " 4: ('today', [3, 10, 8])\n",
      " 5: ('and', [7, 2, 4])\n",
      " 6: ('the', [4, 1, 4])\n",
      " 7: ('sky', [1, 2, 9])\n",
      " 8: ('is', [3, 1, 3])\n",
      " 9: ('grey.', [7, 6, 7])\n",
      "\n",
      " 0: ('His', [0, 10, 3])\n",
      " 1: ('dog', [8, 6, 6])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('brown', [9, 8, 9])\n",
      " 4: ('while', [9, 2, 8])\n",
      " 5: ('hers', [0, 5, 4])\n",
      " 6: ('is', [3, 1, 3])\n",
      " 7: ('grey.', [7, 6, 7])\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for row in [s.split(' ') for s in sents]:\n",
    "    eprint(list(zip(row, (probabilistic_hash_item(word, 11, str) for word in row))))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we show that we can obtain unique representations for >99.9% of words in a vocabulary of 30,000 words with a far smaller embedding matrix. The number of buckets is the number of rows in the embedding matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unique_combos(tups):\n",
    "    return len(set(tuple(sorted(x)) for x in tups))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hash_all_idx(vocab_size, n_buckets, n_hashes):\n",
    "    return [probabilistic_hash_item(i, n_buckets, int, n_hashes) \n",
    "            for i in range(vocab_size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Buckets: 127 \n",
      "Hashes: 5 \n",
      "Unique combos: 29998 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 251 \n",
      "Hashes: 4 \n",
      "Unique combos: 29996 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 997 \n",
      "Hashes: 3 \n",
      "Unique combos: 29997 \n",
      "% unique: 0.9999\n",
      "\n",
      "\n",
      "Buckets: 5003 \n",
      "Hashes: 2 \n",
      "Unique combos: 29969 \n",
      "% unique: 0.999\n"
     ]
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "buckets2hashes = {127: 5,\n",
    "                  251: 4,\n",
    "                  997: 3,\n",
    "                  5_003: 2}\n",
    "for b, h in buckets2hashes.items():\n",
    "    tups = hash_all_idx(vocab_size, b,  h)\n",
    "    unique = unique_combos(tups)\n",
    "    print('\\n\\nBuckets:', b, '\\nHashes:', h, '\\nUnique combos:', unique,\n",
    "          '\\n% unique:', round(unique/30_000, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@auto_repr\n",
    "class LazyDataset(Dataset):\n",
    "    \"\"\"Lazily load batches from an enormous dataframe that can't fit into \n",
    "    memory.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, df_path, length, shuffle, chunksize=1_000, \n",
    "                 c=2, classes=('neg', 'pos'), **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        df_path: str\n",
    "            File path of dataframe to load.\n",
    "        length: int\n",
    "            Number of rows of data to use. This is required so that we don't \n",
    "            have to go through the whole file and count the number of lines,\n",
    "            which can be enormous with a big dataset. It also makes it easy to\n",
    "            work with a subset (the data should already be shuffled, so \n",
    "            choosing the top n rows is fine).\n",
    "        shuffle: bool\n",
    "            If True, shuffle the data in each chunk. Note that if batch size\n",
    "            is close to chunk size, this will have minimal effect. If possible,\n",
    "            the training set should therefore load as large a chunk as \n",
    "            possible if we want to shuffle the data. Shuffling is unnecessary \n",
    "            for the validation set.\n",
    "        chunksize: int\n",
    "            Number of rows of df to load at a time. This should usually \n",
    "            be significantly larger than the batch size in order to retain\n",
    "            some randomness in the batches.\n",
    "        c: int\n",
    "            Number of classes. Used if training with FastAI.\n",
    "        classes: iterable\n",
    "            List of tuple of class names. Used if training with FastAI.\n",
    "        kwargs: any\n",
    "            Additional keyword arguments to pass to `read_csv`, eg. \n",
    "            compression='gzip'.\n",
    "        \"\"\"\n",
    "        if length < chunksize:\n",
    "            warnings.warn('Total # of rows < 1 full chunk. LazyDataset may '\n",
    "                          'not be necessary.')\n",
    "\n",
    "        self.length = length\n",
    "        self.shuffle = shuffle\n",
    "        self.chunksize = chunksize\n",
    "        self.df_path = df_path\n",
    "        self.df = None\n",
    "        self.chunk = None\n",
    "        self.chunk_idx = None\n",
    "        self.df_kwargs = kwargs\n",
    "        \n",
    "        # Additional attributes required by FastAI. \n",
    "        # c: Number of classes in model.\n",
    "        self.c = c\n",
    "        self.classes = list(classes)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"Because not all indices are loaded at once, we must do shuffling\n",
    "        in the dataset rather than the dataloader (e.g. if the loader randomly\n",
    "        samples index 5000 but we have indices 0-500 loaded, it will be\n",
    "        unavailable).\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        idx: int\n",
    "            Retrieve item i in dataset.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        tuple[np.array]: x array, y array\n",
    "        \"\"\"\n",
    "        # Load next chunk of data if necessary. Must specify nrows, otherwise\n",
    "        # we will chunk through the whole file.\n",
    "        if not self.chunk_idx:\n",
    "            while True:\n",
    "                try:\n",
    "                    self.chunk = self.df.get_chunk()\n",
    "                    break\n",
    "                except (AttributeError, StopIteration):\n",
    "                    self.df = pd.read_csv(self.df_path, engine='python',\n",
    "                                          chunksize=self.chunksize,\n",
    "                                          nrows=len(self),\n",
    "                                          **self.df_kwargs)\n",
    "\n",
    "            self.chunk_idx = self.chunk.index.values\n",
    "            if self.shuffle: np.random.shuffle(self.chunk_idx)\n",
    "            self.chunk_idx = deque(self.chunk_idx)\n",
    "            \n",
    "        *x, y = self.chunk.loc[self.chunk_idx.popleft()].values\n",
    "        return np.array(x), y.astype(float)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
