OCR-Transformer-LM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prelims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:02.123407Z",
     "start_time": "2019-05-24T20:53:01.732546Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:03.687774Z",
     "start_time": "2019-05-24T20:53:02.126096Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/adamschiller/anaconda3/envs/fastai-cpu/lib/python3.6/site-packages/sklearn/ensemble/weight_boosting.py:29: DeprecationWarning: numpy.core.umath_tests is an internal NumPy module and should not be imported. It will be removed in a future NumPy release.\n",
      "  from numpy.core.umath_tests import inner1d\n"
     ]
    }
   ],
   "source": [
    "from fastai.conv_learner import *\n",
    "from fastai.text import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:03.756264Z",
     "start_time": "2019-05-24T20:53:03.691776Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:03.815556Z",
     "start_time": "2019-05-24T20:53:03.758926Z"
    }
   },
   "outputs": [],
   "source": [
    "def nonzero(pred):\n",
    "    ints = to_np(pred).astype(int)\n",
    "    return ints[np.nonzero(ints)]\n",
    "\n",
    "def char_label_text(pred, chunk=70):\n",
    "    return ''.join([itos[i] for i in nonzero(pred)])\n",
    "#     return '\\n'.join(textwrap.wrap(st, chunk))\n",
    "\n",
    "def char_split_text(pred):\n",
    "    return [itos[i] for i in nonzero(pred)]\n",
    "\n",
    "def word_label_text(pred, chunk=70):\n",
    "    return ' '.join([w_itos[i] for i in nonzero(pred)])\n",
    "#     return '\\n'.join(textwrap.wrap(st, chunk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:03.884583Z",
     "start_time": "2019-05-24T20:53:03.819592Z"
    }
   },
   "outputs": [],
   "source": [
    "def show_img(im, figsize=None, ax=None, alpha=None, title=None):\n",
    "    if not ax: fig,ax = plt.subplots(figsize=figsize)\n",
    "    ax.imshow(im, alpha=alpha)\n",
    "    if title: ax.set_title(title)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:03.954447Z",
     "start_time": "2019-05-24T20:53:03.889037Z"
    }
   },
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "\n",
    "# pulled from Sean Nareen's deepspeech decoder module\n",
    "# https://github.com/SeanNaren/deepspeech.pytorch/blob/master/decoder.py\n",
    "\n",
    "def cer(t, p):\n",
    "    \"\"\"\n",
    "    Computes the Character Error Rate, defined as the edit distance.\n",
    "    Arguments:\n",
    "        t (string): target space-separated sentence\n",
    "        p (string): prediction space-separated sentence\n",
    "    \"\"\"\n",
    "    t, p, = t.replace(' ', ''), p.replace(' ', '')\n",
    "    return Lev.distance(t, p)/len(t)\n",
    "\n",
    "def wer(s1, s2):\n",
    "    \"\"\"\n",
    "    Computes the Word Error Rate, defined as the edit distance between the\n",
    "    two provided sentences after tokenizing to words.\n",
    "    Arguments:\n",
    "        s1 (string): space-separated sentence\n",
    "        s2 (string): space-separated sentence\n",
    "    \"\"\"\n",
    "\n",
    "    # build mapping of words to integers\n",
    "    b = set(s1.split() + s2.split())\n",
    "    word2char = dict(zip(b, range(len(b))))\n",
    "\n",
    "    # map the words to a char array (Levenshtein package only accepts strings)\n",
    "    w1 = [chr(word2char[w]) for w in s1.split()]\n",
    "    w2 = [chr(word2char[w]) for w in s2.split()]\n",
    "\n",
    "    return Lev.distance(''.join(w1), ''.join(w2))/len(w1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:04.020411Z",
     "start_time": "2019-05-24T20:53:03.957024Z"
    }
   },
   "outputs": [],
   "source": [
    "def char_error_rate(preds, targs):\n",
    "    bs,sl = targs.size()      #=> ([bs, sl])\n",
    "    # preds.size()            #=> ([bs, sl, vs])\n",
    "        \n",
    "    res = torch.argmax(preds, dim=2)\n",
    "    error = 0\n",
    "    for i in range(bs):\n",
    "        p = char_label_text(res[i])\n",
    "        t = char_label_text(targs[i])\n",
    "        error += cer(t,p)\n",
    "    return error/bs\n",
    "\n",
    "def word_error_rate(preds, targs):\n",
    "    bs,sl = targs.size()      #=> ([bs, sl])\n",
    "    # preds.size()            #=> ([bs, sl, vs])\n",
    "        \n",
    "    res = torch.argmax(preds, dim=2)\n",
    "    error = 0\n",
    "    for i in range(bs):\n",
    "        p = word_label_text(res[i])\n",
    "        t = word_label_text(targs[i])\n",
    "        error += wer(t,p)\n",
    "    return error/bs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WikiText"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:20.700760Z",
     "start_time": "2019-05-24T20:53:20.563598Z"
    }
   },
   "outputs": [],
   "source": [
    "PATH = Path('data/wikitext/wikitext-2-raw')\n",
    "# PATH = Path('data/wikitext/wikitext-103-raw')\n",
    "\n",
    "with open(PATH/'wiki.train.raw') as file:  \n",
    "    trn = file.read()\n",
    "with open(PATH/'wiki.valid.raw') as file:  \n",
    "    val = file.read()\n",
    "with open(PATH/'wiki.test.raw') as file:  \n",
    "    tst = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:20.963373Z",
     "start_time": "2019-05-24T20:53:20.912647Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10918892"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(trn)\n",
    "# 2:    10918892\n",
    "# 103: 539566975"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-04T17:14:12.390140Z",
     "start_time": "2019-03-04T17:14:12.278180Z"
    }
   },
   "source": [
    "### clean the text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:24.999485Z",
     "start_time": "2019-05-24T20:53:24.947245Z"
    }
   },
   "outputs": [],
   "source": [
    "# convert spaced out \" strings \" to \"strings\"\n",
    "def despace_quotes(m):\n",
    "    m = m.group(0)   # entire matched string\n",
    "    m = m.replace('\" ','\"')\n",
    "    m = m.replace(' \"','\"')\n",
    "    return m\n",
    "\n",
    "def cleanup(x):\n",
    "    x = x.replace(' @-@ ', '-').replace(' =', '').replace('\\n \\n \\n', '\\n').replace('\\n \\n', '\\n').replace(\n",
    "        \" \\'\", \"\\'\").replace(' ,', ',').replace(' .', '.').replace(' :', ':').replace(' ;', ';').replace(\n",
    "        '( ', '(').replace(' )', ')').replace('[ ', '[').replace(' ]', ']').replace(' @.@ ', '.').replace(\n",
    "        ' @,@ ', ',')\n",
    "    x = re.sub(r'\\\"(.+?)\\\"', despace_quotes, x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:26.079152Z",
     "start_time": "2019-05-24T20:53:25.276117Z"
    }
   },
   "outputs": [],
   "source": [
    "trn = cleanup(trn)\n",
    "val = cleanup(val)\n",
    "tst = cleanup(tst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:26.257022Z",
     "start_time": "2019-05-24T20:53:26.204936Z"
    }
   },
   "outputs": [],
   "source": [
    "trn = trn + tst\n",
    "# len(trn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## IAM chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:04.374393Z",
     "start_time": "2019-05-10T15:02:04.325619Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "IAM_PATH = Path('data/IAM_handwriting')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:05.913160Z",
     "start_time": "2019-05-10T15:02:04.575123Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "maxTextLen = 32\n",
    "samples = []\n",
    "chars = set()\n",
    "\n",
    "with open(IAM_PATH/'ascii/words.txt') as f:\n",
    "    for line in f:\n",
    "        # ignore comment line\n",
    "        if not line or line[0]=='#':\n",
    "            continue\n",
    "\n",
    "        lineSplit = line.strip().split(' ')\n",
    "        assert len(lineSplit) >= 9\n",
    "\n",
    "        fileName = lineSplit[0]\n",
    "\n",
    "        # GT text are columns starting at 9\n",
    "        gtText = ''.join(lineSplit[8:])[:maxTextLen]\n",
    "        char_len = len(gtText)\n",
    "        chars = chars.union(set(list(gtText)))\n",
    "\n",
    "        # put sample into list\n",
    "        samples.append([fileName, gtText, char_len])\n",
    "    \n",
    "samples = np.stack(samples)\n",
    "df = pd.DataFrame(samples, columns=['filename', 'word', 'char_len'], )\n",
    "del samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:06.166064Z",
     "start_time": "2019-05-10T15:02:06.069434Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>filename</th>\n",
       "      <th>word</th>\n",
       "      <th>char_len</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>a01-000u-00-00</td>\n",
       "      <td>A</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>a01-000u-00-01</td>\n",
       "      <td>MOVE</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>a01-000u-00-02</td>\n",
       "      <td>to</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>a01-000u-00-03</td>\n",
       "      <td>stop</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>a01-000u-00-04</td>\n",
       "      <td>Mr.</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         filename  word  char_len\n",
       "0  a01-000u-00-00     A         1\n",
       "1  a01-000u-00-01  MOVE         4\n",
       "2  a01-000u-00-02    to         2\n",
       "3  a01-000u-00-03  stop         4\n",
       "4  a01-000u-00-04   Mr.         3"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['char_len'] = df.char_len.astype('int32')\n",
    "# df = df.loc[df['char_len'] > 3]\n",
    "df = df.loc[df['char_len'] < 20]\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:06.299534Z",
     "start_time": "2019-05-10T15:02:06.236111Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "590841"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iam_trn = ' '.join(df.word.values)\n",
    "len(iam_trn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:06.426762Z",
     "start_time": "2019-05-10T15:02:06.370895Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "iam_val = iam_trn[-60000:]   #last 60000\n",
    "iam_trn = iam_trn[:-60000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-10T15:02:06.610634Z",
     "start_time": "2019-05-10T15:02:06.557024Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "trn = trn+iam_trn\n",
    "val = val+iam_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Str to idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:53:35.486960Z",
     "start_time": "2019-05-24T20:53:35.436636Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "82"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "itos = pickle.load(open('data/IAM_handwriting/tmp/char_itos.pkl', 'rb'))\n",
    "stoi = collections.defaultdict(lambda: 0, {v:k for k,v in enumerate(itos)})\n",
    "len(itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:55:07.524625Z",
     "start_time": "2019-05-24T20:53:36.061250Z"
    }
   },
   "outputs": [],
   "source": [
    "#convert text into idxs\n",
    "trn_idx = np.array([stoi[c] for c in trn])\n",
    "val_idx = np.array([stoi[c] for c in val])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:55:08.469099Z",
     "start_time": "2019-05-24T20:55:08.241438Z"
    }
   },
   "outputs": [],
   "source": [
    "# remove unknown chars\n",
    "trn_mask = trn_idx.nonzero()\n",
    "trn_idx = trn_idx[trn_mask]\n",
    "\n",
    "val_mask = val_idx.nonzero()\n",
    "val_idx = val_idx[val_mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T20:55:09.529593Z",
     "start_time": "2019-05-24T20:55:09.475708Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' \\n Valkyria Chronicles III \\n Senj no Valkyria 3: Unrecorded Chronicles (Japanese: 3, lit. Valkyria of the Battlefield 3), commonly referred to as Valkyria Chronicles III outside Japan, is a tactical role-playing video game developed by Sega and Media.Vision for the PlayStation Portable. Released in January 2011 in Japan, it is the third game in the Valkyria series. Employing the same fusion of tactical and real-time gameplay as its predecessors, the story runs parallel to the first game and foll'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "''.join([itos[i] for i in trn_idx[:500]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Language Model Loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-04T17:37:28.301710Z",
     "start_time": "2019-03-04T17:37:28.212731Z"
    },
    "heading_collapsed": true
   },
   "source": [
    "## AWD-LSTM Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:19:22.083889Z",
     "start_time": "2019-05-21T15:19:22.034485Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "wd=1e-7\n",
    "bptt=30  # back prop through time\n",
    "bs=50\n",
    "opt_fn = partial(optim.Adam, betas=(0.8, 0.99))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:19:22.338588Z",
     "start_time": "2019-05-21T15:19:22.086471Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "trn_dl = LanguageModelLoader(trn_idx, bs, bptt)\n",
    "val_dl = LanguageModelLoader(val_idx, bs, bptt)\n",
    "md = LanguageModelData(PATH, 0, len(itos), trn_dl, val_dl, bs=bs, bptt=bptt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:19:22.398749Z",
     "start_time": "2019-05-21T15:19:22.341152Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# overfitting - increase multiplier (0.7)\n",
    "# underfitting - decrease multiplier (0.7)\n",
    "drops = np.array([0.25, 0.1, 0.2, 0.02, 0.15])*0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:19:22.640758Z",
     "start_time": "2019-05-21T15:19:22.402132Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "em_sz,nh,nl = 400,1150,3\n",
    "learner= md.get_model(opt_fn, em_sz, nh, nl, \n",
    "    dropouti=drops[0], dropout=drops[1], wdrop=drops[2], dropoute=drops[3], dropouth=drops[4])\n",
    "\n",
    "learner.metrics = [accuracy]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T03:17:48.528046Z",
     "start_time": "2019-05-20T19:07:04.843904Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20d010d77aaa42f4be75155537dc577e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=2), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss   accuracy                             \n",
      "    0      1.134826   1.115255   0.663817  \n",
      "    1      1.070876   1.001363   0.696245                             \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[1.0013631663798634, 0.6962449800136478]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr=1e-3\n",
    "learner.fit(lr, 1, wds=wd, use_clr=(10,2), cycle_len=2, best_save_name='wiki103_lm')\n",
    "\n",
    "# wikitext2  15cycle(20,10), 1e-3\n",
    "# 1.18086    1.189847   0.646    fastai LM    'wiki2_lm'\n",
    "\n",
    "# wikitext103\n",
    "# 1.070876   1.001363   0.696245"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:20:01.098761Z",
     "start_time": "2019-05-21T15:20:00.961663Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.load('wiki103_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Transformer Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:12.665577Z",
     "start_time": "2019-05-24T21:11:12.609206Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LMLoader():\n",
    "    \"\"\" Returns a language model iterator that iterates through batches that are of length N(bptt,5)\n",
    "    The first batch returned is always bptt+25; the max possible width.  This is done because of the way that pytorch\n",
    "    allocates cuda memory in order to prevent multiple buffers from being created as the batch width grows.\n",
    "    \"\"\"\n",
    "    def __init__(self, nums, vocab_len, bs, bptt):\n",
    "        self.bs,self.bptt = bs,bptt\n",
    "        self.vocab_len = vocab_len\n",
    "        self.data = self.batchify(nums)\n",
    "        self.i,self.iter = 0,0\n",
    "        self.n = len(self.data)\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.i,self.iter = 0,0\n",
    "        while self.i < self.n-1 and self.iter<len(self):\n",
    "            if self.i == 0:\n",
    "                seq_len = self.bptt + 5 * 5\n",
    "            else:\n",
    "                bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2.\n",
    "                seq_len = max(5, int(np.random.normal(bptt, 5)))\n",
    "            res = self.get_pair(self.i, seq_len)\n",
    "            self.i += seq_len\n",
    "            self.iter += 1\n",
    "            yield res\n",
    "\n",
    "    def __len__(self): return self.n // self.bptt - 1\n",
    "    \n",
    "    def batchify(self, data):\n",
    "        nb = data.shape[0] // self.bs        # integer division into batches\n",
    "        data = np.array(data[:nb*self.bs])   # remove remainder\n",
    "        data = data.reshape(self.bs, -1).T   # reshape and transpose\n",
    "        return T(data)                       # output a tensor\n",
    "        \n",
    "    def get_pair(self, i, seq_len):\n",
    "        source = self.data\n",
    "        seq_len = min(seq_len, len(source) - 1 - i)\n",
    "        return source[i:i+seq_len].transpose(1,0), source[i+1:i+seq_len+1].transpose(1,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:04:43.022954Z",
     "start_time": "2019-05-24T21:04:42.976251Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bs, bptt = 50, 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:17.518720Z",
     "start_time": "2019-05-24T21:11:17.335120Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "trn_dl = LMLoader(trn_idx, len(itos), bs, bptt)\n",
    "val_dl = LMLoader(val_idx, len(itos), bs, bptt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:17.719450Z",
     "start_time": "2019-05-24T21:11:17.665230Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "md = LanguageModelData(PATH, 0, len(itos), trn_dl, val_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:18.105187Z",
     "start_time": "2019-05-24T21:11:18.053154Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 125]), torch.Size([50, 125]))"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ii = iter(md.trn_dl)\n",
    "x,y = next(ii)\n",
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:21.503840Z",
     "start_time": "2019-05-24T21:11:21.454266Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'hitecture competition for its design, no construction took place. Instead, the parcel was turned into a parking lot, which it'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(x[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-24T21:11:24.760667Z",
     "start_time": "2019-05-24T21:11:24.711881Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'itecture competition for its design, no construction took place. Instead, the parcel was turned into a parking lot, which it '"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(y[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Denoising AutoEncoder LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:44.783460Z",
     "start_time": "2019-05-15T15:24:44.746883Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DenoisingAutoEncoderLoader():\n",
    "    \"\"\" Returns a language model iterator that iterates through batches that are of length N(bptt,5)\n",
    "    The first batch returned is always bptt+25; the max possible width.  This is done because of the way that pytorch\n",
    "    allocates cuda memory in order to prevent multiple buffers from being created as the batch width grows.\n",
    "    \"\"\"\n",
    "    def __init__(self, nums, vocab_len, bs, bptt):\n",
    "        self.bs,self.bptt = bs,bptt\n",
    "        self.vocab_len = vocab_len\n",
    "        self.data = self.batchify(nums)\n",
    "        self.i,self.iter = 0,0\n",
    "        self.n = len(self.data)\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.i,self.iter = 0,0\n",
    "        while self.i < self.n-1 and self.iter<len(self):\n",
    "            if self.i == 0:\n",
    "                seq_len = self.bptt + 5 * 5\n",
    "            else:\n",
    "                bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2.\n",
    "                seq_len = max(5, int(np.random.normal(bptt, 5)))\n",
    "            res = self.get_pair(self.i, seq_len)\n",
    "            self.i += seq_len\n",
    "            self.iter += 1\n",
    "            yield res\n",
    "\n",
    "    def __len__(self): return self.n // self.bptt - 1\n",
    "    \n",
    "    def batchify(self, data):\n",
    "        nb = data.shape[0] // self.bs        # integer division into batches\n",
    "        data = np.array(data[:nb*self.bs])   # remove remainder\n",
    "        data = data.reshape(self.bs, -1).T   # reshape and transpose\n",
    "        return data                          # output a tensor\n",
    "\n",
    "    def get_pair(self, i, seq_len):\n",
    "        seq_len = min(seq_len, self.n - 1 - i)\n",
    "        arr = self.data[i:i+seq_len]    # (~125,50)\n",
    "\n",
    "        # seq: (~bptt, bs)\n",
    "        # need different scramble for each bs\n",
    "        res,src = [],[]\n",
    "        for b in range(arr.shape[1]):\n",
    "            source = arr[:,b]\n",
    "            \n",
    "            # remove partial words from beginning and end\n",
    "            spaces = np.where(source==1)[0]\n",
    "            source = source[spaces[0]+1:spaces[-1]]\n",
    "            src.append(source)\n",
    "            \n",
    "            seq = source.copy()\n",
    "            \n",
    "            # scramble\n",
    "            num = random.randint(0, math.floor(self.bptt * 0.15))\n",
    "            idxs = np.random.randint(len(seq), size=num)\n",
    "            for i in idxs:\n",
    "                if seq[i] not in [1,2]:   # don't modify ' ' or '\\n'\n",
    "                    prob = random.random()\n",
    "                    if prob < 0.3:           #replacement\n",
    "                        seq[i] = random.randrange(56,82) #self.vocab_len)  # only lowercase letters\n",
    "                    elif 0.3 <= prob < 0.6:  #removal\n",
    "                        seq[i] = 0\n",
    "                    elif 0.6 <= prob < 0.9:  #addition\n",
    "                        seq = np.insert(seq, i, random.randrange(56,82)) #self.vocab_len))\n",
    "            mask = seq.nonzero()\n",
    "            res.append(seq[mask])\n",
    "        \n",
    "        # convert res to 2d numpy array w/ zero padding\n",
    "        out_res = np.zeros([len(res), len(max(res,key = lambda x: len(x)))], dtype=int)\n",
    "        for i,j in enumerate(res):\n",
    "            out_res[i][0:len(j)] = j\n",
    "            \n",
    "        out_src = np.zeros([len(src), len(max(src,key=lambda x: len(x)))], dtype=int)\n",
    "        for i,j in enumerate(src):\n",
    "            out_src[i][0:len(j)] = j\n",
    "            \n",
    "        return T(out_res), T(out_src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:44.809408Z",
     "start_time": "2019-05-15T15:24:44.785311Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bs, bptt = 50, 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:45.397967Z",
     "start_time": "2019-05-15T15:24:44.811472Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "trn_dl = DenoisingAutoEncoderLoader(trn_idx, len(itos), bs, bptt)\n",
    "val_dl = DenoisingAutoEncoderLoader(val_idx, len(itos), bs, bptt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:45.436628Z",
     "start_time": "2019-05-15T15:24:45.400472Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "md = LanguageModelData(PATH, 0, len(itos), trn_dl, val_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.579264Z",
     "start_time": "2019-05-15T15:24:45.439069Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([50, 125]), torch.Size([50, 122]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ii = iter(md.trn_dl)\n",
    "x,y = next(ii)\n",
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.607243Z",
     "start_time": "2019-05-15T15:24:49.580938Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"it had dnissipted over open watrs. \\n Preparations and impact \\n Upon the cyclone's formatyon, the Bureau of Meteorflogy\""
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(x[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.631330Z",
     "start_time": "2019-05-15T15:24:49.608961Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"it had dissipated over open waters. \\n Preparations and impact \\n Upon the cyclone's formation, the Bureau of Meteorology\""
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(y[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helpers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss and Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.654471Z",
     "start_time": "2019-05-15T15:24:49.632995Z"
    }
   },
   "outputs": [],
   "source": [
    "def loss_prep(input, target):\n",
    "    \"equalize input/target sl; combine bs/sl dimensions\"\n",
    "    bs,tsl = target.shape\n",
    "    _ ,sl,vocab = input.shape\n",
    "        \n",
    "    # F.pad( front,back for dimensions: 1,0,2 )\n",
    "    if sl>tsl: target = F.pad(target, (0,sl-tsl))\n",
    "        \n",
    "    # this should only be used when testing for small seq_lens\n",
    "    # if tsl>sl: target = target[:,:sl]\n",
    "    \n",
    "    if tsl>sl: input = F.pad(input, (0,0,0,tsl-sl))\n",
    "    # not ideal => adds 82 logits all 0s...\n",
    "        \n",
    "    targ = target.contiguous().view(-1).long()\n",
    "    pred = input.contiguous().view(-1, vocab)\n",
    "    return pred, targ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.675461Z",
     "start_time": "2019-05-15T15:24:49.655871Z"
    }
   },
   "outputs": [],
   "source": [
    "# def loss_prep(input, target):\n",
    "#     \"equalize input/target sl; combine bs/sl dimensions\"\n",
    "# #     bs,tsl = target.shape\n",
    "#     _,sl,vocab = input.shape\n",
    "        \n",
    "# #     # F.pad( front,back for dimensions: 1,0,2 )\n",
    "# #     if sl>tsl: target = F.pad(target, (0,sl-tsl))\n",
    "# #     if tsl>sl: target = target[:,:sl]\n",
    "# # #     if tsl>sl: input = F.pad(input, (0,0,0,0,0,tsl-sl))\n",
    "        \n",
    "#     targ = target.contiguous().view(-1).long()\n",
    "#     pred = input.contiguous().view(-1, vocab)\n",
    "#     return pred, targ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.697920Z",
     "start_time": "2019-05-15T15:24:49.677218Z"
    }
   },
   "outputs": [],
   "source": [
    "class LabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(LabelSmoothing, self).__init__()\n",
    "        self.confidence = 1.0 - smoothing\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, target):\n",
    "        pred,targ = loss_prep(pred, target)\n",
    "        pred = F.log_softmax(pred, dim=-1)  # need this for KLDivLoss\n",
    "        true_dist = pred.data.clone()\n",
    "        true_dist.fill_(self.smoothing / pred.size(1))                  # fill with 0.0012\n",
    "        true_dist.scatter_(1, targ.data.unsqueeze(1), self.confidence)  # [0.0012, 0.0012, 0.90, 0.0012]\n",
    "        return F.kl_div(pred, true_dist, reduction='sum')/bs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stepper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.771076Z",
     "start_time": "2019-05-15T15:24:49.750918Z"
    }
   },
   "outputs": [],
   "source": [
    "def subsequent_mask(size):\n",
    "    attn_shape = torch.ones((size,size), dtype=torch.int, device=device)\n",
    "    mask = torch.tril(attn_shape).unsqueeze(0)\n",
    "    return mask\n",
    "\n",
    "def make_tgt_mask(tgt, pad=0):\n",
    "    \"Create a mask to hide padding and future words.\"\n",
    "    tgt_mask = (tgt != pad).unsqueeze(-2)\n",
    "    tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))\n",
    "    return tgt_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.792638Z",
     "start_time": "2019-05-15T15:24:49.772550Z"
    }
   },
   "outputs": [],
   "source": [
    "def rshift(tgt, token=1):\n",
    "    \"Shift y to the right by prepending token\"\n",
    "    return torch.cat((torch.ones((tgt.size(0),token), device=device, dtype=torch.long), tgt[:,:-1]), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.816619Z",
     "start_time": "2019-05-15T15:24:49.794354Z"
    }
   },
   "outputs": [],
   "source": [
    "class TfmrStepper(Stepper):\n",
    "    def step(self, xs, y, epoch):\n",
    "        xtra = []\n",
    "        shifted_y = rshift(y).long()\n",
    "        tgt_mask = subsequent_mask(shifted_y.size(-1)) #make_tgt_mask(shifted_y)\n",
    "        output = self.m(*xs, shifted_y, tgt_mask)\n",
    "        \n",
    "        if isinstance(output,tuple): output,*xtra = output\n",
    "        self.opt.zero_grad()\n",
    "        loss = raw_loss = self.crit(output, y)\n",
    "        if self.reg_fn: loss = self.reg_fn(output, xtra, raw_loss)\n",
    "        loss.backward()        \n",
    "        if self.clip:   # Gradient clipping\n",
    "            nn.utils.clip_grad_norm_(trainable_params_(self.m), self.clip)\n",
    "        self.opt.step()\n",
    "        return raw_loss.item()\n",
    "    \n",
    "    def evaluate(self, xs, y):\n",
    "        shifted_y = rshift(y).long()\n",
    "        tgt_mask = subsequent_mask(shifted_y.size(-1)) #make_tgt_mask(shifted_y)\n",
    "        preds = self.m(*xs, shifted_y, tgt_mask)\n",
    "        if isinstance(preds,tuple): preds=preds[0]\n",
    "        return preds, self.crit(preds, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Transformer Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.839755Z",
     "start_time": "2019-05-15T15:24:49.818586Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# similar to batchnorm but on a layer level\n",
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.a_2 = nn.Parameter(torch.ones(features))\n",
    "        self.b_2 = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.862180Z",
     "start_time": "2019-05-15T15:24:49.841343Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SublayerConnection(nn.Module):\n",
    "    \"A residual connection followed by a layer norm.  Note: (for code simplicity) norm is first.\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super(SublayerConnection, self).__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        return x + self.dropout(sublayer(self.norm(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.884193Z",
     "start_time": "2019-05-15T15:24:49.863848Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def clones(module, N):\n",
    "    \"Produce N identical layers.\"\n",
    "    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.911682Z",
     "start_time": "2019-05-15T15:24:49.886455Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.937474Z",
     "start_time": "2019-05-15T15:24:49.913822Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    \"Encoder: self-attn and feed forward\"\n",
    "    def __init__(self, size, self_attn, feed_forward, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x))\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.963288Z",
     "start_time": "2019-05-15T15:24:49.940123Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, layer, N):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src, tgt_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:49.990886Z",
     "start_time": "2019-05-15T15:24:49.965384Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    \"Decoder: self-attn, src-attn, and feed forward\"\n",
    "    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 3)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))\n",
    "        x = self.sublayer[1](x, lambda x: self.src_attn(x, src, src))\n",
    "        return self.sublayer[2](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.019117Z",
     "start_time": "2019-05-15T15:24:49.993071Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    \"Compute 'Scaled Dot Product Attention'\"\n",
    "    depth = query.size(-1)\n",
    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(depth)\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -1e9)    \n",
    "    p_attn = F.softmax(scores, dim=-1)\n",
    "    if dropout is not None:\n",
    "        p_attn = dropout(p_attn)\n",
    "    return torch.matmul(p_attn, value), p_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.079752Z",
     "start_time": "2019-05-15T15:24:50.021244Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SingleHeadedAttention(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2):\n",
    "        super(SingleHeadedAttention, self).__init__()\n",
    "        self.linears = clones(nn.Linear(d_model, d_model), 4)\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        query, key, value = [l(x) for l, x in zip(self.linears, (query, key, value))]\n",
    "        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)\n",
    "        return self.linears[-1](x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.110658Z",
     "start_time": "2019-05-15T15:24:50.082961Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2, mult=4):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_model*mult)\n",
    "        self.w_2 = nn.Linear(d_model*mult, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.dropout(F.relu(self.w_1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.142419Z",
     "start_time": "2019-05-15T15:24:50.112798Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2, max_len=2000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "    \n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0.0, max_len).unsqueeze(1)\n",
    "        log_increment = math.log(1e4) / d_model\n",
    "        div_term = torch.exp(torch.arange(0.0, d_model, 2) * -log_increment)  \n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe.unsqueeze_(0)\n",
    "        self.register_buffer('pe', pe)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.180562Z",
     "start_time": "2019-05-15T15:24:50.144897Z"
    }
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_embed(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.210059Z",
     "start_time": "2019-05-15T15:24:50.182907Z"
    }
   },
   "outputs": [],
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super(Embeddings, self).__init__()\n",
    "        self.lut = nn.Embedding(vocab, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.lut(x) * 18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.241659Z",
     "start_time": "2019-05-15T15:24:50.212493Z"
    }
   },
   "outputs": [],
   "source": [
    "def make_language_model(vocab, d_model=512, N=4, drop=0.2):\n",
    "    c = copy.deepcopy\n",
    "    attn = SingleHeadedAttention(d_model)\n",
    "#     attn = MultiHeadedAttention(d_model, 8)\n",
    "    ff = PositionwiseFeedForward(d_model, drop)\n",
    "    pos_enc = PositionalEncoding(d_model, drop, 2000)\n",
    "    \n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drop), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drop), N),\n",
    "        nn.Sequential(nn.Embedding(vocab, d_model), pos_enc),\n",
    "        nn.Sequential(nn.Embedding(vocab, d_model), pos_enc),\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.274480Z",
     "start_time": "2019-05-15T15:24:50.243912Z"
    }
   },
   "outputs": [],
   "source": [
    "# denoising auto-encoder\n",
    "# Want to predict entire output including masked words\n",
    "\n",
    "class BetterSpeller(nn.Module):\n",
    "    def __init__(self, lm):\n",
    "        super(BetterSpeller, self).__init__()\n",
    "        self.lm = lm\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None):\n",
    "        return self.lm.generate(self.lm(src, tgt, tgt_mask))\n",
    "\n",
    "    def greedy_decode(self, src):\n",
    "        with torch.no_grad():\n",
    "            feats = self.lm.encode(src)\n",
    "            bs,sl = src.shape\n",
    "            tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "            res = []                \n",
    "            for i in tqdm(range(sl+5)):\n",
    "                mask = subsequent_mask(tgt.size(-1))\n",
    "                dec_outs = self.lm.decode(feats, Variable(tgt), Variable(mask))\n",
    "                prob = self.lm.generate(dec_outs[:,-1])\n",
    "                res.append(prob)\n",
    "                pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                if (pred==0).all(): break\n",
    "                tgt = torch.cat([tgt,pred], dim=-1)\n",
    "            out = torch.stack(res).transpose(1,0).contiguous()\n",
    "            return out      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:24:50.828707Z",
     "start_time": "2019-05-15T15:24:50.276969Z"
    }
   },
   "outputs": [],
   "source": [
    "d_model = 512\n",
    "lm = make_language_model(len(itos), d_model)\n",
    "net = BetterSpeller(lm)\n",
    "\n",
    "wd=1e-7\n",
    "opt_fn = partial(optim.Adam, betas=(0.8, 0.99))\n",
    "\n",
    "learn = Learner(md, BasicModel(to_gpu(net)), opt_fn=opt_fn)\n",
    "\n",
    "learn.clip = 0.25\n",
    "learn.crit = LabelSmoothing(smoothing=0.1)\n",
    "learn.metrics = [char_error_rate]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-14T18:40:59.896282Z",
     "start_time": "2019-05-14T18:40:43.093671Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find(stepper=TfmrStepper)\n",
    "learn.sched.plot(n_skip=0, n_skip_end=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T21:51:13.563832Z",
     "start_time": "2019-05-15T15:35:47.253607Z"
    },
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db63877360e84bd394b65eeb3a8f14c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=1), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss   char_error_rate                     \n",
      "    0      7.036011   7.017283   0.022698  \n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[7.017283363079806, 0.022697553340128872]"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr=1e-3\n",
    "learn.fit(lr, 1, wds=wd, use_clr=(20,10), cycle_len=1, stepper=TfmrStepper, best_save_name='better_speller_103')\n",
    "\n",
    "# wikitext2\n",
    "# 39.696276  38.299323  0.549255   0.535159     15cycles\n",
    "\n",
    "# 6.512353   5.655003   0.035785   tfmr  15cycles(20,10)   'better_speller'\n",
    "\n",
    "\n",
    "# wikitext103\n",
    "# 10.755489  9.274058   0.029507   tfmr  1cycle(20,10)   'LM_103'\n",
    "# 7.036011   7.017283   0.022698   2nd cycle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:01.388980Z",
     "start_time": "2019-05-15T22:04:01.046093Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('LM_103')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T15:27:22.501520Z",
     "start_time": "2019-05-15T15:27:22.115362Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('LM_103')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:09.186007Z",
     "start_time": "2019-05-15T22:04:09.150189Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(iter(md.val_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:13.428021Z",
     "start_time": "2019-05-15T22:04:10.454976Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████| 129/129 [00:02<00:00, 44.02it/s]\n"
     ]
    }
   ],
   "source": [
    "learn.model.eval()\n",
    "preds = learn.model.greedy_decode(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:14.742677Z",
     "start_time": "2019-05-15T22:04:14.712440Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "probs = torch.argmax(preds, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:53.900445Z",
     "start_time": "2019-05-15T22:04:53.876552Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:54.180709Z",
     "start_time": "2019-05-15T22:04:54.154271Z"
    },
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ater eading boeks about Nazi physician Josef Mengele while on tour with the band: \"I remember stopping someplace'"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(x[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:54.457486Z",
     "start_time": "2019-05-15T22:04:54.429238Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'after leading books about Nazi physician Josef Mengele while on tour with the band: \"I remember stopping someplace'"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(probs[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-15T22:04:54.839428Z",
     "start_time": "2019-05-15T22:04:54.807857Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'after reading books about Nazi physician Josef Mengele while on tour with the band: \"I remember stopping someplace'"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_label_text(y[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-13T17:17:54.817096Z",
     "start_time": "2019-03-13T17:15:47.564877Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd0be332c4f24fcba4b03ce3213eacb1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Epoch', max=10), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch      trn_loss   val_loss   char_error_rate accuracy   \n",
      "    0      43.857646  42.990693  0.627222   0.484561  \n",
      "    1      40.316483  43.975653  0.621577   0.498255        \n",
      "    2      40.887276  41.461205  0.623438   0.497235        \n",
      "    3      39.405278  41.622659  0.617584   0.502438        \n",
      "    4      40.867689  40.411398  0.625081   0.497289        \n",
      "    5      39.15356   41.386987  0.618365   0.500589        \n",
      "    6      39.567341  38.565258  0.62212    0.499334        \n",
      "    7      40.321067  40.285084  0.621228   0.500227        \n",
      "    8      39.070562  42.125116  0.619125   0.50083         \n",
      "    9      39.798026  40.925197  0.617509   0.502518        \n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[40.92519678009881, 0.617509345910516, 0.5025176637702518]"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lr=2e-5\n",
    "learn.fit(lr, 1, wds=wd, use_clr=(20,10), cycle_len=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-08T19:46:50.881388Z",
     "start_time": "2019-03-08T19:46:50.734519Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = learn.model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-08T19:46:51.535729Z",
     "start_time": "2019-03-08T19:46:51.443383Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6253905094664887"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "char_error_rate(preds,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-03-08T19:46:52.340777Z",
     "start_time": "2019-03-08T19:46:52.299933Z"
    },
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.4497, device='cuda:0')"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy(preds,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:26:53.167479Z",
     "start_time": "2019-05-21T15:26:53.117805Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learner.model.eval()\n",
    "learner.model.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:42:05.154896Z",
     "start_time": "2019-05-21T15:42:05.101927Z"
    }
   },
   "outputs": [],
   "source": [
    "def next_with_creativity(preds, k=5, thresh=.05):\n",
    "    probs, idxs = torch.topk(F.softmax(preds, dim=-1), k, dim=-1)\n",
    "    d = {itos[k]: round(v.item(), 3) for k,v in zip(idxs,probs)}\n",
    "    print(d)\n",
    "    \n",
    "    seq = np.array([], dtype=np.long)\n",
    "    for p,i in zip(probs,idxs):\n",
    "        num = int(p * 100)\n",
    "        seq = np.append(seq, [i.item()] * num)\n",
    "    \n",
    "    return random.choice(seq.flatten())\n",
    "    \n",
    "#     return{k:v if v>=thresh else None for k,v in d}\n",
    "#     mask = [probs >= thresh] \n",
    "#     m_probs, m_idxs = probs[mask], idxs[mask]\n",
    "    \n",
    "#     if len(m_idxs) > 0:\n",
    "#         # simple weighted choice\n",
    "#         seq = \n",
    "#         random.choice(seq)\n",
    "#         idx = random.randint(0,len(m_idxs))\n",
    "#         return m_idxs[idx]\n",
    "#     else:\n",
    "#         return idxs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:43:16.123308Z",
     "start_time": "2019-05-21T15:43:16.072461Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_next(inp):\n",
    "    idxs = T(np.array([stoi[c] for c in inp])).unsqueeze(0)\n",
    "    p = learner.model(Variable(idxs))\n",
    "#     i = torch.argmax(p[0][-1], dim=-1)\n",
    "#     i = torch.multinomial(p[0].exp(), 1)[-1]\n",
    "    i = next_with_creativity(p[0][-1])\n",
    "    return itos[i.item()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T16:11:51.509729Z",
     "start_time": "2019-05-21T16:11:51.449377Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{' ': 0.196, 'n': 0.121, 's': 0.112, 'l': 0.078, 'd': 0.057}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'d'"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_next('whe')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:43:32.728425Z",
     "start_time": "2019-05-21T15:43:32.677257Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_next_n(inp, n):\n",
    "    res = inp\n",
    "    for i in range(n):\n",
    "        c = get_next(res)\n",
    "        res += c\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-05-21T15:43:34.240836Z",
     "start_time": "2019-05-21T15:43:34.125066Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'a': 0.296, 'e': 0.189, 'u': 0.12, 'o': 0.103, ' ': 0.09}\n",
      "{'n': 0.102, 'b': 0.081, 'r': 0.075, 's': 0.056, 'm': 0.056}\n",
      "{'a': 0.17, 'o': 0.121, 'e': 0.115, ' ': 0.092, 'i': 0.086}\n",
      "{'a': 0.078, 'w': 0.051, 'i': 0.048, 'o': 0.046, ' ': 0.044}\n",
      "{'n': 0.039, 'w': 0.039, 'r': 0.031, 'u': 0.028, 'p': 0.027}\n",
      "{'i': 0.251, 'a': 0.197, 'e': 0.162, 'h': 0.153, 'o': 0.126}\n",
      "{'n': 0.058, ' ': 0.049, 'm': 0.039, 'r': 0.039, 'l': 0.035}\n",
      "{'a': 0.078, 'w': 0.051, 'i': 0.048, 'o': 0.046, ' ': 0.044}\n",
      "{'n': 0.058, ' ': 0.049, 'm': 0.039, 'r': 0.039, 'l': 0.035}\n",
      "{'g': 0.053, ' ': 0.047, 'i': 0.03, 'e': 0.03, 'a': 0.029}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'thum owa ana'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_next_n('th', 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  },
  "notify_time": "10",
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}