{
 "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",
    "import boto3\n",
    "from collections import deque\n",
    "from functools import partial\n",
    "import mmh3\n",
    "from multiprocessing import Pool\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm.auto import tqdm\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## File Handling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class BotoUploader:\n",
    "    \"\"\"Uploads files to S3. Built as a public alternative to Accio. Note to \n",
    "    self: the interfaces are not identical so be careful to know which you're\n",
    "    using.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, bucket, verbose=True):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        bucket: str\n",
    "            Name of s3 bucket to upload to. For a single project, I generally\n",
    "            stick to a single bucket so we can usually keep this fixed. We can\n",
    "            always change the attribute later if necessary.\n",
    "        verbose: bool\n",
    "            If True, print message when downloading each file.\n",
    "        \"\"\"\n",
    "        self.s3 = boto3.resource('s3')\n",
    "        self.bucket = bucket\n",
    "        self.verbose = verbose\n",
    "        \n",
    "    def upload_file(self, path, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Upload a single local file. By default, its path in S3 will be the \n",
    "        same as its local path. Usually, this means your S3 bucket will have a \n",
    "        single directory called \"data\" which corresponds exactly to your local \n",
    "        \"data\" directory.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        path: str or Path\n",
    "            Local path to the file to upload. This must be relative to the \n",
    "            project root directory (e.g. \"data/models/v1/model.gz\").\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}.\n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. All four combinations of retain_tree \n",
    "            (True/False) and s3_dir (empty/non-empty) are supported.\n",
    "        \"\"\"\n",
    "        path = str(path)\n",
    "        s3_path = self._convert_local_path(path, s3_dir, retain_tree)\n",
    "        if self.verbose: print(f'Uploading {path} -> {s3_path}.')\n",
    "        self.s3.meta.client.upload_file(path, self.bucket, s3_path)\n",
    "        \n",
    "    def upload_files(self, paths, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Upload multiple files. Currently does not support parallelized\n",
    "        uploads: need to figure out a workaround because self.s3 cannot be\n",
    "        pickled.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        paths: Iterable[str or Path]\n",
    "            Sequence of file paths to upload.\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}.\n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. All four combinations of retain_tree \n",
    "            (True/False) and s3_dir (empty/non-empty) are supported.\n",
    "        \"\"\"\n",
    "        for p in paths:\n",
    "            self.upload_file(p, s3_dir=s3_dir, retain_tree=retain_tree)\n",
    "            \n",
    "    def upload_folder(self, dirname, s3_dir, retain_tree=True, recurse=True,\n",
    "                      keep_fn=None):\n",
    "        \"\"\"Upload all files in a directory.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        dirname: str or Path\n",
    "            Directory to upload.\n",
    "        s3_dir: str\n",
    "            If provided, this will pre prepended to each path: {s3_dir}/{path}. \n",
    "            Otherwise, s3 paths will be the same as local paths.\n",
    "        retain_tree: bool\n",
    "            If True, the local file structure will be retained. Otherwise,\n",
    "            only the base name is kept. When uploading recursively, we require\n",
    "            that the file tree is retained.\n",
    "        recurse: bool\n",
    "            If True, upload all files in subdirectories as well.\n",
    "        keep_fn: None or callable\n",
    "            If provided, this should be a function that accepts a filename as\n",
    "            input and returns a boolean specifying whether to include it in the\n",
    "            upload or not. For example:\n",
    "            \n",
    "            lambda x: os.path.splitext(x)[-1] != '.pkl'\n",
    "            \n",
    "            keeps all files except those with an '.pkl' extension. \n",
    "        \"\"\"\n",
    "        if recurse and not retain_tree:\n",
    "            raise ValueError('retain_tree must be True when uploading '\n",
    "                             'recursively.')\n",
    "        pat = os.path.join(str(dirname), '**' if recurse else '*')\n",
    "        # glob's recursive option only has an effect when using '**'.\n",
    "        paths = (o for o in glob(pat, recursive=True) if os.path.isfile(o))\n",
    "        if keep_fn: paths = filter(keep_fn, paths)\n",
    "        self.upload_files(paths, s3_dir, retain_tree)\n",
    "            \n",
    "    def _convert_local_path(self, path, s3_dir='', retain_tree=True):\n",
    "        \"\"\"Convert local path to s3 path. See public methods for parameter\n",
    "        documentation.\n",
    "        \"\"\"\n",
    "        path = path if retain_tree else os.path.basename(path)\n",
    "        return os.path.join(s3_dir, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "up = BotoUploader('gg-datascience')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No S3 prefix, Yes retain file tree:\n",
      "data/v1/history.csv\n",
      "\n",
      "Yes S3 prefix, Yes retain file tree:\n",
      "hmamin/data/v1/history.csv\n",
      "\n",
      "Yes S3 prefix, No retain file tree:\n",
      "hmamin/history.csv\n",
      "\n",
      "No S3 prefix, No retain file tree:\n",
      "history.csv\n"
     ]
    }
   ],
   "source": [
    "ft = up._convert_local_path('data/v1/history.csv')\n",
    "tt = up._convert_local_path('data/v1/history.csv', 'hmamin')\n",
    "tf = up._convert_local_path('data/v1/history.csv', 'hmamin', retain_tree=False)\n",
    "ff = up._convert_local_path('data/v1/history.csv', retain_tree=False)\n",
    "\n",
    "print('No S3 prefix, Yes retain file tree:\\n' + ft)\n",
    "print('\\nYes S3 prefix, Yes retain file tree:\\n' + tt)\n",
    "print('\\nYes S3 prefix, No retain file tree:\\n' + tf)\n",
    "print('\\nNo S3 prefix, No retain file tree:\\n' + ff)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
