1.0--LM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prelims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# from fastai import *\n",
    "from fastai.text import *\n",
    "from fastai.callbacks.tracker import *\n",
    "import pdb\n",
    "import textwrap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "PATH = Path('data/IAM_handwriting')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Setup Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def remove_newline_spaces(x):\n",
    "    x = x.replace(' \\n','\\n').replace('\\n ','\\n')\n",
    "    return re.sub(r'(\\n)+','\\n',x)\n",
    "\n",
    "def remove_equals(x):\n",
    "    return x.replace(' =', '').replace('= ', '')\n",
    "\n",
    "# convert spaced out \" strings \" to \"strings\"\n",
    "def despace_quotes(m):\n",
    "    m = m.group(0)   # entire matched string\n",
    "    return m.replace('\" ','\"').replace(' \"','\"')\n",
    "\n",
    "def cleanup(x):\n",
    "    x = fix_html(x)\n",
    "    x = remove_newline_spaces(x)\n",
    "    x = remove_equals(x)\n",
    "    x = x.replace( \" \\'\", \"\\'\").replace(' ,', ',').replace(' .', '.').replace(' :', ':').replace(\n",
    "        ' ;', ';').replace('( ', '(').replace(' )', ')').replace('[ ', '[').replace(' ]', ']')\n",
    "    x = re.sub(r'\\\"(.+?)\\\"', despace_quotes, x)\n",
    "    x = re.sub(r'[^\\x00-\\x7F]+','', x)   # remove all non ascii characters\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## WikiText"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# wiki_path = Path('data/wikitext/wikitext-2-raw')\n",
    "wiki_path = Path('data/wikitext/wikitext-103-raw')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "with open(wiki_path/'wiki.train.raw') as file:  \n",
    "    trn = file.read()\n",
    "with open(wiki_path/'wiki.valid.raw') as file:  \n",
    "    val = file.read()\n",
    "with open(wiki_path/'wiki.test.raw') as file:  \n",
    "    tst = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "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"
    },
    "hidden": true
   },
   "source": [
    "### clean text and save CSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full = trn + val + tst\n",
    "# len(trn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full = cleanup(full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full[0:5000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lines = full.split('\\n')\n",
    "len(lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "text = [textwrap.wrap(line, 700) for line in lines]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "flat_list = []\n",
    "for sublist in text:\n",
    "    for item in sublist:\n",
    "        flat_list.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(flat_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# add '\\n' characters\n",
    "text = [textwrap.fill(t, width=60) for t in flat_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'text': text})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# remove small lines\n",
    "df['char_len'] = df.text.map(len)\n",
    "df = df[df['char_len'] > 15]\n",
    "df.sort_values('char_len', inplace=True)\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df['labels'] = [0] * len(df)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'wiki103.csv'\n",
    "# df = pd.read_csv(PATH/CSV)\n",
    "df.to_csv(PATH/CSV, columns=['labels', 'text'], index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## IMDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb_path = untar_data(URLs.IMDB)\n",
    "imdb_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'texts.csv'\n",
    "df = pd.read_csv(imdb_path/CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df['text'] = df.text.apply(lambda x: cleanup(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# this also handles \\n\n",
    "text = [textwrap.wrap(line, 1000) for line in df.text.values]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "flat_list = []\n",
    "for sublist in text:\n",
    "    for item in sublist:\n",
    "        flat_list.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(flat_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'text': flat_list})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# remove small lines\n",
    "df['char_len'] = df.text.map(len)\n",
    "df = df[df['char_len'] > 15]\n",
    "df.sort_values('char_len', inplace=True)\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df['labels'] = [0] * len(df)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'imdb.csv'\n",
    "# imdb = pd.read_csv(PATH/CSV)\n",
    "df.to_csv(PATH/CSV, columns=['labels', 'text'], index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Combo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full = pd.concat([wiki, imdb], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full.sort_values('char_len', inplace=True)\n",
    "full.reset_index(inplace=True, drop=True)\n",
    "full.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full['labels'] = [0] * len(full)\n",
    "full.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "PATH = Path('data/IAM_handwriting')\n",
    "full.to_csv(PATH/'wiki103_imdb.csv', columns=['labels', 'text'], index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### create csv from folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for directory in os.listdir(PATH):\n",
    "    if os.path.isdir(os.path.join(PATH, directory)):\n",
    "        print(directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def add_to_df(path):\n",
    "    df = pd.DataFrame()\n",
    "    for file in os.listdir(path):\n",
    "        with open(path/file,'r') as txt:\n",
    "            df = df.append({'text': txt.read()}, ignore_index=True)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tst_neg = add_to_df(PATH/'test/neg')\n",
    "tst_pos = add_to_df(PATH/'test/pos')\n",
    "trn_neg = add_to_df(PATH/'train/neg')\n",
    "trn_pos = add_to_df(PATH/'train/neg')\n",
    "unsup = add_to_df(PATH/'unsup')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb = pd.concat([tst_neg, tst_pos, trn_neg, trn_pos, unsup], ignore_index=True)\n",
    "len(imdb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb.to_csv(PATH/'texts.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Huggingface BeRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from transformers import DistilBertTokenizer, DistilBertForMaskedLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Original Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "vocab_len = len(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FastAiBertTokenizer(BaseTokenizer): \n",
    "    def __init__(self, tokenizer, max_seq_len:int=300, **kwargs): \n",
    "        self._pretrained_tokenizer = tokenizer \n",
    "        self.max_seq_len = max_seq_len \n",
    "    def __call__(self, *args, **kwargs):  return self \n",
    "    def tokenizer(self, t:str) -> List[str]: \n",
    "        return [\"[CLS]\"] + self._pretrained_tokenizer.tokenize(t)[:self.max_seq_len - 2] + [\"[SEP]\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "f_tokenizer = Tokenizer(tok_func=FastAiBertTokenizer(tokenizer, max_seq_len=300),\n",
    "                        pre_rules=[], post_rules=[])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "f_vocab = Vocab(list(tokenizer.vocab.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'wiki2.csv' #'wiki103_imdb.csv' #'wiki103.csv' #'wiki2.csv'\n",
    "data = TextLMDataBunch.from_csv(PATH, CSV, tokenizer=f_tokenizer, vocab=f_vocab, include_bos=False, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data w/ added tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bert_tok = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "num_added_tokens = bert_tok.add_tokens(['\\n',' ','[UP]','[MAJ]'])\n",
    "vocab_len = len(bert_tok) #+ num_added_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def add_cap_tokens(text):  # before encode\n",
    "    re_caps = re.compile(r'[A-Z]+')\n",
    "    return re_caps.sub(_replace_caps, text)\n",
    "    \n",
    "def _replace_caps(m):\n",
    "    tok = '[UP]' if m.end()-m.start() > 1 else '[MAJ]'\n",
    "    return tok + m.group().lower()\n",
    "\n",
    "def remove_cap_tokens(text):  # after decode\n",
    "    text = re.sub(r'\\[UP\\]\\w+', lambda m: m.group()[4:].upper(), text)  #cap entire word\n",
    "    text = re.sub(r'\\[MAJ\\]\\w?', lambda m: m.group()[5:].upper(), text) #cap first letter\n",
    "    return text\n",
    "\n",
    "def remove_special_toks(text):\n",
    "    text = re.sub(r'\\[CLS\\]\\s*', '', text)  #[CLS] (w/ following whitespace)\n",
    "    text = re.sub(r'\\s*\\[SEP\\]', '', text)  #[SEP] (w/ preceding whitespace)\n",
    "    return text\n",
    "\n",
    "def remove_wordpiece_toks(text):\n",
    "    return re.sub(r'##', '', text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BertTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[str]: return [\"[CLS]\"] + bert_tok.tokenize(t) + [\"[SEP]\"]\n",
    "\n",
    "class BertVocab(Vocab):\n",
    "    def __init__(self):\n",
    "        self.itos = list(bert_tok.vocab.keys()) + ['\\n',' ','[UP]','[MAJ]']\n",
    "        self.stoi = collections.defaultdict(lambda: 100, {v:k for k,v in enumerate(self.itos)})\n",
    "\n",
    "    def textify(self, nums:Collection[int], sep=''):\n",
    "        st = sep.join([self.itos[i] for i in nums])\n",
    "        st = remove_wordpiece_toks(st)\n",
    "        st = remove_cap_tokens(st)\n",
    "#         st = remove_special_toks(st)\n",
    "        return st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "toknizr = Tokenizer(tok_func=BertTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "                    post_rules=[], special_cases=[])\n",
    "\n",
    "procs = [TokenizeProcessor(tokenizer=toknizr, include_bos=False), NumericalizeProcessor(vocab=BertVocab())]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def label_collater(samples:BatchSamples):\n",
    "    data = to_data(samples)\n",
    "    x,y = zip(*data)\n",
    "    if len(data) is 1 and y[0] is 0:  #inference only -> data.one_item()\n",
    "        y = torch.zeros(1,1).long()\n",
    "        return torch.from_numpy(x[0]), y\n",
    "    return collate(x), collate(y)\n",
    "    \n",
    "def collate(lbls):   #default: pad_idx=0\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long()  # add 1 to max_len to account for bos token\n",
    "    for i,lbl in enumerate(lbls):\n",
    "        labels[i,:len(lbl)] = torch.from_numpy(lbl)  #padding end    \n",
    "    return labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BasicTextList(TextList):\n",
    "    _bunch = DataBunch\n",
    "    _processor = [TokenizeProcessor, NumericalizeProcessor]\n",
    "    _is_lm = False\n",
    "\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.vocab=BertVocab()\n",
    "        self.pad_idx=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'clean_pg.csv' #'output.csv'\n",
    "bs = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (BasicTextList.from_csv(PATH, CSV, cols=0, processor=procs)\n",
    "        .split_by_rand_pct(valid_pct=0.10, seed=42)\n",
    "        .label_from_df(cols=0, label_cls=BasicTextList, processor=procs)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bert_mlm = DistilBertForMaskedLM.from_pretrained(\"distilbert-base-uncased\")\n",
    "bert_mlm.resize_token_embeddings(vocab_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BertAdaptor(Module):\n",
    "    def __init__(self):\n",
    "        self.model = bert_mlm\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.model(x)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MLM_Mask(LearnerCallback):\n",
    "    def __init__(self, learn:Learner, mlm_probability=0.15, mask_tok=103, vocab_len=vocab_len):\n",
    "        super().__init__(learn)\n",
    "        self.mask_tok = mask_tok\n",
    "        self.mlm_probability = mlm_probability\n",
    "        self.vocab_len = vocab_len\n",
    "        \n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        new_input,new_target = self.mask_tokens(last_input)\n",
    "        return {'last_input':new_input, 'last_target':new_target}\n",
    "    \n",
    "    def mask_tokens(self, inputs):\n",
    "        \"\"\" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. \"\"\"\n",
    "        labels = inputs.clone()\n",
    "        # We sample a few tokens in each sequence for masked-LM training\n",
    "        masked_indices = torch.bernoulli(torch.full(labels.shape, self.mlm_probability)).bool()\n",
    "        labels[~masked_indices] = -1  # We only compute loss on masked tokens\n",
    "        \n",
    "        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n",
    "        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n",
    "        inputs[indices_replaced] = self.mask_tok\n",
    "\n",
    "        # 10% of the time, we replace masked input tokens with random word\n",
    "        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n",
    "        random_words = torch.randint(self.vocab_len, labels.shape, dtype=torch.long, device=device)\n",
    "        inputs[indices_random] = random_words[indices_random]\n",
    "\n",
    "        # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n",
    "        return inputs, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def bert_acc(input:Tensor, targs:Tensor, ignore_index=-1)->Rank0Tensor:\n",
    "    mask = (targs!=ignore_index)\n",
    "    preds = input.argmax(dim=-1)\n",
    "    return (preds[mask]==targs[mask]).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "\n",
    "class CER(Callback):\n",
    "    def __init__(self, fn, ignore_index=-1):\n",
    "        super().__init__()\n",
    "        self.name = 'cer'\n",
    "        self.ignore_index = ignore_index\n",
    "        self.fn = fn\n",
    "\n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.errors, self.total = 0, 0\n",
    "    \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        mask = (last_target!=self.ignore_index)\n",
    "        last_output = last_output[mask]\n",
    "        last_target = last_target[mask]\n",
    "        error,size = cer(last_output, last_target, self.fn)\n",
    "        self.errors += error\n",
    "        self.total += size\n",
    "    \n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, self.errors/self.total)\n",
    "\n",
    "def cer(preds, targs, fn):\n",
    "    res = torch.argmax(preds, dim=-1)\n",
    "    p = fn(res)   #.replace(' ', '')\n",
    "    t = fn(targs) #.replace(' ', '')\n",
    "    return Lev.distance(t, p)/len(t), 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learner = Learner(data, BertAdaptor(), metrics=[bert_acc, CER(data.vocab.textify)],\n",
    "                  callback_fns=[MLM_Mask], loss_func=CrossEntropyFlat(ignore_index=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learner.validate(callbacks=[MLM_Mask(learner)], metrics=[bert_acc])\n",
    "# original: 3.1909459, tensor(0.4557)\n",
    "# w/ added: 6.9017134, tensor(0.0869)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learner.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 66,988,606"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.load('distilbert_mlm_few'); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "! mkdir {PATH/'distilbert'}\n",
    "learner.model.model.save_pretrained(PATH/'distilbert')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "learner.lr_find()\n",
    "learner.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learner.fit_one_cycle(5, 1e-4)\n",
    "# clean_pg.csv - 5(1e-4)\n",
    "# 1.006225\t0.938545\t0.818746\t0.251780\t00:22\n",
    "\n",
    "# clean_few.csv (pg+books), preload distilbert_mlm_xtra, 5cycle(1e-4)\n",
    "# 0.375431\t0.357671\t0.925877\t0.078226\t08:11    'distilbert_mlm_few'\n",
    "\n",
    "# 2.278534\t2.196675\t0.570621\t02:51  2cycle(1e-4) - original   'distilbert_mlm_orig'\n",
    "# 1.333142\t1.271659\t0.742742\t05:27  2cycle(1e-4) - w/ added   'distilbert_mlm_xtra'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.save('distilbert_mlm_few')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "st = \"Music Association Awards on November 10, 2010, it was released as te second single due to strong demands of radio stations on the following day. The song contains elements of country\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xb,yb = learner.data.one_item(st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.model.eval()\n",
    "probs = learner.model(xb)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = torch.argmax(probs, dim=-1)\n",
    "data.vocab.textify(preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Train on Model Text outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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, *args):\n",
    "        pred,targ = self.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\n",
    "    \n",
    "    def loss_prep(self, pred, target):\n",
    "        \"equalize input/target sl; combine bs/sl dimensions\"\n",
    "        bs,tsl = target.shape\n",
    "        _ ,sl,vocab = pred.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: pred = F.pad(pred, (0,0,0,tsl-sl))\n",
    "\n",
    "        targ = target.contiguous().view(-1).long()\n",
    "        pred = pred.contiguous().view(-1, vocab)\n",
    "        return pred, targ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "\n",
    "class CER(Callback):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.name = 'cer'\n",
    "        self.fn = fn\n",
    "\n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.errors, self.total = 0, 0\n",
    "    \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        error,size = cer(last_output, last_target, self.fn)\n",
    "        self.errors += error\n",
    "        self.total += size\n",
    "    \n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, self.errors/self.total)\n",
    "    \n",
    "def cer(preds, targs, fn):\n",
    "    bs = targs.size(0)\n",
    "    res = torch.argmax(preds, dim=-1)\n",
    "    error = 0\n",
    "    for i in range(bs):\n",
    "        p = str(fn(res[i]))\n",
    "        t = str(fn(targs[i]))\n",
    "        error += Lev.distance(t, p)/(len(t) or 1)\n",
    "    return error, bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learner = Learner(data, BertAdaptor(), metrics=[CER(data.y.reconstruct)], loss_func=LabelSmoothing())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.validate(metrics=[CER(data.y.reconstruct)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.load('distilbert_mlm_xtra')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "learner.lr_find()\n",
    "learner.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.fit_one_cycle(3, 1e-4)\n",
    "\n",
    "# 130.916901\t157.076309\t0.108203\t41:19    'distilbert_tfmr_out'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.save('distilbert_tfmr_out')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vdl = iter(data.valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(vdl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x[0],y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "mlm = MLM_Mask(learner)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "nx,ny = mlm.mask_tokens(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.model.eval()\n",
    "probs = learner.model(nx[None])[0]\n",
    "preds = torch.argmax(probs, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.reconstruct(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.reconstruct(y[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.reconstruct(preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bert_acc(probs, ny).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g = [learner.loss_func(probs, ny).item(), bert_acc(probs, ny).item()]\n",
    "print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')\n",
    "input_ids = torch.tensor(bert_tok.encode(\"Hello, my dog is cute\")).unsqueeze(0)  # Batch size 1\n",
    "outputs = bert_mlm(input_ids, masked_lm_labels=input_ids)\n",
    "loss, probs = outputs[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "st = \"Music Association Awards on November 10, 2010, it was released as the second single due to strong demands of radio stations on the following day. The song contains elements of country\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def mask_tokens(inputs):\n",
    "    mlm_probability = 0.30\n",
    "    mask_tok = 103\n",
    "    vocab_len = 30522\n",
    "    \n",
    "    labels = inputs.clone()\n",
    "    # We sample a few tokens in each sequence for masked-LM training\n",
    "    masked_indices = torch.bernoulli(torch.full(labels.shape, mlm_probability)).bool()\n",
    "    labels[~masked_indices] = -1  # We only compute loss on masked tokens\n",
    "\n",
    "    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n",
    "    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n",
    "    inputs[indices_replaced] = mask_tok\n",
    "\n",
    "    # 10% of the time, we replace masked input tokens with random word\n",
    "    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n",
    "    random_words = torch.randint(vocab_len, labels.shape, dtype=torch.long, device=device)\n",
    "    inputs[indices_random] = random_words[indices_random]\n",
    "\n",
    "    # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n",
    "    return inputs, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "input_ids = torch.tensor(tokenizer.encode(st)).unsqueeze(0)  # Batch size 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "nx,ny = mask_tokens(input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "nx,ny"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "outputs = model(nx, masked_lm_labels=ny)\n",
    "loss, probs = outputs[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = torch.argmax(probs, dim=-1)[0]\n",
    "tokenizer.decode(preds.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "probs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "input_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "F.cross_entropy(probs.view(-1,30522), input_ids.flatten())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "text = \"My computer can read language!  What a time to be alive?!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xb,yb = learner.data.one_item(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xb[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xb = xb[:,:-1]\n",
    "xb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.data.vocab.textify(xb[0], sep=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "res = learner.pred_batch(batch=(xb,yb))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "torch.multinomial(res[0], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = torch.argmax(res[0], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.data.vocab.textify(preds, sep=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def predict(self, text:str, n_words:int=1, no_unk:bool=True, temperature:float=1., min_p:float=None, sep:str=' ',\n",
    "            decoder=decode_spec_tokens):\n",
    "    \"Return `text` and the `n_words` that come after\"\n",
    "    xb,yb = self.data.one_item(text)\n",
    "    xb = xb[:,:-1]      # remove the eos/SEP token which is automatically added\n",
    "    print(xb)\n",
    "\n",
    "    \n",
    "    new_idx = []\n",
    "    for _ in range(n_words):\n",
    "        res = self.pred_batch(batch=(xb,yb))[0][-1]\n",
    "        #if len(new_idx) == 0: self.model[0].select_hidden([0])\n",
    "        if no_unk: res[self.data.vocab.stoi[UNK]] = 0.\n",
    "        if min_p is not None:\n",
    "            if (res >= min_p).float().sum() == 0:\n",
    "                warn(f\"There is no item with probability >= {min_p}, try a lower value.\")\n",
    "            else: res[res < min_p] = 0.\n",
    "        if temperature != 1.: res.pow_(1 / temperature)\n",
    "        idx = torch.multinomial(res, 1).item()\n",
    "        new_idx.append(idx)\n",
    "        xb = xb.new_tensor([idx])[None]\n",
    "    return text + sep + self.data.vocab.textify(new_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "predict(learner, text, n_words=5, sep=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
    "model = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "input_ids = torch.tensor(tokenizer.encode(\"My car can read a car!\", add_special_tokens=True)).unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "np.array(tokenizer.encode(\"My car can read a car!\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "outputs = model(input_ids, masked_lm_labels=input_ids)\n",
    "loss, prediction_scores = outputs[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = torch.argmax(prediction_scores[0], dim=-1)\n",
    "tokenizer.decode(preds.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DistilBertMLM(nn.Module):\n",
    "    def __init__(self, name='distilbert-base-uncased'):\n",
    "        super().__init__()\n",
    "        self.tokenizer = DistilBertTokenizer.from_pretrained(name)\n",
    "        self.model = DistilBertForMaskedLM.from_pretrained(name)\n",
    "        self.model.train()\n",
    "        \n",
    "    def forward(self, text):\n",
    "        input_ids = torch.tensor(self.tokenizer.encode(text, add_special_tokens=True), device=device)\n",
    "        outputs = self.model(input_ids, masked_lm_labels=input_ids)\n",
    "        return outputs[:2]   # loss, prediction_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = prediction_scores.argmax(dim=-1)[0];\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tokenizer.decode(preds.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# DataBunch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CustomVocab(Vocab):\n",
    "    def __init__(self, itos:Collection[str]):\n",
    "        self.itos = itos\n",
    "        self.stoi = collections.defaultdict(lambda: 3, {v:k for k,v in enumerate(self.itos)})\n",
    "\n",
    "    def textify(self, nums:Collection[int], sep=''):\n",
    "        return sep.join([self.itos[i] for i in nums]) if sep is not None else [self.itos[i] for i in nums]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CustomTokenizer(BaseTokenizer):\n",
    "    \"Split words but keep original spacing\"\n",
    "    def tokenizer(self, t:str) -> List[str]: \n",
    "        chars = list(t)\n",
    "        res = []\n",
    "        tok = \"\"\n",
    "        for c in chars:\n",
    "            if c.isalnum():\n",
    "                tok+=c\n",
    "            else:\n",
    "                if tok.isalnum(): res.append(tok) \n",
    "                res.append(c)\n",
    "                tok = \"\"\n",
    "        if tok.isalnum(): res.append(tok)\n",
    "        return ['xxbos'] + res + ['xxeos']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def label_text(pred, sep=''):\n",
    "    ints = to_np(pred).astype(int)\n",
    "    nonzero = ints[np.nonzero(ints)] #[:-1]  #remove eos token\n",
    "    return sep.join([itos[i] for i in nonzero])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "itos = pickle.load(open(PATH/'combo_itos_60k.pkl', 'rb'))\n",
    "itos = itos[:30000]\n",
    "vocab = CustomVocab(itos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "toknizr = Tokenizer(tok_func=CustomTokenizer, pre_rules=[rm_useless_spaces],\n",
    "                    special_cases=['xxbos','xxeos','xxmask','xxunk','xxpad','xxmaj','xxup','\\n'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### characters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def characterize(x:Collection[str]) -> Collection[str]:\n",
    "    \"Separate word tokens into letters. (Keep special modifiers: xxmaj, xxup)\"\n",
    "    res = []\n",
    "    for t in x:\n",
    "        [res.append(c) for c in list(t)]\n",
    "    return res\n",
    "\n",
    "toknizr = Tokenizer(tok_func=CustomTokenizer, pre_rules=[rm_useless_spaces],\n",
    "                    post_rules=[replace_all_caps, deal_caps, characterize],\n",
    "                    special_cases=['xxbos','xxeos','xxmask','xxunk','xxpad','xxmaj','xxup'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Databunch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'wiki2.csv' #'wiki103_imdb.csv' #'wiki103.csv' #'wiki2.csv'\n",
    "data = TextLMDataBunch.from_csv(PATH, CSV, tokenizer=toknizr, vocab=vocab, include_bos=False, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "\n",
    "class CER(Callback):\n",
    "    def __init__(self, ignore_index=-1):\n",
    "        super().__init__()\n",
    "        self.name = 'cer'\n",
    "        self.ignore_index = ignore_index\n",
    "\n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.errors, self.total = 0, 0\n",
    "    \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        mask = (last_target!=self.ignore_index)\n",
    "        last_output = last_output[mask]\n",
    "        last_target = last_target[mask]\n",
    "        error,size = cer(last_output, last_target)\n",
    "        self.errors += error\n",
    "        self.total += size\n",
    "    \n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, self.errors/self.total)\n",
    "\n",
    "def cer(preds, targs):\n",
    "    res = torch.argmax(preds, dim=-1)\n",
    "    p = label_text(res)   #.replace(' ', '')\n",
    "    t = label_text(targs) #.replace(' ', '')\n",
    "    return Lev.distance(t, p)/len(t), 1\n",
    "\n",
    "# def cer(preds, targs):\n",
    "#     bs = targs.size(0)\n",
    "#     res = torch.argmax(preds, dim=-1)\n",
    "#     error = 0\n",
    "#     for i in range(bs):\n",
    "#         p = label_text(res[i])   #.replace(' ', '')\n",
    "#         t = label_text(targs[i]) #.replace(' ', '')\n",
    "#         error += Lev.distance(t, p)/len(t)\n",
    "#     return error, bs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# v1 ULMFit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "config = dict(emb_sz=512, n_hid=1400, n_layers=3, pad_token=0, qrnn=False, bidir=False, output_p=0.2,\n",
    "              hidden_p=0.2, input_p=0.5, embed_p=0.1, weight_p=0.4, tie_weights=True, out_bias=True)\n",
    "\n",
    "learn = language_model_learner(data, AWD_LSTM, config=config, drop_mult=0.5,\n",
    "                               pretrained=False, metrics=[accuracy, CER()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 30,378,720"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# v1 Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Transformer Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "LayerNorm = partial(nn.LayerNorm, eps=1e-4)  # accomodates mixed precision training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def clones(module, N):\n",
    "    \"Produce N identical layers.\"\n",
    "    return nn.ModuleList([deepcopy(module) for _ in range(N)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "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, mask=None):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "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, mask=None):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "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, -1e4)    \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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, d_model, h=8, dropout=0.2):\n",
    "        super(MultiHeadedAttention, self).__init__()\n",
    "        assert d_model % h == 0\n",
    "        self.d_k = d_model // h        # assume d_v always equals d_k\n",
    "        self.h = h\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, q, k, v, mask=None):\n",
    "        if mask is not None: mask = mask.unsqueeze(1)\n",
    "        bs = q.size(0)\n",
    "        \n",
    "        # 1) Do all the linear projections in batch from d_model => h x d_k \n",
    "        q, k, v = [l(x).view(bs, -1, self.h, self.d_k).transpose(1,2) for l, x in zip(self.linears, (q, k, v))]\n",
    "        \n",
    "        # 2) Apply attention on all the projected vectors in batch. \n",
    "        x, self.attn = attention(q, k, v, mask=mask, dropout=self.dropout)\n",
    "        \n",
    "        # 3) \"Concat\" using a view and apply a final linear. \n",
    "        x = x.transpose(1, 2).contiguous().view(bs, -1, self.h * self.d_k)\n",
    "        return self.linears[-1](x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_model*4)\n",
    "        self.w_2 = nn.Linear(d_model*4, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.dropout(F.gelu(self.w_1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.1, 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",
    "\n",
    "        self.register_buffer('pe', pe)    #(1,max_len,d_model)\n",
    "        # registered buffers are Tensors (not Variables)\n",
    "        # not a parameter but still want in the state_dict\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.size(1)]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class TransformerLM(Module):\n",
    "    def __init__(self, vocab, d_model=512, N=4, drops=0.2, attn_type='multi', attn_heads=8):        \n",
    "        if attn_type=='multi':\n",
    "            attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "        else:\n",
    "            attn = SingleHeadedAttention(d_model)\n",
    "        ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "        self.tgt_embed = nn.Sequential(Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000))\n",
    "        self.encoder = Encoder(EncoderLayer(d_model, attn, ff, drops), N)\n",
    "        \n",
    "    def reset(self): pass\n",
    "\n",
    "    def forward(self, x):\n",
    "        bs,x_len = x.size()\n",
    "        mask = self.subsequent_mask(x_len)\n",
    "        inp = self.encoder(self.tgt_embed(x), mask=mask)\n",
    "        return ([inp],[inp]) #For the LinearDecoder\n",
    "    \n",
    "    def subsequent_mask(self, size):\n",
    "#         return torch.tril(torch.ones((size,size), device=device)).byte()[None]  # original...\n",
    "        return None #MLM\n",
    "\n",
    "        # only next output is masked\n",
    "#         upper = torch.triu(torch.ones((size,size), device=device), diagonal=1).bool()\n",
    "#         lower = torch.tril(torch.ones((size,size), device=device), diagonal=1).bool()\n",
    "#         mask = upper & lower\n",
    "#         return (~mask).byte()[None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def init_tfmr_lm(m):\n",
    "    for p in m.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "            \n",
    "#     classname = m.__class__.__name__\n",
    "#     if classname.find('Linear') != -1:\n",
    "#         if hasattr(m, 'weight') and m.weight is not None: nn.init.normal_(m.weight, 0., 0.02)\n",
    "#         if hasattr(m, 'bias') and m.bias is not None:     nn.init.constant_(m.bias, 0.)\n",
    "#     elif classname.find('LayerNorm') != -1:\n",
    "#         if hasattr(m, 'weight') and m.weight is not None: nn.init.normal_(m.weight, 1., 0.02)\n",
    "#         if hasattr(m, 'bias') and m.bias is not None:     nn.init.constant_(m.bias, 0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model=512, N=6, drops=0.2, attn_type='multi', attn_heads=8, **learn_kwargs):\n",
    "    vocab_sz = len(data.vocab.itos)\n",
    "    encoder = TransformerLM(vocab_sz, d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    decoder = LinearDecoder(vocab_sz, d_model, drops, tie_encoder=encoder.tgt_embed[0].lut, bias=True)\n",
    "    model = SequentialRNN(encoder, decoder)\n",
    "    model.apply(init_tfmr_lm)\n",
    "    return LanguageLearner(data, model, **learn_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 6, metrics=[accuracy, CER()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 12,659,808"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# TransformerXL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def init_transformer(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Linear') != -1:\n",
    "        if hasattr(m, 'weight') and m.weight is not None: nn.init.normal_(m.weight, 0., 0.02)\n",
    "        if hasattr(m, 'bias') and m.bias is not None:     nn.init.constant_(m.bias, 0.)\n",
    "    elif classname.find('LayerNorm') != -1:\n",
    "        if hasattr(m, 'weight') and m.weight is not None: nn.init.normal_(m.weight, 1., 0.02)\n",
    "        if hasattr(m, 'bias') and m.bias is not None:     nn.init.constant_(m.bias, 0.)\n",
    "    elif classname.find('TransformerXL') != -1:\n",
    "        if hasattr(m, 'u'): nn.init.normal_(m.u, 0., 0.02)\n",
    "        if hasattr(m, 'v'): nn.init.normal_(m.v, 0., 0.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## fastai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "config = dict(ctx_len=150, n_layers=6, n_heads=8, d_model=512, d_head=64, d_inner=2048, resid_p=0.1, attn_p=0.1,\n",
    "              ff_p=0.1, embed_p=0.1, output_p=0.1, bias=False, scale=True, act=Activation.GeLU, double_drop=True,\n",
    "              tie_weights=True, out_bias=True, init=init_transformer, mem_len=150, mask=True)\n",
    "\n",
    "learn = language_model_learner(data, TransformerXL, config=config, drop_mult=1,\n",
    "                               pretrained=False, metrics=[accuracy, CER()], callback_fns=[MLM_Mask])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 30,378,720"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {},
    "hidden": true
   },
   "source": [
    "## manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(Module):\n",
    "    \"Encode the position with a sinusoid.\"\n",
    "    def __init__(self, d_model:int):\n",
    "        self.register_buffer('freq', 1 / (10000 ** (torch.arange(0., d_model, 2.) / d_model)))\n",
    "\n",
    "    def forward(self, pos:Tensor):\n",
    "        inp = torch.ger(pos, self.freq)\n",
    "        enc = torch.cat([inp.sin(), inp.cos()], dim=-1)\n",
    "        return enc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FeedForward(Module):\n",
    "    def __init__(self, d_model:int, drops=0.1):\n",
    "        self.core = nn.Sequential(\n",
    "            nn.Linear(d_model, d_model*4), nn.ReLU(inplace=True), nn.Dropout(drops),\n",
    "            nn.Linear(d_model*4, d_model), nn.Dropout(drops)\n",
    "        )\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.core(x)\n",
    "        return self.ln(x + out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiHeadRelativeAttention(Module):\n",
    "    def __init__(self, n_heads:int, d_model:int, drops=0.1, bias=False):\n",
    "        d_head = d_model//n_heads\n",
    "        self.n_heads, self.d_head = n_heads, d_head\n",
    "        self.attention = nn.Linear(d_model, 3 * n_heads * d_head, bias=bias)\n",
    "        self.out = nn.Linear(n_heads * d_head, d_model, bias=bias)\n",
    "        self.drop_att,self.drop_res = nn.Dropout(drops),nn.Dropout(drops)\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        self.r_attn = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ln(x + self.drop_res(self.out(self._mhra(x, mask=mask, **kwargs))))\n",
    "\n",
    "    def _mhra(self, x:Tensor, r:Tensor=None, u:Tensor=None, v:Tensor=None, mask:Tensor=None, mem:Tensor=None):\n",
    "        #Notations from the paper:\n",
    "        #x: input, r: vector of relative distance between two elements\n",
    "        #u,v: learnable parameters of the model common between layers\n",
    "        #mask: to avoid cheating, mem: previous hidden states\n",
    "                \n",
    "        #x: [bs, sl, d_model]\n",
    "        #u/v: [n_heads, 1, d_head]\n",
    "        #r: [sl, d_model]\n",
    "        #mem: 1st:[0]; 2nd:[bs, sl, d_model]; nth: sl*i-1 up to mem_len\n",
    "        bs,x_len,seq_len = x.size(0),x.size(1),r.size(0)\n",
    "        context = x if mem is None else torch.cat([mem, x], dim=1)\n",
    "        # after 1st iteration: mem => \n",
    "        wq,wk,wv = torch.chunk(self.attention(context), 3, dim=-1)\n",
    "        wq = wq[:,-x_len:]\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        # [bs, sl, n_heads, d_head]\n",
    "        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)   #wk: transposed(-2,-1)\n",
    "        wkr = self.r_attn(r)\n",
    "        wkr = wkr.view(seq_len, self.n_heads, self.d_head)\n",
    "        wkr = wkr.permute(1,2,0)  #transposed ala wk w/out bs\n",
    "        #### compute attention score (AC is (a) + (c) and BD is (b) + (d) in the paper)\n",
    "        AC = torch.matmul(wq+u,wk)\n",
    "        BD = _line_shift(torch.matmul(wq+v, wkr))\n",
    "        attn_score = (AC + BD).div_(self.d_head ** 0.5)  #scale\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask, -float('inf')).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))\n",
    "        attn_vec = torch.matmul(attn_prob, wv)\n",
    "        return attn_vec.permute(0, 2, 1, 3).contiguous().view(bs, x_len, -1)\n",
    "    \n",
    "def _line_shift(x:Tensor):\n",
    "    \"Shift the line i of `x` by p-i elements to the left, is `mask` puts 0s on the diagonal.\"\n",
    "    bs,nh,n,p = x.size()\n",
    "    x_pad = torch.cat([x.new_zeros(bs,nh,n,1), x], dim=3)\n",
    "    x_shift = x_pad.view(bs,nh,p + 1,n)[:,:,1:].view_as(x)\n",
    "    return x_shift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(Module):\n",
    "    def __init__(self, n_heads:int, d_model:int, drops=0.1):\n",
    "        self.mhra = MultiHeadRelativeAttention(n_heads, d_model, drops=drops)\n",
    "        self.ff   = FeedForward(d_model, drops=drops)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ff(self.mhra(x, mask=mask, **kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class TransformerXL(Module):\n",
    "    \"TransformerXL model: https://arxiv.org/abs/1901.02860.\"\n",
    "    def __init__(self, vocab_sz:int, d_model:int, n_layers:int, n_heads:int, drops:float=0.1, mem_len:int=150):\n",
    "        d_head = d_model//n_heads\n",
    "        self.encoder = nn.Embedding(vocab_sz, d_model)\n",
    "        self.pos_enc = PositionalEncoding(d_model)\n",
    "        self.drop_emb = nn.Dropout(drops)\n",
    "        self.u = nn.Parameter(torch.Tensor(n_heads, 1, d_head)) #Remove 1 for einsum implementation of attention\n",
    "        self.v = nn.Parameter(torch.Tensor(n_heads, 1, d_head)) #Remove 1 for einsum implementation of attention\n",
    "        self.mem_len,self.n_layers,self.d_model = mem_len,n_layers,d_model\n",
    "        self.init = False\n",
    "        self.layers = nn.ModuleList([DecoderLayer(n_heads, d_model, drops) for k in range(n_layers)])\n",
    "\n",
    "    def reset(self):\n",
    "        \"Reset the internal memory.\"\n",
    "        self.hidden = [next(self.parameters()).data.new(0) for i in range(self.n_layers+1)]\n",
    "\n",
    "    def _update_mems(self, hids):\n",
    "        if not getattr(self, 'hidden', False): return None\n",
    "        assert len(hids) == len(self.hidden), 'len(hids) != len(self.hidden)'\n",
    "        with torch.no_grad():\n",
    "            for i in range(len(hids)):\n",
    "                cat = torch.cat([self.hidden[i], hids[i]], dim=1)\n",
    "                self.hidden[i] = cat[:,-self.mem_len:].detach()\n",
    "\n",
    "    def select_hidden(self, idxs): self.hidden = [h[idxs] for h in self.hidden]\n",
    "\n",
    "    def forward(self, x):\n",
    "        #The hidden state has to be initiliazed in the forward pass for nn.DataParallel\n",
    "        if self.mem_len > 0 and not self.init:\n",
    "            self.reset()\n",
    "            self.init = True\n",
    "        bs,x_len = x.size()\n",
    "        inp = self.drop_emb(self.encoder(x)) #.mul_(self.d_model ** 0.5)\n",
    "        m_len = self.hidden[0].size(1) if hasattr(self, 'hidden') and len(self.hidden[0].size()) > 1 else 0\n",
    "        seq_len = m_len + x_len\n",
    "        \n",
    "        mask = None  #MLM\n",
    "#         upper = torch.tril(x.new_ones(x_len, seq_len), diagonal=1+m_len).bool()[None,None]\n",
    "#         mask = torch.triu(x.new_ones(x_len, seq_len), diagonal=1+m_len).bool()[None,None]  # regular tfmrXL\n",
    "#         mask = upper & mask\n",
    "        \n",
    "        hids = []\n",
    "        pos = torch.arange(seq_len-1, -1, -1, device=inp.device, dtype=inp.dtype) #[len, len-1, len-2, len-3,...]\n",
    "        pos_enc = self.pos_enc(pos)\n",
    "        hids.append(inp)\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            mem = self.hidden[i] if self.mem_len > 0 else None\n",
    "            inp = layer(inp, r=pos_enc, u=self.u, v=self.v, mask=mask, mem=mem)\n",
    "            hids.append(inp)\n",
    "        core_out = inp[:,-x_len:]\n",
    "        if self.mem_len > 0 : self._update_mems(hids)\n",
    "        return (self.hidden if self.mem_len > 0 else [core_out]),[core_out]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vocab_sz = len(data.vocab.itos)\n",
    "encoder = TransformerXL(vocab_sz, 512, 4, n_heads=8, drops=0.1)\n",
    "decoder = LinearDecoder(vocab_sz, 512, 0.1, tie_encoder=encoder.encoder, bias=True)\n",
    "model = SequentialRNN(encoder, decoder)\n",
    "model.apply(init_transformer)\n",
    "learn = LanguageLearner(data, model, split_func=tfmerXL_lm_split, metrics=[accuracy, CER()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 20,525,152"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# BeRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MLM_Mask(LearnerCallback):\n",
    "    def __init__(self, learn:Learner, mlm_probability=0.3, mask_tok=7):\n",
    "        super().__init__(learn)\n",
    "        self.mask_tok = mask_tok\n",
    "        self.mlm_probability = mlm_probability\n",
    "        self.itos = learn.data.vocab.itos\n",
    "        \n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        new_input,new_target = self.mask_tokens(last_input)\n",
    "        return {'last_input':new_input, 'last_target':new_target}\n",
    "    \n",
    "    def mask_tokens(self, inputs):\n",
    "        \"\"\" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. \"\"\"\n",
    "        labels = inputs.clone()\n",
    "        # We sample a few tokens in each sequence for masked-LM training\n",
    "        masked_indices = torch.bernoulli(torch.full(labels.shape, self.mlm_probability)).bool()\n",
    "        labels[~masked_indices] = -1  # We only compute loss on masked tokens\n",
    "        \n",
    "        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n",
    "        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n",
    "        inputs[indices_replaced] = self.mask_tok\n",
    "\n",
    "        # 10% of the time, we replace masked input tokens with random word\n",
    "        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n",
    "        random_words = torch.randint(len(self.itos), labels.shape, dtype=torch.long, device=device)\n",
    "        inputs[indices_random] = random_words[indices_random]\n",
    "\n",
    "        # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n",
    "        return inputs, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def bert_acc(input:Tensor, targs:Tensor, ignore_index=-1)->Rank0Tensor:\n",
    "    mask = (targs!=ignore_index)\n",
    "    preds = input.argmax(dim=-1)\n",
    "    return (preds[mask]==targs[mask]).float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tfmr w/ MLM\n",
    "learn = make_learner(data, 512, 6, metrics=[bert_acc, CER()],\n",
    "                     callback_fns=[MLM_Mask], loss_func=CrossEntropyFlat(ignore_index=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tfmrXL w/ MLM\n",
    "vocab_sz = len(data.vocab.itos)\n",
    "encoder = TransformerXL(vocab_sz, 512, 6, n_heads=8, drops=0.1, mem_len=30)\n",
    "decoder = LinearDecoder(vocab_sz, 512, 0.1, tie_encoder=encoder.encoder, bias=True)\n",
    "model = SequentialRNN(encoder, decoder)\n",
    "model.apply(init_transformer)\n",
    "learn = LanguageLearner(data, model, split_func=tfmerXL_lm_split, metrics=[bert_acc, CER()],\n",
    "                        callback_fns=[MLM_Mask], loss_func=CrossEntropyFlat(ignore_index=-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('wiki2_bert_tfmr')#, strict=False)\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-3\n",
    "learn.fit_one_cycle(5, lr, callbacks=[SaveModelCallback(learn, name='wiki2_bert_tfmr8')])\n",
    "### Words: combo_60k ###\n",
    "# wiki2\n",
    "\n",
    "# MLM(.15),N=4,multi(8)\n",
    "# 2.953938\t2.823303\t0.088051\t0.907012\t03:24   tfmr  'wiki2_lm_bert'\n",
    "# 2.149526\t2.129630\t0.101498\t0.869657\t03:34   tfmrXL  'wiki2_lm_bertXL'\n",
    "# 0.204440\t0.175675\t0.979842\t0.041045\t04:13   diagonal+1  'wiki2_eye_tfmrXL'\n",
    "# 2.724019\t2.668955\t0.092031\t0.874019\t03:43   tfmrXL pretrained on wiki2_eye_tfmrXL\n",
    "\n",
    "# 5cycle(1e-3); 30k; N:6\n",
    "# 2.815044\t2.807196\t0.537950\t0.632317\t02:50   vanilla tfmr,  'wiki2_bert_tfmr'\n",
    "# 2.558876\t2.362679\t0.095183\t0.885061\t02:50   preload wiki2_bert_tfmr w/ MLM  'wiki2_bert_tfmr2'\n",
    "# 2.436609\t2.309809\t0.188842\t0.819951\t02:50   preload wiki2_bert_tfmr2 w/ MLM(.30)  'wiki2_bert_tfmr3'\n",
    "# 3.051376\t2.938654\t0.280082\t0.772065\t02:51   MLM(.50)  'wiki2_bert_tfmr4'\n",
    "# 2.513130\t2.432869\t0.185478\t0.823640\t02:51   preload wiki2_bert_tfmr4 w/ MLM(.30)  'wiki2_bert_tfmr5'\n",
    "\n",
    "# fixed acc & CER\n",
    "# 3.077021\t2.917398\t0.574049\t0.593222\t02:50   tfmr MLM(.15)\n",
    "# 2.528980\t2.358310\t0.629164\t0.541393\t02:51   preload wiki2_bert_tfmr w/ MLM(.15)  'wiki2_bert_tfmr6'\n",
    "# 2.575874\t2.412832\t0.614977\t0.549122\t02:54   preload wiki2_bert_tfmr w/ MLM(.30)  'wiki2_bert_tfmr7'\n",
    "# 2.165149\t2.158344\t0.659073\t0.503470\t03:06   tfmrXL MLM(.15)   'wiki2_bert_tfmr8'\n",
    "\n",
    "### Chars ###\n",
    "# wiki2: 10cycle, 1e-3\n",
    "\n",
    "# AWD-LSTM\n",
    "# 1.144401\t1.092210\t0.671697\t0.327961   (512/1400) 'wiki2_lm'\n",
    "\n",
    "# Tfmr \n",
    "# 1.432683\t1.367244\t0.594961\t0.405437   N=4,multi(8) 'wiki2_lm_tfmr'\n",
    "# 1.371769\t1.318286\t0.607220\t0.394283   N=6,multi(8) 'wiki2_lm_tfmr2'\n",
    "\n",
    "# TfmrXL\n",
    "# 1.174197\t1.160675\t0.652781\t0.349375   N=6,multi(8)  'wiki2_lm_tfmrXL'\n",
    "# 1.168891\t1.129729\t0.659086\t0.343158   5cycle,1e-6   'wiki2_lm_tfmrXL_v2'\n",
    "# 1.171422\t1.126684\t0.660261\t0.342019   pretrained on above; N=10, 3cycle,1e-4   wiki2_lm_tfmrXL_v3\n",
    "# 1.204486\t1.190764\t0.645694\t0.356218   N=10,multi(8)   wiki2_lm_tfmrXL_v4\n",
    "\n",
    "# 1cycle, 1e-3\n",
    "# 1.466067\t1.401443\t0.587846\t0.413188   fastai TransformerXL\n",
    "# 1.462828\t1.403346\t0.587787\t0.413295   manual TransformerXL\n",
    "\n",
    "# wiki103: 3cycle, 1e-4, (stopped after 1st cycle)\n",
    "# 1.216455\t1.101982\t0.666687\t0.334717   'wiki103_lm'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(iter(data.valid_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x[2],y[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "preds = learn.model(x[2][None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pred = torch.argmax(preds[1][0][0], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "[itos[word.item()] for word in pred]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def predict(self:learn, text:str, n_words:int=1, no_unk:bool=True, temperature:float=1., min_p:float=None, sep:str=' ',\n",
    "            decoder=decode_spec_tokens):\n",
    "    \"Return `text` and the `n_words` that come after\"\n",
    "    self.model.reset()\n",
    "    xb,yb = self.data.one_item(text)\n",
    "    \n",
    "    # remove the eos token which is automatically added\n",
    "    xb = xb[:,:-1]\n",
    "    print(xb)\n",
    "\n",
    "    \n",
    "    new_idx = []\n",
    "    for _ in range(n_words): #progress_bar(range(n_words), leave=False):\n",
    "        res = self.pred_batch(batch=(xb,yb))[0][-1]\n",
    "        #if len(new_idx) == 0: self.model[0].select_hidden([0])\n",
    "        if no_unk: res[self.data.vocab.stoi[UNK]] = 0.\n",
    "        if min_p is not None:\n",
    "            if (res >= min_p).float().sum() == 0:\n",
    "                warn(f\"There is no item with probability >= {min_p}, try a lower value.\")\n",
    "            else: res[res < min_p] = 0.\n",
    "        if temperature != 1.: res.pow_(1 / temperature)\n",
    "        idx = torch.multinomial(res, 1).item()\n",
    "        new_idx.append(idx)\n",
    "        xb = xb.new_tensor([idx])[None]\n",
    "    return text + sep + sep.join(decoder(self.data.vocab.textify(new_idx, sep=None)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "predict(learn, \"This is a wonderful\", n_words=3, sep='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save_encoder('wiki2_lm_enc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learner.model.eval()\n",
    "learner.model.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "get_next('whe')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "get_next_n('th', 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "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.7.3"
  },
  "notify_time": "30",
  "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
}