{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp fast_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast Utils\n",
    "\n",
    "> Wrappers and helpers for interacting with fastai."
   ]
  },
  {
   "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 numpy as np\n",
    "import torch\n",
    "\n",
    "from htools import auto_repr, valuecheck"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def most_common_errors(interp):\n",
    "    \"\"\"More concise version of `most_confused`. Find the single most common\n",
    "    error for each true class.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    interp: fastai ClassificationInterpretation\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    dict[str, str]: Map true class to its most common incorrectly predicted\n",
    "        class.\n",
    "    \"\"\"\n",
    "    res = {}\n",
    "    for label, pred, _ in interp.most_confused():\n",
    "        if label not in res: res[label] = pred\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def n_groups(learn):\n",
    "    \"\"\"Quickly get the number of layer groups in a fastai learner.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    learn: fastai Learner\n",
    "        It seems that the models themselves don't actually specify groups the\n",
    "        way I've been doing in torch. Guessing this is because they've reworked\n",
    "        the API so everything is Sequential and can be indexed into. This does\n",
    "        make it easier to experiment with different groupings.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    int: 4 for default awd_lstm.\n",
    "    \"\"\"\n",
    "    return len(learn.opt.param_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def complete_sentence(learn, start_phrase, n_words=30, n_samples=3, temp=.75):\n",
    "    \"\"\"Generate text from a given input. This is a decent way to get a glimpse\n",
    "    of how a language model is doing.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    learn: fastai Learner\n",
    "    start_phrase: str\n",
    "        The prompt (start of a sentence) that you want the model to complete.\n",
    "    n_words: int\n",
    "        Number of words to generate after the given prompt.\n",
    "    n_samples: int\n",
    "        Number of sample sentences to generate.\n",
    "    temp: float\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    list[str]: Each item is one completed sentence.\n",
    "    \"\"\"\n",
    "    return [learn.predict(start_phrase, n_words, temperature=temp)\n",
    "            for _ in range(n_samples)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@auto_repr\n",
    "class LRPicker:\n",
    "    \"\"\"Take user-suggested LR into account as an \"anchor\" to ensure `lr_find`\n",
    "    doesn't choose something too far out of the ordinary.\n",
    "    \"\"\"\n",
    "\n",
    "    @valuecheck\n",
    "    def __init__(self, sugg_type:('lr_min', 'lr_steep', 'avg')='lr_min',\n",
    "                 resolve:('avg', 'auto', 'manual')='avg', tol_ratio=.1):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        sugg_type: str\n",
    "            Determines type of auto LR to use. \"avg\" will take the mean of\n",
    "            lr_min and lr_steep.\n",
    "        resolve: str\n",
    "            Determines method of resolving \"disagreements\" between default LR\n",
    "            and suggested LR. You can default to the auto-LR, the\n",
    "            user-suggested LR, or the mean of the two.\n",
    "        tol_ratio: float\n",
    "            Specifies minimum ratio of LR to user-provided LR to consider the\n",
    "            auto-suggestion as acceptable. This should always be between 0 and\n",
    "            1 (the ratio is calculated such that the smaller value is in the\n",
    "            numerator).\n",
    "        \"\"\"\n",
    "        self.sugg_type = sugg_type\n",
    "        self.resolve = resolve\n",
    "        self.tol_ratio = tol_ratio\n",
    "\n",
    "    def lr_find(self, learn, manual_sugg=None, find=True):\n",
    "        \"\"\"Find learning rate, incorporating manual suggestion and results of \n",
    "        `learn.lr_find()`.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        learn: fastai Learner\n",
    "        manual_sugg: float or None\n",
    "            User-suggested learning rate.\n",
    "        find: bool\n",
    "            If True, run lr_find to get an auto-suggestion. Otherwise just use\n",
    "            user-suggested LR.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        float: Learning rate to use.\n",
    "        \"\"\"\n",
    "        if find:\n",
    "            auto_sugg = learn.lr_find()\n",
    "            print('Auto suggestion:', auto_sugg)\n",
    "            return self._approve_lr(auto_sugg, manual_sugg)\n",
    "        return manual_sugg\n",
    "\n",
    "    def _approve_lr(self, auto_sugg, manual_sugg=None):\n",
    "        \"\"\"Take auto LR and user-suggested into account to choose a final LR.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        auto_sugg: SuggestedLRs\n",
    "            Object returned by learn.lr_find(). Basically a namedtuple I think.\n",
    "        manual_sugg: float or None\n",
    "            User-suggested learning rate.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        float: learning rate\n",
    "        \"\"\"\n",
    "        if self.sugg_type == 'avg':\n",
    "            sugg_lr = (auto_sugg.lr_min + auto_sugg.lr_steep) / 2\n",
    "        else:\n",
    "            sugg_lr = getattr(auto_sugg, self.sugg_type)\n",
    "\n",
    "        if not manual_sugg or self.resolve == 'auto' or \\\n",
    "                np.divide(*sorted([sugg_lr, manual_sugg])) >= self.tol_ratio:\n",
    "            return sugg_lr\n",
    "        if self.resolve == 'avg':\n",
    "            return (sugg_lr + manual_sugg) / 2\n",
    "        if self.resolve == 'manual':\n",
    "            return manual_sugg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@auto_repr\n",
    "class ULMFineTuner:\n",
    "    \"\"\"Fine tune language model using ULM Fit procedure. I noticed the built-in\n",
    "    `fine_tune` method does not unfreeze 1 layer at a time as the paper\n",
    "    describes - not sure if they found that to be a better practice or if it's\n",
    "    just simpler for an automated method.\n",
    "    \n",
    "    Originally, part of the reason for building this was to also decrease the\n",
    "    batch size at each stage since unfreezing eats up more memory with stored\n",
    "    gradients. However, I decided I'd rather not have to account for changing\n",
    "    batch size when selecting each stage's LR (we could run lr_find before each\n",
    "    stage but I opted for the simpler approach).\n",
    "    \"\"\"\n",
    "    moms = (0.8, 0.7, 0.8)\n",
    "    lr_div = 2.6\n",
    "\n",
    "    def __init__(self, learn, name_fmt='cls_stage_{}'):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        learn: fastai Learner\n",
    "        name_fmt: str\n",
    "            Determines name of model that will be saved at each stage.\n",
    "        \"\"\"\n",
    "        self.learn = learn\n",
    "        # Usually 4 layer groups in fastai text models, but check just in case.\n",
    "        self.groups = len(learn.opt.param_groups)\n",
    "        self.name_fmt = name_fmt\n",
    "\n",
    "    def fine_tune(self, lr, unfrozen_epochs=15, **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        lr: float\n",
    "            The learning rate to use for the last layer in the first stage of\n",
    "            training. The learning rate is halved after each stage. When\n",
    "            earlier layers are unfrozen, they will use a smaller LR as\n",
    "            described in the ULM Fit paper.\n",
    "        unfrozen_epochs: int\n",
    "            Number of epochs to fit after unfreezing everything. Couldn't find\n",
    "            a good guideline on this but I think it's fine to be fairly high\n",
    "            because we can use a callback to save the best model and stop early\n",
    "            if appropriate.\n",
    "        kwargs: any\n",
    "            Passed on to `learn.fit_one_cycle()`.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        None\n",
    "        \"\"\"\n",
    "        for stage in range(1, self.groups+1):\n",
    "            epochs = 1 if stage < self.groups else unfrozen_epochs\n",
    "            self.learn.freeze_to(-stage)\n",
    "            self.learn.fit_one_cycle(epochs,\n",
    "                                     slice(lr/self.lr_div**self.groups, lr),\n",
    "                                     moms=self.moms, **kwargs)\n",
    "            self.learn.save(self.name_fmt.format(stage))\n",
    "            lr /= 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class FastLabelEncoder:\n",
    "    \"\"\"Use a fastai learner to mimic a sklearn label encoder. This can be\n",
    "    useful because our standard evaluation code is often built when we are\n",
    "    trying out simple baseline models (e.g. logistic regression using sklearn).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, learn):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        learn: fastai Learner\n",
    "            The dataloaders must have a `vocab` attribute (text dataloaders do,\n",
    "            not sure about others).\n",
    "        \"\"\"\n",
    "        self.i2o = self.classes_ = learn.dls.vocab[1]\n",
    "        self.o2i = {o: i for i, o in enumerate(self.i2o)}\n",
    "\n",
    "    def transform(self, label_names):\n",
    "        \"\"\"Convert strings to integer indices.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        label_names: Iterable[str]\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        list[int]\n",
    "        \"\"\"\n",
    "        return [self.o2i[o] for o in label_names]\n",
    "\n",
    "    def inverse_transform(self, label_idx):\n",
    "        \"\"\"Convert integer indices to strings.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        label_idx: Iterable[int]\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        list[str]\n",
    "        \"\"\"\n",
    "        return [self.i2o[i] for i in label_idx]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
