1--SentencePiece-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\n",
    "import sentencepiece as spm"
   ]
  },
  {
   "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(' ]', ']').replace(\n",
    "        ' %', '%').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": "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": [
    "wiki_df = DataFrame({'text': lines[1:-1]})\n",
    "wiki_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "wiki_df.to_csv(PATH/'clean_wiki_103.csv', 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",
    "imdb_df = pd.read_csv(imdb_path/CSV)\n",
    "len(imdb_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb_df['text'] = imdb_df.text.apply(lambda x: cleanup(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "len(imdb_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb_df.text.values[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "imdb_df.to_csv(PATH/'clean_imdb.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Book texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FPATH = Path('data/fonts/')\n",
    "\n",
    "a = pd.read_csv(FPATH/'knot.csv', usecols=['filename', 'label'])\n",
    "c = pd.read_csv(FPATH/'adrift.csv', usecols=['filename', 'label'])\n",
    "d = pd.read_csv(FPATH/'zane.csv', usecols=['filename', 'label'])\n",
    "e = pd.read_csv(FPATH/'american.csv', usecols=['filename', 'label'])\n",
    "f = pd.read_csv(FPATH/'age.csv', usecols=['filename', 'label'])\n",
    "g = pd.read_csv(FPATH/'room.csv', usecols=['filename', 'label'])\n",
    "book_df = pd.concat([a,c,d,e,f,g], ignore_index=True)\n",
    "len(book_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "book_df['text'] = book_df.label.apply(lambda x: x.replace('\\n',' '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "book_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "book_df.to_csv(PATH/'clean_books.csv', columns=['text'], index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Paragraphs (w/out test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pg_df = pd.read_csv(PATH/'edited_pg.csv')\n",
    "pg_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pg_df['text'] = pg_df.text.apply(lambda x: x.replace('\\n',' '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pg_df.to_csv(PATH/'clean_pg.csv', columns=['text'], index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Combine and process (remove caps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "a = pd.read_csv(PATH/'clean_pg.csv')\n",
    "b = pd.read_csv(PATH/'clean_books.csv')\n",
    "c = pd.read_csv(PATH/'clean_imdb.csv')\n",
    "d = pd.read_csv(PATH/'clean_wiki_103.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "c.text.values[300]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full = pd.concat([a, b, c, d], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full.dropna(inplace=True)\n",
    "full.reset_index(inplace=True, drop=True)\n",
    "len(full)"
   ]
  },
  {
   "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full['text'] = full.text.apply(lambda x: rm_useless_spaces(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full['text'] = full.text.apply(lambda x: add_cap_tokens(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full.tail()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Write to raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "name = str(PATH/'spm_full')\n",
    "# name = str(PATH/'spm_pg')\n",
    "fname = name + '.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# puts dataset into format expected by sentencepiece:\n",
    "# .txt file entries separated by \\n\n",
    "def write_text(texts, filename=fname):\n",
    "    with open(filename, 'w', encoding='utf-8') as f:\n",
    "        for text in texts:\n",
    "            f.write(text + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "write_text(full.text.values)\n",
    "# write_text(pg_df.text.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true
   },
   "source": [
    "# Train SPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "name = str(PATH/'spm_full')\n",
    "fname = name + '.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "symbols = \"\\n,[UP],[MAJ],▁,:,;,!,?,(,),[,],{,},<,>,@,#,$,%,^,&,*,-,_,+,=,/,~\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "spm.SentencePieceTrainer.Train(\n",
    "    f\"--unk_id=3 --pad_id=0 --input={fname} --model_prefix={name+'_30k'} --vocab_size=30000 --user_defined_symbols={symbols} --input_sentence_size=1500000 --shuffle_input_sentence=True\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(name+'_30k.model')"
   ]
  },
  {
   "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": [
    "df = pd.read_csv(fname, sep='\\n', header=None, names=['text'])\n",
    "st = df.text.values[-2]; st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pieces = sp.encode_as_pieces(st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "st2 = sp.decode_pieces(pieces[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "st == st2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sp.EncodeAsPieces('adaptability')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for n in range(5):\n",
    "    print(sp.SampleEncodeAsPieces('adaptability', -1, .1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vocab = {i: sp.id_to_piece(i) for i in range(len(sp))}\n",
    "vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DataBunch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "\n",
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(str(PATH/'spm_full_10k.model'))\n",
    "sp.SetEncodeExtraOptions(\"eos\")\n",
    "sp.SetDecodeExtraOptions(\"bos:eos\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# itos = {i:sp.id_to_piece(i) for i in range(len(sp))}"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SPMTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[int]: return [1] + sp.EncodeAsIds(t)[1:]  #remove initial space\n",
    "\n",
    "class SPMProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.toknizr = Tokenizer(tok_func=SPMTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "                                 post_rules=[], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "    \n",
    "    def process(self, ds):\n",
    "        tokens = []\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            tokens += self.toknizr.process_all(ds.items[i:i+self.chunksize])\n",
    "        ds.items = tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SPMTextList(TextList):\n",
    "    _bunch = TextLMDataBunch   # this is the databunch created when calling .databunch()\n",
    "    _processor = []\n",
    "    _is_lm = True\n",
    "\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.vocab = sp\n",
    "        self.pad_idx = 0\n",
    "        self.copy_new += ['vocab']\n",
    "    \n",
    "    def get(self, i): return Text(i, self.textify(i))\n",
    "    \n",
    "    def reconstruct(self, x:Tensor):\n",
    "        nonzero_idxs = (x != self.pad_idx).nonzero()\n",
    "        idx_max = nonzero_idxs.max() if len(nonzero_idxs) > 0 else 0\n",
    "        return Text(x[0:idx_max+1], self.textify(x[0:idx_max+1]))\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "        \n",
    "    def textify(self, ids):\n",
    "        if isinstance(ids, torch.Tensor): ids = ids.tolist()\n",
    "        st = self.vocab.DecodeIds(ids)\n",
    "        st = remove_cap_tokens(st)\n",
    "        return st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LMLabelList(EmptyLabelList):\n",
    "    \"Basic `ItemList` for dummy labels.\"\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.loss_func = CrossEntropyFlat()\n",
    "        \n",
    "    def reconstruct(self, t:Tensor, x:Tensor=None):\n",
    "        if len(t.size()) == 0: return EmptyLabel()\n",
    "        return self.x.reconstruct(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = 'clean_pg.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (SPMTextList.from_csv(PATH, CSV, cols=0, processor=SPMProcessor())\n",
    "        .split_by_rand_pct(valid_pct=0.10, seed=42)\n",
    "        .label_const(0, label_cls=LMLabelList)\n",
    "        .databunch(bs=64, device=device)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=5, ds_type=DatasetType.Train)"
   ]
  },
  {
   "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, 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
   },
   "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": [
    "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.vocab_sz = len(learn.data.vocab)\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_sz, 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": "markdown",
   "metadata": {},
   "source": [
    "# BeRT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": 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 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": {},
   "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_heads=8):        \n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "        ff = PositionwiseFeedForward(d_model, drops)\n",
    "        self.embedding = Embeddings(d_model, vocab)\n",
    "        self.pos_enc = PositionalEncoding(d_model, drops)\n",
    "        self.encoder = Encoder(EncoderLayer(d_model, attn, ff, drops), N)\n",
    "        \n",
    "        self.generator = nn.Linear(d_model, vocab)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.decode(self.encode(x))\n",
    "    \n",
    "    def encode(self, x):\n",
    "        return self.encoder(self.pos_enc(self.embedding(x)))\n",
    "        \n",
    "    def decode(self, x):\n",
    "        return self.generator(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def init_tfmr_lm(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.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model=512, N=4, drops=0.1, attn_heads=8, tie_weights=True, **learn_kwargs):\n",
    "    vocab_sz = len(sp)\n",
    "    model = TransformerLM(vocab_sz, d_model, N=N, drops=drops, attn_heads=attn_heads)\n",
    "    model.apply(init_tfmr_lm)\n",
    "    if tie_weights: model.generator.weight = model.embedding.lut.weight\n",
    "    return Learner(data, model, **learn_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 4, metrics=[bert_acc, CER(data.x.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": [
    "# true number of trainable params\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)\n",
    "\n",
    "# Total trainable params: 22,860,560"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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
   },
   "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-2\n",
    "learn.fit_one_cycle(5, lr)#, callbacks=[SaveModelCallback(learn, name='pg_bert_tfmr')])\n",
    "#clean_pg.csv\n",
    "\n",
    "# 6.215954\t4.791131\t0.426439\t0.768457\t00:05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-2\n",
    "learn.fit_one_cycle(5, lr)#, callbacks=[SaveModelCallback(learn, name='pg_bert_tfmr')])\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": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Test"
   ]
  },
  {
   "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": "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": []
  },
  {
   "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.6.5"
  },
  "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
}