1--Chars&Words.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.vision import *\n",
    "from fastai.text import *\n",
    "from fastai.callbacks.tracker import *\n",
    "from fastai.callbacks.hooks import *\n",
    "\n",
    "import pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from transformers import DistilBertTokenizer, DistilBertForMaskedLM"
   ]
  },
  {
   "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": {},
   "source": [
    "## Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def show_img(im, figsize=None, ax=None, alpha=None, title=None):\n",
    "    if not ax: fig,ax = plt.subplots(figsize=figsize)\n",
    "    ax.imshow(image2np(im.data), alpha=alpha)\n",
    "    if title: ax.set_title(title)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def rshift(tgt, bos_token=1):\n",
    "    \"Shift y to the right by prepending token\"\n",
    "    bos = torch.zeros((tgt.size(0),1), device=device).type_as(tgt) + bos_token\n",
    "    return torch.cat((bos, tgt[:,:-1]), dim=-1)\n",
    "\n",
    "def subsequent_mask(size):\n",
    "    return torch.tril(torch.ones((1,size,size), device=device).byte())\n",
    "    \n",
    "def parallelogram_mask(size, diagonal):\n",
    "    mask = torch.ones((1,size,size), device=device).byte()\n",
    "    upper = torch.tril(mask).bool()\n",
    "    lower = torch.triu(mask, diagonal=-diagonal).bool()\n",
    "    return (upper & lower).byte()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss, Metrics, Callbacks"
   ]
  },
  {
   "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):\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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
   },
   "outputs": [],
   "source": [
    "class TeacherForce(LearnerCallback):\n",
    "    def __init__(self, learn:Learner):\n",
    "        super().__init__(learn)\n",
    "        \n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        return {'last_input':(last_input, last_target), 'last_target':last_target}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## paragraphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'edited_pg.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'paragraphs'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 512,10\n",
    "seq_len,word_len = 700,300"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## sm synth dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'edited_sm_synth.csv' #'small_synth_words.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'edited_sm_synth'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 256,50#100\n",
    "seq_len,word_len = 100,50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## combo_cat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#font generated\n",
    "fname = 'font_mix_129k.csv'\n",
    "FOLDER = 'combo_cat'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#handwriting only\n",
    "fname = 'hand_mix_25k.csv'\n",
    "FOLDER = 'combo_cat'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = PATH/fname\n",
    "df = pd.read_csv(CSV)\n",
    "\n",
    "sz,bs = 512,10\n",
    "seq_len,word_len = 750,300\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### old"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'combo_145k.csv'\n",
    "FOLDER = 'combo_cat'\n",
    "\n",
    "CSV = PATH/fname\n",
    "df = pd.read_csv(CSV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# full\n",
    "sz,bs = 512,16\n",
    "seq_len,word_len = 750,300\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 6 and fewer\n",
    "df = df[df.num_lines <= 6]\n",
    "sz,bs = 512,20\n",
    "seq_len,word_len = 600,200\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 6 and greater\n",
    "df = df[df.num_lines >= 6]\n",
    "sz,bs = 512,15\n",
    "seq_len,word_len = 750,300\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## test combo (no fonts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'test_combo'\n",
    "CSV = PATH/'test_combo.csv'\n",
    "df = pd.read_csv(CSV)\n",
    "\n",
    "sz,bs = 512,20\n",
    "seq_len,word_len = 750,300\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ModelData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tfms = get_transforms(do_flip=False, max_zoom=1, max_rotate=0, max_warp=0.1)\n",
    "\n",
    "def force_gray(image): return image.convert('L').convert('RGB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Char or Word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def label_collater(samples:BatchSamples, pad_idx:int=0):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1 and lbls[0] is 0:   #predict\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels    \n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(data), max_len+1).long() + pad_idx  # 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 imgs, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SequenceList(TextList):    \n",
    "    def __init__(self, items:Iterator, vocab:Vocab, tokenizer:Tokenizer, **kwargs):\n",
    "        toknizr = Tokenizer(tok_func=tokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        procs = [TokenizeProcessor(tokenizer=toknizr, include_bos=False), NumericalizeProcessor(vocab=vocab)]\n",
    "        super().__init__(items, vocab, sep='', pad_idx=0, processor=procs)\n",
    "    \n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "itos = pickle.load(open(PATH/'itos.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CharTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[str]: return list(t)+['xxeos']\n",
    "            \n",
    "class CharVocab(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",
    "        st = sep.join([self.itos[i] for i in nums])\n",
    "#         st = remove_special_toks(st)\n",
    "        return st\n",
    "    \n",
    "def remove_special_toks(text):\n",
    "    #text = re.sub(r'(xxbos)', '', text)  #[CLS] (w/ following whitespace)\n",
    "    text = re.sub(r'(xxeos)', '', text)  #[SEP] (w/ preceding whitespace)\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SequenceList, vocab=CharVocab(itos), tokenizer=CharTokenizer)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "#         .transform(tfms, size=sz, resize_method=ResizeMethod.PAD, padding_mode='border')\n",
    "        # maintains aspect ratio but too small for good results => mostly whitespace\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "        #.normalize() # this sets x values to an odd range (~.3,-6)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "word_itos = pickle.load(open(PATH/'word_itos_60k_mod.pkl', 'rb'))\n",
    "#word_itos = word_itos[:10000]\n",
    "len(word_itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[str]: \n",
    "        chars = list(t.lower())\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 res + ['xxeos']\n",
    "\n",
    "# class WordTokenizer(BaseTokenizer):\n",
    "#     def tokenizer(self, t:str) -> List[str]: return t.lower().replace('\\n', ' \\n ').split(' ') + ['xxeos']\n",
    "            \n",
    "class WordVocab(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",
    "        st = sep.join([self.itos[i] for i in nums])\n",
    "        st = remove_special_toks(st)\n",
    "        return st\n",
    "    \n",
    "    def remove_special_toks(text):\n",
    "        text = re.sub(r'(xxbos)', '', text)  #[CLS] (w/ following whitespace)\n",
    "        text = re.sub(r'(xxeos)', '', text)  #[SEP] (w/ preceding whitespace)\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# w_vocab = Vocab(word_itos)\n",
    "# w_procs = [TokenizeProcessor(include_bos=False, include_eos=True), NumericalizeProcessor(vocab=w_vocab)]\n",
    "\n",
    "words = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SequenceList, vocab=WordVocab(word_itos), tokenizer=WordTokenizer)\n",
    "        #.label_from_df(label_cls=TextList, pad_idx=0, vocab=w_vocab, processor=w_procs)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "words.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Words (bert_tok)"
   ]
  },
  {
   "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]'])"
   ]
  },
  {
   "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 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": [
    "# class BertWordList(TextList):    \n",
    "#     def __init__(self, items:Iterator, vocab, **kwargs):\n",
    "#         toknizr = Tokenizer(tok_func=BertTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "#                             post_rules=[], special_cases=[])\n",
    "#         procs = [TokenizeProcessor(tokenizer=toknizr, include_bos=False), NumericalizeProcessor(vocab=vocab)]\n",
    "#         super().__init__(items, vocab, sep='', pad_idx=0, processor=procs)\n",
    "    \n",
    "#     def analyze_pred(self, pred:Tensor):\n",
    "#         return torch.argmax(pred, dim=-1)\n",
    "    \n",
    "# # .label_from_df(label_cls=BertWordList, vocab=BertVocab())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.toknizr = Tokenizer(tok_func=BertTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "                                 post_rules=[], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\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\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.vocab = ds.vocab\n",
    "        \n",
    "    def process_one(self,item):\n",
    "        return np.array(self.vocab.numericalize(item), dtype=np.int64)\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])    \n",
    "        \n",
    "\n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "\n",
    "    def get(self, i):\n",
    "        w = self.items[i]\n",
    "        return Text(w, self.vocab.textify(w))\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        idx_min,idx_max = (t != self.pad_idx).nonzero().min(), (t != self.pad_idx).nonzero().max()\n",
    "        return Text(t[idx_min:idx_max+1], self.vocab.textify(t[idx_min:idx_max+1]))\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList, vocab=BertVocab(), pad_idx=0)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Chars (bert_tok)"
   ]
  },
  {
   "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]'])"
   ]
  },
  {
   "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 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": [
    "def characterize(x:Collection[str]) -> Collection[str]:\n",
    "    \"Separate word tokens into letters.\"\n",
    "    res = []\n",
    "    for t in x:\n",
    "        if t in ['\\n',' ','[UP]','[MAJ]','[UNK]', '[CLS]', '[MASK]', '[PAD]', '[SEP]']:\n",
    "            res.append(t)\n",
    "        elif t.startswith('##'):  # wordpiece\n",
    "            [res.append(c) for c in list(t[2:])]\n",
    "        else:\n",
    "            [res.append(c) for c in list(t)]\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.toknizr = Tokenizer(tok_func=BertTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "                                 post_rules=[characterize], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\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\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.vocab = ds.vocab\n",
    "        \n",
    "    def process_one(self,item):\n",
    "        return np.array(self.vocab.numericalize(item), dtype=np.int64)\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])    \n",
    "        \n",
    "\n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "\n",
    "    def get(self, i):\n",
    "        w = self.items[i]\n",
    "        return Text(w, self.vocab.textify(w))\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        idx_min,idx_max = (t != self.pad_idx).nonzero().min(), (t != self.pad_idx).nonzero().max()\n",
    "        return Text(t[idx_min:idx_max+1], self.vocab.textify(t[idx_min:idx_max+1]))\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList, vocab=BertVocab(), pad_idx=0)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### SentencePiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(str(PATH/'spm_full_30k.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": [
    "from string import ascii_lowercase"
   ]
  },
  {
   "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[str]: return sp.EncodeAsIds(t)\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_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\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\n",
    "\n",
    "class SPMList(ItemList):\n",
    "    _processor = [SPMProcessor]\n",
    "\n",
    "    def __init__(self, items:Iterator, sp, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.vocab = sp\n",
    "        self.copy_new += ['vocab']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        return Text(o, self.textify(o))\n",
    "\n",
    "    def reconstruct(self, t:Tensor):\n",
    "        return Text(t, self.textify(t))\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": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SPMList, sp=sp)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Char and Word"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Trio w/ Bert tokenizer"
   ]
  },
  {
   "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 on 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": [
    "bert_tok = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BertTokenizer(BaseTokenizer): \n",
    "    def __init__(self, tokenizer=bert_tok, **kwargs):\n",
    "        self._tok_func = tokenizer\n",
    "        \n",
    "    def __call__(self, *args, **kwargs): return self \n",
    "    \n",
    "    def tokenizer(self, t:str) -> List[str]: \n",
    "        return [\"[CLS]\"] + self._tok_func.tokenize(t) + [\"[SEP]\"]"
   ]
  },
  {
   "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": "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls, bert_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1:\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls), c_pad(bert_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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 MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.c_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','\\n'])\n",
    "        self.w_toknizr = Tokenizer(tok_func=CustomTokenizer, pre_rules=[rm_useless_spaces],\n",
    "                                   special_cases=['xxbos','xxeos','xxmask','xxunk','xxpad','xxmaj','xxup','\\n'])\n",
    "        self.b_toknizr = Tokenizer(tok_func=BertTokenizer(bert_tok),\n",
    "                           pre_rules=[], post_rules=[], special_cases=['[CLS], [SEP]'])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\n",
    "    \n",
    "    def process(self, ds):\n",
    "        ds.items = list(zip(self._p(self.c_toknizr, ds), self._p(self.w_toknizr, ds), self._p(self.b_toknizr, ds)))\n",
    "    \n",
    "    def _p(self, toknizr, ds):\n",
    "        res = []\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            res += toknizr.process_all(ds.items[i:i+self.chunksize])\n",
    "        return res\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.vocab = CustomVocab(itos)\n",
    "        self.bert_vocab = Vocab(list(bert_tok.vocab.keys()))\n",
    "        \n",
    "    def process_one(self,item):\n",
    "        chars = np.array(self.vocab.numericalize(item[0]), dtype=np.int64)\n",
    "        words = np.array(self.vocab.numericalize(item[1]), dtype=np.int64)\n",
    "        bert_words = np.array(self.bert_vocab.numericalize(item[2]), dtype=np.int64)\n",
    "        return [chars,words,bert_words]\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])\n",
    "        \n",
    "        \n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "    \n",
    "    def __init__(self, items:Iterator, pad_idx=0, sep='', **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.vocab = CustomVocab(itos)\n",
    "        self.bert_tok = bert_tok\n",
    "        self.bert_vocab = Vocab(list(bert_tok.vocab.keys()))\n",
    "        self.pad_idx, self.sep = pad_idx,sep\n",
    "        self.copy_new += ['vocab', 'bert_vocab', 'bert_tok', 'pad_idx', 'sep']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = self.items[i]\n",
    "        return [self._get(o[0]), self._get(o[1]), self._get(o[2], self.bert_vocab, ' ')]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        \"t: List of tensors -> [char, word, bert]\"        \n",
    "        return [self._recon(t[0]), self._recon(t[1]), self._recon(t[2], self.bert_vocab)]\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "    \n",
    "    def _get(self, x, vocab=None, sep=None):\n",
    "        vocab = ifnone(vocab, self.vocab)\n",
    "        sep = ifnone(sep, self.sep)\n",
    "        return Text(x, vocab.textify(x, sep))\n",
    "    \n",
    "    def _recon(self, x, vocab=None):\n",
    "        vocab = ifnone(vocab, self.vocab)\n",
    "        idx_min,idx_max = (x != self.pad_idx).nonzero().min(), (x != self.pad_idx).nonzero().max()\n",
    "        return Text(x[idx_min:idx_max+1], vocab.textify(x[idx_min:idx_max+1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words,berts = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words) + '\\n\\n' + str(berts))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Combo (bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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": [
    "bert_tok = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "num_added_tokens = bert_tok.add_tokens(['\\n',' ','[UP]','[MAJ]'])"
   ]
  },
  {
   "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": [
    "def characterize(x:Collection[str]) -> Collection[str]:\n",
    "    \"Separate word tokens into letters.\"\n",
    "    res = []\n",
    "    for t in x:\n",
    "        if t in ['\\n',' ','[UP]','[MAJ]','[UNK]', '[CLS]', '[MASK]', '[PAD]', '[SEP]']:\n",
    "            res.append(t)\n",
    "        elif t.startswith('##'):  # wordpiece\n",
    "            [res.append(c) for c in list(t[2:])]\n",
    "        else:\n",
    "            [res.append(c) for c in list(t)]\n",
    "    return res"
   ]
  },
  {
   "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]:\n",
    "        w_toks = bert_tok.tokenize(t) + [\"[SEP]\"]\n",
    "        c_toks = characterize(w_toks)\n",
    "        return [c_toks, w_toks]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": [
    "class MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.tokenizer = Tokenizer(tok_func=BertTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],\n",
    "                                 post_rules=[], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\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.tokenizer.process_all(ds.items[i:i+self.chunksize])\n",
    "        ds.items = tokens\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.vocab = ds.vocab\n",
    "        \n",
    "    def process_one(self,item):\n",
    "        chars = np.array(self.vocab.numericalize(item[0]), dtype=np.int64)\n",
    "        words = np.array(self.vocab.numericalize(item[1]), dtype=np.int64)\n",
    "        return [chars,words]\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])    \n",
    "        \n",
    "\n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "\n",
    "    def get(self, i):\n",
    "        c,w = self.items[i]\n",
    "        return [Text(c, self.vocab.textify(c)), Text(w, self.vocab.textify(w))]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        c,w = t\n",
    "        return [self.reconstruct_one(c),self.reconstruct_one(w)]\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "    \n",
    "    def reconstruct_one(self, x):\n",
    "        nonzero_idxs = (x != self.pad_idx).nonzero()\n",
    "        idx_min = 0  #(x != self.pad_idx).nonzero().min()\n",
    "        idx_max = nonzero_idxs.max() if len(nonzero_idxs) > 0 else 0\n",
    "        return Text(x[idx_min:idx_max+1], self.vocab.textify(x[idx_min:idx_max+1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList, vocab=BertVocab(), pad_idx=0)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combo SentencePiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(str(PATH/'spm_full_30k.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 = {}\n",
    "for i in range(len(sp)):\n",
    "    itos[i] = sp.id_to_piece(i)\n",
    "\n",
    "c_itos={}\n",
    "for k,v in itos.items():\n",
    "    if k<7:\n",
    "        c_itos[k] = [k]\n",
    "    else:\n",
    "        c_itos[k] = [sp.piece_to_id(c) for c in list(v)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def characterize(t:Collection[int]) -> Collection[int]:\n",
    "    return functools.reduce(operator.iconcat, [c_itos[c] for c in t], [])\n",
    "    # flatten nested list - fastest\n",
    "    # https://stackoverflow.com/questions/952914/how-to-make-a-flat-list-out-of-list-of-lists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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": [
    "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]:\n",
    "        w_toks = sp.EncodeAsIds(t)\n",
    "        c_toks = characterize(w_toks)\n",
    "        return [c_toks, w_toks]\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_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\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 SPMMultiList(ItemList):\n",
    "    _processor = [SPMProcessor]\n",
    "\n",
    "    def __init__(self, items:Iterator, sp, **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):\n",
    "        c,w = self.items[i]\n",
    "        return [Text(c, self.textify(c)), Text(w, self.textify(w))]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        c,w = t\n",
    "        return [self.reconstruct_one(c),self.reconstruct_one(w)]\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "    \n",
    "    def reconstruct_one(self, x):\n",
    "        nonzero_idxs = (x != self.pad_idx).nonzero()\n",
    "        idx_min = 0  #(x != self.pad_idx).nonzero().min()\n",
    "        idx_max = nonzero_idxs.max() if len(nonzero_idxs) > 0 else 0\n",
    "        return Text(x[idx_min:idx_max+1], self.textify(x[idx_min:idx_max+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 ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SPMMultiList, sp=sp)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Combo (shared itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def remove_cap_tokens(text):  # after decode\n",
    "    text = re.sub(r'xxup\\w+', lambda m: m.group()[4:].upper(), text)  #cap entire word\n",
    "    text = re.sub(r'xxmaj\\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 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], handle_cap_tokens:bool=False, sep=''):\n",
    "        res = sep.join([self.itos[i] for i in nums]) if sep is not None else [self.itos[i] for i in nums]\n",
    "        if handle_cap_tokens:\n",
    "            return remove_cap_tokens(res)\n",
    "        else:\n",
    "            return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CustomTokenizer(BaseTokenizer):\n",
    "    \"Split on 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": [
    "itos = pickle.load(open(PATH/'combo_itos_60k.pkl', 'rb'))\n",
    "itos = itos[:30000]\n",
    "#vocab = CustomVocab(itos)"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1:\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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 MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.c_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','\\n'])\n",
    "        self.w_toknizr = Tokenizer(tok_func=CustomTokenizer, pre_rules=[rm_useless_spaces],\n",
    "                                   special_cases=['xxbos','xxeos','xxmask','xxunk','xxpad','xxmaj','xxup','\\n'])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\n",
    "    \n",
    "    def process(self, ds):\n",
    "        ds.items = list(zip(self._p(self.c_toknizr, ds), self._p(self.w_toknizr, ds)))\n",
    "    \n",
    "    def _p(self, toknizr, ds):\n",
    "        res = []\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            res += toknizr.process_all(ds.items[i:i+self.chunksize])\n",
    "        return res\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.vocab = ds.vocab #CustomVocab(itos)\n",
    "        \n",
    "    def process_one(self,item):\n",
    "        chars = np.array(self.vocab.numericalize(item[0]), dtype=np.int64)\n",
    "        words = np.array(self.vocab.numericalize(item[1]), dtype=np.int64)\n",
    "        return [chars,words]\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])\n",
    "        \n",
    "        \n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "    \n",
    "    def __init__(self, items:Iterator, pad_idx=0, sep='', **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.vocab = CustomVocab(itos)        \n",
    "        self.pad_idx, self.sep = pad_idx,sep\n",
    "        self.copy_new += ['vocab', 'pad_idx', 'sep']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = self.items[i]\n",
    "        return [self._get(o[0]),self._get(o[1])]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        \"t: List of tensors -> [char, word]\"        \n",
    "        c,w = t\n",
    "        return [self._recon(c),self._recon(w)]\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "    \n",
    "    def _get(self, x):\n",
    "        return Text(x, self.vocab.textify(x, self.sep))\n",
    "    \n",
    "    def _recon(self, x):\n",
    "        idx_min,idx_max = (x != self.pad_idx).nonzero().min(), (x != self.pad_idx).nonzero().max()\n",
    "        return Text(x[idx_min:idx_max+1], self.vocab.textify(x[idx_min:idx_max+1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Both (char itos + bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1:\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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": [
    "bert_tok = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "class BertTokenizer(BaseTokenizer): \n",
    "    def tokenizer(self, t:str) -> List[str]: \n",
    "        return bert_tok.tokenize(t) + [\"[SEP]\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.c_toknizr = Tokenizer(tok_func=CharTokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        self.w_toknizr = Tokenizer(tok_func=BertTokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\n",
    "    \n",
    "    def process(self, ds):\n",
    "        ds.chars = ds.items\n",
    "        ds.words = ds.items\n",
    "\n",
    "        char_tokens = []\n",
    "        word_tokens = []\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            char_tokens += self.c_toknizr.process_all(ds.chars[i:i+self.chunksize])\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            word_tokens += self.w_toknizr.process_all(ds.words[i:i+self.chunksize])\n",
    "        ds.items = list(zip(char_tokens, word_tokens))\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.char_vocab = ds.char_vocab\n",
    "        self.word_vocab = ds.word_vocab\n",
    "\n",
    "    def process_one(self,item):\n",
    "        chars = np.array(self.char_vocab.numericalize(item[0]), dtype=np.int64)\n",
    "        words = np.array(self.word_vocab.numericalize(item[1]), dtype=np.int64)\n",
    "        return [chars,words]\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])\n",
    "        \n",
    "        \n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "            \n",
    "    def __init__(self, items:Iterator, pad_idx=0, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.char_vocab = CharVocab(itos)\n",
    "        self.word_vocab = Vocab(list(bert_tok.vocab.keys()))\n",
    "        \n",
    "        self.pad_idx = pad_idx\n",
    "        self.copy_new += ['char_vocab', 'word_vocab', 'pad_idx']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        chars = Text(o[0], self.char_vocab.textify(o[0], ''))\n",
    "        words = Text(o[1], self.word_vocab.textify(o[1], ' '))\n",
    "        return [chars,words]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        \"t: List of tensors -> [char, word]\"        \n",
    "        c,w = t\n",
    "        idx_min,idx_max = (c != self.pad_idx).nonzero().min(), (c != self.pad_idx).nonzero().max()\n",
    "        chars = Text(c[idx_min:idx_max+1], self.char_vocab.textify(c[idx_min:idx_max+1]))\n",
    "        \n",
    "        idx_min,idx_max = (w != self.pad_idx).nonzero().min(), (w != self.pad_idx).nonzero().max()\n",
    "        words = Text(w[idx_min:idx_max+1], self.word_vocab.textify(w[idx_min:idx_max+1]))\n",
    "        return [chars,words]\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Both (separate itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def multi_label_collater(samples:BatchSamples):\n",
    "    \"Function that collect samples and pads ends of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    char_lbls, word_lbls = zip(*lbls)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1:\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels\n",
    "    return imgs, (c_pad(char_lbls), c_pad(word_lbls))\n",
    "    \n",
    "def c_pad(lbls, pad_idx=0):\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(lbls), max_len+1).long() + pad_idx  # 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 MultiTokenizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None, chunksize:int=10000):\n",
    "        self.c_toknizr = Tokenizer(tok_func=CharTokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        self.w_toknizr = Tokenizer(tok_func=WordTokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        self.chunksize = chunksize\n",
    "        \n",
    "    def process_one(self, item):\n",
    "        raise Exception(\"This isn't implemented!  I didn't think it was necessary...\")\n",
    "    \n",
    "    def process(self, ds):\n",
    "        ds.chars = ds.items\n",
    "        ds.words = ds.items\n",
    "\n",
    "        char_tokens = []\n",
    "        word_tokens = []\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            char_tokens += self.c_toknizr.process_all(ds.chars[i:i+self.chunksize])\n",
    "        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):\n",
    "            word_tokens += self.w_toknizr.process_all(ds.words[i:i+self.chunksize])\n",
    "        ds.items = list(zip(char_tokens, word_tokens))\n",
    "        \n",
    "class MultiNumericalizeProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.char_vocab = CharVocab(itos)\n",
    "        self.word_vocab = WordVocab(word_itos)\n",
    "\n",
    "    def process_one(self,item):\n",
    "        chars = np.array(self.char_vocab.numericalize(item[0]), dtype=np.int64)\n",
    "        words = np.array(self.word_vocab.numericalize(item[1]), dtype=np.int64)\n",
    "        return [chars,words]\n",
    "            \n",
    "    def process(self, ds):\n",
    "        ds.items = array([self.process_one(item) for item in ds.items])\n",
    "        \n",
    "        \n",
    "class MultiSequenceList(TextList):\n",
    "    _processor = [MultiTokenizeProcessor, MultiNumericalizeProcessor]\n",
    "            \n",
    "    def __init__(self, items:Iterator, pad_idx=0, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.char_vocab = CharVocab(itos)\n",
    "        self.word_vocab = WordVocab(word_itos)\n",
    "        \n",
    "        self.pad_idx = pad_idx\n",
    "        self.copy_new += ['char_vocab', 'word_vocab', 'pad_idx']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        chars = Text(o[0], self.char_vocab.textify(o[0], ''))\n",
    "        words = Text(o[1], self.word_vocab.textify(o[1], ''))\n",
    "        return [chars,words]\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        c,w = t\n",
    "        return [self.reconstruct_one(c, self.char_vocab),self.reconstruct_one(w, self.word_vocab)]\n",
    "    \n",
    "    def reconstruct_one(self, x, vocab):\n",
    "        nonzero_idxs = (x != self.pad_idx).nonzero()\n",
    "        idx_min = 0  #(x != self.pad_idx).nonzero().min()\n",
    "        idx_max = nonzero_idxs.max() if len(nonzero_idxs) > 0 else 0\n",
    "        return Text(x[idx_min:idx_max+1], vocab.textify(x[idx_min:idx_max+1]))\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageMultiList(ImageList):  \n",
    "    def show_xys(self, xs, ys, figsize:Tuple[int,int]=(9,10), **kwargs):\n",
    "        \"Show the `xs` and `ys` on a figure of `figsize`. `kwargs` are passed to the show method.\"\n",
    "        rows = int(math.sqrt(len(xs)))\n",
    "        fig, axs = plt.subplots(rows,rows,figsize=figsize)\n",
    "        for i, ax in enumerate(axs.flatten() if rows > 1 else [axs]):\n",
    "            chars,words = ys[i]\n",
    "            combined = Text([], str(chars) + '\\n\\n' + str(words))\n",
    "            xs[i].show(ax=ax, y=combined, **kwargs)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=MultiSequenceList)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "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 = nn.LayerNorm\n",
    "LayerNorm = partial(nn.LayerNorm, eps=1e-4)  # accomodates mixed precision training\n",
    "# LayerNorm = partial(nn.BatchNorm2d, eps=1e-4)"
   ]
  },
  {
   "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):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\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": [
    "class Decoder(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, layer, N):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src, tgt_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    \"Decoder: self-attn, src-attn, and feed forward\"\n",
    "    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 3)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))  # acts as a weak LM\n",
    "        x = self.sublayer[1](x, lambda x: self.src_attn(x, src, src))\n",
    "        return self.sublayer[2](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 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)  #changed from: -1e9 to accomodate mixed precision  \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": {
    "heading_collapsed": true
   },
   "source": [
    "# Word Arch (bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# class ResnetBase(nn.Module):\n",
    "#     def __init__(self, em_sz):\n",
    "#         super().__init__()\n",
    "        \n",
    "#         slices = {128: -4, 256: -3, 512: -2}\n",
    "#         s = slices[em_sz]\n",
    "\n",
    "#         net = models.resnet34(True)\n",
    "#         modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "#         self.base = nn.Sequential(*modules)\n",
    "        \n",
    "#         self.conv = conv_layer(em_sz,em_sz)\n",
    "#         self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "#     def forward(self, x):\n",
    "#         x = self.base(x)\n",
    "#         x = self.pool(self.conv(x))\n",
    "#         return x.flatten(2,3).permute(0,2,1)\n",
    "\n",
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)                  #32x32 : 256\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.base(x)\n",
    "\n",
    "class Adaptor(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "                        \n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.lin = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv(self.pool(x))\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.lin(x).mul(8)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        #self.src_adapt = src_adapt\n",
    "        \n",
    "        self.w_decoder = decoder\n",
    "        self.w_embed = embed\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt, 101).long()    # CLS tok\n",
    "        mask = parallelogram_mask(tgt.size(-1), 20)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        return self.w_decoder(self.w_embed(tgt), feats, mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "#         nn.Sequential(\n",
    "#             nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #increases gradients on weights by 8!\n",
    "#         ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(60,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# class Img2Seq(nn.Module):\n",
    "#     def __init__(self, img_encoder, transformer):\n",
    "#         super(Img2Seq, self).__init__()\n",
    "#         self.img_enc = img_encoder\n",
    "#         self.transformer = transformer\n",
    "        \n",
    "#     def forward(self, src, tgt):\n",
    "#         feats = self.img_enc(src)\n",
    "#         outs = self.transformer(feats, tgt)\n",
    "#         return self.transformer.generate(outs)\n",
    "\n",
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.adaptor(self.img_enc(src))\n",
    "        outs = self.transformer(feats, tgt)\n",
    "        return self.transformer.generate(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DistilBertMLM(nn.Module):\n",
    "    def __init__(self, vocab, name='distilbert-base-uncased'):\n",
    "        super().__init__()\n",
    "        self.tokenizer = DistilBertTokenizer.from_pretrained(name)\n",
    "        self.model = DistilBertForMaskedLM.from_pretrained(name)\n",
    "        self.vocab = vocab\n",
    "        self.model.train()\n",
    "        \n",
    "    def forward(self, t:tensor):\n",
    "        preds = torch.argmax(t, dim=-1)\n",
    "        lbls = []\n",
    "        for p in preds:\n",
    "            words = decode_spec_tokens(self.vocab.textify(p))\n",
    "            lbls.append(np.array(self.tokenizer.encode(words, add_special_tokens=True)))\n",
    "        input_ids = c_pad(np.array(lbls))\n",
    "        return self.model(input_ids.to(device=device))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer, lm):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        with torch.no_grad():\n",
    "            feats = self.adaptor(self.img_enc(src))\n",
    "            outs = self.transformer(feats, tgt)\n",
    "            outs = self.transformer.generate(outs)\n",
    "        bert_outs = self.lm(inputs_embeds=outs)\n",
    "        return bert_outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor(em_sz, d_model)\n",
    "    lm = DistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')\n",
    "    lm.resize_token_embeddings(len(itos))\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, adaptor, transformer, lm)\n",
    "    learn = Learner(data, net, loss_func=LabelSmoothing(smoothing),\n",
    "                    metrics=[CER(data.y.reconstruct)], callback_fns=[TeacherForce])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# number of trainable parameters\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Char Arch (bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, src_adapt, embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.src_adapt = src_adapt\n",
    "        \n",
    "        self.c_decoder = decoder\n",
    "        self.c_embed = embed\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt, 101).long()\n",
    "        mask = parallelogram_mask(tgt.size(-1), 25)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        return self.c_decoder(self.c_embed(tgt), feats, mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #increases gradients on weights by 8!\n",
    "        ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(100,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        outs = self.transformer(feats, tgt)\n",
    "        return self.transformer.generate(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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]))   #.replace(' ', '')\n",
    "        t = str(fn(targs[i])) #.replace(' ', '')\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
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    learn = Learner(data, net, loss_func=LabelSmoothing(smoothing),\n",
    "                    metrics=[CER(data.y.reconstruct)],\n",
    "                    callback_fns=[TeacherForce, BnFreeze, partial(AccumulateScheduler, n_step=10)]\n",
    "                   )\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Char Arch (itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, src_adapt, embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.src_adapt = src_adapt\n",
    "        \n",
    "        self.c_decoder = decoder\n",
    "        self.c_embed = embed\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt).long()\n",
    "        mask = parallelogram_mask(tgt.size(-1), 25)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        return self.c_decoder(self.c_embed(tgt), feats, mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #increases gradients on weights by 8!\n",
    "        ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(100,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        outs = self.transformer(feats, tgt)\n",
    "        return self.transformer.generate(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    learn = Learner(data, net, loss_func=LabelSmoothing(smoothing),\n",
    "                    metrics=[CER(data.y.reconstruct)], callback_fns=[TeacherForce])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "#     itos = data.vocab.itos\n",
    "#     img_encoder = ResnetBase(em_sz)\n",
    "#     transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "#     net = Img2Seq(img_encoder, transformer)\n",
    "#     learn = Learner(data, net, loss_func=LabelSmoothing(smoothing),\n",
    "#                     metrics=[CER(data.y.reconstruct)],\n",
    "#                     callback_fns=[TeacherForce, BnFreeze, partial(AccumulateScheduler, n_step=3)])\n",
    "#     return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Combo Arch (itos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "length_results = []\n",
    "\n",
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "                \n",
    "        row_t = self.rows(rows)\n",
    "        col_t = self.cols(cols)\n",
    "        \n",
    "        length_results.append(cols.max().item())  # this is only for verification of lengths on cuda error\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "\n",
    "# c_ler = ler[::2]\n",
    "# w_ler = ler[1::2]\n",
    "# max(c_ler), max(w_ler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, src_adapt, c_emb, w_emb, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.src_adapt = src_adapt\n",
    "        \n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        \n",
    "        self.c_embed = c_emb\n",
    "        self.w_embed = w_emb\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.generator(c_outs), self.generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt).long()\n",
    "        w_tgt = rshift(w_tgt).long()\n",
    "        \n",
    "        c_mask = parallelogram_mask(c_tgt.size(-1), 25)\n",
    "        w_mask = parallelogram_mask(w_tgt.size(-1), 10)\n",
    "#         c_mask = subsequent_mask(c_tgt.size(-1)) \n",
    "#         w_mask = subsequent_mask(w_tgt.size(-1)) \n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    c_col_emb = nn.Embedding(100, d_model//2, 0)\n",
    "    w_col_emb = nn.Embedding(60, d_model//2, 0)\n",
    "    \n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8))\n",
    "            # .mul_(8) increases the gradients on these weights by 8!\n",
    "        ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, c_col_emb),   \n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, w_col_emb),\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiCER(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn, itos):\n",
    "        super().__init__(learn)\n",
    "        self.itos = itos\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out = last_output\n",
    "        c_targ, w_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.itos)\n",
    "        w_error,_    = cer(w_out, w_targ, self.itos)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.total = 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, [self.c_errors/self.total, self.w_errors/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing), callback_fns=[TeacherForce])\n",
    "    learn.callbacks.append(MultiCER(learn, itos))\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Fixup Combo Arch (bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def conv3x3(in_planes, out_planes, stride=1):\n",
    "    \"\"\"3x3 convolution with padding\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                     padding=1, bias=False)\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
    "\n",
    "class FixupBasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(FixupBasicBlock, self).__init__()\n",
    "        # Both self.conv1 and self.downsample layers downsample the input when stride != 1\n",
    "        self.bias1a = nn.Parameter(torch.zeros(1))\n",
    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
    "        self.bias1b = nn.Parameter(torch.zeros(1))\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.bias2a = nn.Parameter(torch.zeros(1))\n",
    "        self.conv2 = conv3x3(planes, planes)\n",
    "        self.scale = nn.Parameter(torch.ones(1))\n",
    "        self.bias2b = nn.Parameter(torch.zeros(1))\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x + self.bias1a)\n",
    "        out = self.relu(out + self.bias1b)\n",
    "\n",
    "        out = self.conv2(out + self.bias2a)\n",
    "        out = out * self.scale + self.bias2b\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x + self.bias1a)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n",
    "    \n",
    "class FixupBottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None):\n",
    "        super(FixupBottleneck, self).__init__()\n",
    "        # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n",
    "        self.bias1a = nn.Parameter(torch.zeros(1))\n",
    "        self.conv1 = conv1x1(inplanes, planes)\n",
    "        self.bias1b = nn.Parameter(torch.zeros(1))\n",
    "        self.bias2a = nn.Parameter(torch.zeros(1))\n",
    "        self.conv2 = conv3x3(planes, planes, stride)\n",
    "        self.bias2b = nn.Parameter(torch.zeros(1))\n",
    "        self.bias3a = nn.Parameter(torch.zeros(1))\n",
    "        self.conv3 = conv1x1(planes, planes * self.expansion)\n",
    "        self.scale = nn.Parameter(torch.ones(1))\n",
    "        self.bias3b = nn.Parameter(torch.zeros(1))\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.conv1(x + self.bias1a)\n",
    "        out = self.relu(out + self.bias1b)\n",
    "\n",
    "        out = self.conv2(out + self.bias2a)\n",
    "        out = self.relu(out + self.bias2b)\n",
    "\n",
    "        out = self.conv3(out + self.bias3a)\n",
    "        out = out * self.scale + self.bias3b\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x + self.bias1a)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FixupResNet(nn.Module):\n",
    "    def __init__(self, block, layers, d_model=512):\n",
    "        super().__init__()\n",
    "        self.num_layers = sum(layers)\n",
    "        self.inplanes = 64\n",
    "        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "        self.bias1 = nn.Parameter(torch.zeros(1))\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
    "        \n",
    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=1)\n",
    "        #self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.maxpool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        self.bias2 = nn.Parameter(torch.zeros(1))\n",
    "        self.fc = nn.Linear(512 * block.expansion, d_model)\n",
    "        self.apply_init()\n",
    "\n",
    "    def _make_layer(self, block, planes, blocks, stride=1):\n",
    "        downsample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = conv1x1(self.inplanes, planes * block.expansion, stride)\n",
    "\n",
    "        layers = []\n",
    "        layers.append(block(self.inplanes, planes, stride, downsample))\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for _ in range(1, blocks):\n",
    "            layers.append(block(self.inplanes, planes))\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.relu(x + self.bias1)\n",
    "        x = self.maxpool(x)\n",
    "\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "\n",
    "        x = self.maxpool(x) #self.avgpool(x)\n",
    "        #x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x.flatten(2,3).permute(0,2,1) + self.bias2)\n",
    "        return x\n",
    "    \n",
    "    def apply_init(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, FixupBasicBlock):\n",
    "                nn.init.normal_(m.conv1.weight, mean=0, std=np.sqrt(\n",
    "                    2 / (m.conv1.weight.shape[0] * np.prod(m.conv1.weight.shape[2:]))) * self.num_layers ** (-0.5))\n",
    "                nn.init.constant_(m.conv2.weight, 0)\n",
    "                if m.downsample is not None:\n",
    "                    nn.init.normal_(m.downsample.weight, mean=0, std=np.sqrt(\n",
    "                        2 / (m.downsample.weight.shape[0] * np.prod(m.downsample.weight.shape[2:]))))\n",
    "            elif isinstance(m, FixupBottleneck):\n",
    "                nn.init.normal_(m.conv1.weight, mean=0, std=np.sqrt(\n",
    "                    2 / (m.conv1.weight.shape[0] * np.prod(m.conv1.weight.shape[2:]))) * self.num_layers ** (-0.25))\n",
    "                nn.init.normal_(m.conv2.weight, mean=0, std=np.sqrt(\n",
    "                    2 / (m.conv2.weight.shape[0] * np.prod(m.conv2.weight.shape[2:]))) * self.num_layers ** (-0.25))\n",
    "                nn.init.constant_(m.conv3.weight, 0)\n",
    "                if m.downsample is not None:\n",
    "                    nn.init.normal_(m.downsample.weight, mean=0, std=np.sqrt(\n",
    "                        2 / (m.downsample.weight.shape[0] * np.prod(m.downsample.weight.shape[2:]))))\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                nn.init.constant_(m.weight, 0)\n",
    "                nn.init.constant_(m.bias, 0)              \n",
    "\n",
    "def fixup_resnet34(**kwargs):\n",
    "    \"\"\"Constructs a Fixup-ResNet-34 model.\n",
    "    \"\"\"\n",
    "    model = FixupResNet(FixupBasicBlock, [3, 4, 6, 3], **kwargs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, c_emb, w_emb, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        \n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        \n",
    "        self.c_embed = c_emb\n",
    "        self.w_embed = w_emb\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.generator(c_outs), self.generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt, 101).long()\n",
    "        w_tgt = rshift(w_tgt, 101).long()\n",
    "        \n",
    "        c_mask = parallelogram_mask(c_tgt.size(-1), 25)\n",
    "        w_mask = parallelogram_mask(w_tgt.size(-1), 10)\n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(100, d_model//2, 0)),  #char\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(60,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiCER(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn):\n",
    "        super().__init__(learn)\n",
    "        self.recon = learn.data.y.reconstruct_one\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out = last_output\n",
    "        c_targ, w_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.recon)\n",
    "        w_error,_    = cer(w_out, w_targ, self.recon)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.total = 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, [self.c_errors/self.total, self.w_errors/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiAccumulateScheduler(AccumulateScheduler):\n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        \"accumulate samples and batches\"\n",
    "        self.acc_samples += last_input[0].shape[0]\n",
    "        self.acc_batches += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = fixup_resnet34(d_model=d_model) #ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing),\n",
    "                    callback_fns=[TeacherForce, MultiCER, partial(MultiAccumulateScheduler, n_step=5)])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Integrated Combo Arch (bert_tok)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# original version w/ saved model: test_combo_bert_tok_word\n",
    "\n",
    "# class ResnetBase(nn.Module):\n",
    "#     def __init__(self, em_sz):\n",
    "#         super().__init__()\n",
    "        \n",
    "#         slices = {128: -4, 256: -3, 512: -2}\n",
    "#         s = slices[em_sz]\n",
    "\n",
    "#         net = models.resnet34(True)\n",
    "#         modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "#         self.base = nn.Sequential(*modules)\n",
    "                \n",
    "#         self.conv = conv_layer(em_sz,em_sz)\n",
    "#         self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "#     def forward(self, x):\n",
    "#         x = self.base(x)\n",
    "#         x = self.pool(self.conv(x))\n",
    "#         return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)                  #32x32 : 256\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.base(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Adaptor(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "                        \n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.lin = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv(self.pool(x))\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.lin(x).mul(8)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, c_emb, w_emb, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        #self.src_adapt = src_adapt\n",
    "        \n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        \n",
    "        self.c_embed = c_emb\n",
    "        self.w_embed = w_emb\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "        #return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.generator(c_outs), self.generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt, 101).long()\n",
    "        w_tgt = rshift(w_tgt, 101).long()\n",
    "        \n",
    "        c_mask = parallelogram_mask(c_tgt.size(-1), 20)   # char needs word context\n",
    "        w_mask = parallelogram_mask(w_tgt.size(-1), 20)   # word needs sentence context\n",
    "#         c_mask = subsequent_mask(c_tgt.size(-1)) \n",
    "#         w_mask = subsequent_mask(w_tgt.size(-1)) \n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "#         nn.Sequential(\n",
    "#             nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #increases gradients on weights by 8!\n",
    "#         ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(100, d_model//2, 0)),  #char\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(60,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.adaptor(self.img_enc(src))\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiCER(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn):\n",
    "        super().__init__(learn)\n",
    "        self.recon = learn.data.y.reconstruct_one\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out = last_output\n",
    "        c_targ, w_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.recon)\n",
    "        w_error,_    = cer(w_out, w_targ, self.recon)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.total = 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, [self.c_errors/self.total, self.w_errors/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    itos = data.vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, adaptor, transformer)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing),\n",
    "                    callback_fns=[TeacherForce, MultiCER])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# number of trainable parameters\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('hw_combo_bt', strict=False); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split(lambda m: [m.adaptor,m.lm]); None\n",
    "len(learn.layer_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze()\n",
    "# learn.unfreeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# lrs = slice(1e-4,1e-3)\n",
    "lrs = 2e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(2, max_lr=lrs, callbacks=[SaveModelCallback(learn, name='hw_word_bt')])\n",
    "\n",
    "# small, bs: 50\n",
    "\n",
    "# maxpool(16,None), original adaptor\n",
    "# 52.609947\t48.825050\t0.691070\t0.818217\t15:10   88M; freeze(), 1e-3   \n",
    "# 7.000093\t8.394917\t0.048174\t0.052259\t14:23   unfreeze(), (1e-4,1e-3)   *sm_combo_bt\n",
    "# chars:    1.73652   .07256\n",
    "# words:    0.89602   .03005\n",
    "\n",
    "# adaptor from multi-resolution; remove src_adapt\n",
    "# 58.962891\t53.468037\t0.750514\t0.815653\t11:11   91M; freeze(), 1e-3\n",
    "# 8.015573\t9.585506\t0.054404\t0.065465\t12:05   unfreeze(), (1e-4,1e-3)   *sm_combo_bt2\n",
    "# chars:    1.18025   .09011\n",
    "# words:    0.96089   .05492\n",
    "\n",
    "# adaptor: conv_layer(em_sz,em_sz,(2,1)) + lin/mul(8)\n",
    "# 58.688446\t54.525970\t0.677837\t0.841201\t12:07   86M; freeze(), 1e-3\n",
    "# 7.844480\t9.105121\t0.054534\t0.064099\t12:44   unfreeze(), (1e-4,1e-3)   *sm_combo_bt3\n",
    "\n",
    "# preload pg_combo_bt4; maxpool(16,None) -> conv\n",
    "# 54.845032\t50.746662\t0.684326\t0.817202\t12:39   86M; freeze(), 1e-3\n",
    "# 6.565070\t8.298703\t0.048370\t0.050501\t13:18   unfreeze(), (1e-4,1e-3)   *sm_combo_bt4\n",
    "# chars:    1.22381   .09270\n",
    "# words:    0.94450   .06052\n",
    "\n",
    "\n",
    "# pg, bs: 10\n",
    "\n",
    "# conv -> maxpool(16,None), original + lin/mul(8)\n",
    "# 1794.547241\t1622.572144\t1.033109\t0.881144\t02:04\n",
    "# 1015.301025\t1080.083496\t0.564250\t0.680124\t02:09   *pg_combo_bt\n",
    "\n",
    "# maxpool(16,None) -> conv\n",
    "# 1776.183105\t1585.937256\t1.065259\t0.884869\t02:05\n",
    "# 1010.742371\t1077.592163\t0.563446\t0.678190\t02:08   *pg_combo_bt4\n",
    "\n",
    "# adaptor from multi-resolution\n",
    "# 1789.552856\t1590.777710\t1.138431\t0.911390\t02:00\n",
    "# 1027.244995\t1084.395264\t0.569967\t0.679046\t02:03   *pg_combo_bt2\n",
    "\n",
    "# adaptor: conv(em_sz,em_sz,(2,1))\n",
    "# 1791.784668\t1596.507812\t1.053432\t0.880950\t02:04\n",
    "# 1043.639893\t1096.874268\t0.565953\t0.697545\t02:07   *pg_combo_bt3\n",
    "\n",
    "# font_mix, bs: 10, preload 'sm_combo_bt4'\n",
    "\n",
    "# 75.363541\t62.632046\t0.038649\t0.099920\t2:56:03   2epochs (1e-4,1e-3)   *font_combo_bt\n",
    "# 42.600681\t39.815552\t0.015061\t0.045635\t3:04:19   2epochs (1e-5,1e-4)   *font_combo_bt2\n",
    "# chars:    41.8378   .00182\n",
    "# words:    83.1500   .00275\n",
    "\n",
    "# handwriting_mix, bs: 10, preload 'font_combo_bt2'\n",
    "\n",
    "# 11.269932\t10.436209\t0.004604\t0.002646\t35:18   5cycle(1e-4,1e-3)    *hw_combo_bt\n",
    "# chars:    85.4927   .00315\n",
    "# words:    0.61776   .00065\n",
    "\n",
    "#TEST pg\n",
    "# chars:    138.678   .04029\n",
    "# words:    69.4958   .05652\n",
    "#TEST upload\n",
    "# chars:    128.604   .23494\n",
    "# words:    56.7417   .29080\n",
    "\n",
    "# word only, preload 'hw_combo_bt'\n",
    "# 4.775697\t5.477561\t0.000647\t18:20    2cycle(2e-5)    *hw_word_bt\n",
    "#     pg:    68.1078   .05712\n",
    "# upload:    50.6502   .29197"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# word char (separate vocabs) Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab, num_cols, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "\n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, src_adapt, c_emb, w_emb, c_gen, w_gen):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.src_adapt = src_adapt\n",
    "        \n",
    "        self.w_decoder = w_dec\n",
    "        self.c_decoder = c_dec\n",
    "        \n",
    "        self.w_embed = w_emb\n",
    "        self.c_embed = c_emb\n",
    "        self.w_generator = w_gen\n",
    "        self.c_generator = c_gen\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.c_generator(c_outs), self.w_generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt).long()\n",
    "        w_tgt = rshift(w_tgt).long()\n",
    "        \n",
    "        c_mask = parallelogram_mask(c_tgt.size(-1), 25)\n",
    "        w_mask = parallelogram_mask(w_tgt.size(-1), 10)\n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(c_vocab, w_vocab, d_model, em_sz, N=4, drops=0.1, heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    \n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential( nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) ),  # multiplies gradients by 8!\n",
    "        LearnedPositionalEmbeddings(d_model, c_vocab, 100),   \n",
    "        LearnedPositionalEmbeddings(d_model, w_vocab, 60),\n",
    "        nn.Linear(d_model, c_vocab),\n",
    "        nn.Linear(d_model, w_vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CERMetric(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn, c_itos, w_itos):\n",
    "        super().__init__(learn)\n",
    "        self.c_itos = c_itos\n",
    "        self.w_itos = w_itos\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out = last_output\n",
    "        c_targ, w_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.c_itos)\n",
    "        w_error,_    = cer(w_out, w_targ, self.w_itos)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.total = 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, [self.c_errors/self.total, self.w_errors/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    c_vocab = data.train_dl.dl.dataset.y.char_vocab.itos\n",
    "    w_vocab = data.train_dl.dl.dataset.y.word_vocab.itos\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(c_vocab), len(w_vocab), d_model, em_sz, N, drops, attn_type, heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing), callback_fns=[TeacherForce])\n",
    "    learn.callbacks.append(CERMetric(learn, c_vocab, w_vocab))\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# word char (combo vocab) w/ DistilBert LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "length_results = []\n",
    "\n",
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        \n",
    "        row_t = self.rows(rows)\n",
    "        col_t = self.cols(cols)\n",
    "        \n",
    "        length_results.append(cols.max().item())  # this is only for verification of lengths on cuda error\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "\n",
    "# c_ler = ler[::2]\n",
    "# w_ler = ler[1::2]\n",
    "# max(c_ler), max(w_ler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, src_adapt, c_emb, w_emb, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.src_adapt = src_adapt\n",
    "        \n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        \n",
    "        self.c_embed = c_emb\n",
    "        self.w_embed = w_emb\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encode(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.generator(c_outs), self.generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt,_ = mix_tgt\n",
    "        c_tgt = rshift(c_tgt).long()\n",
    "        w_tgt = rshift(w_tgt).long()\n",
    "        \n",
    "        c_mask = parallelogram_mask(c_tgt.size(-1), 25)\n",
    "        w_mask = parallelogram_mask(w_tgt.size(-1), 10)\n",
    "#         c_mask = subsequent_mask(c_tgt.size(-1)) \n",
    "#         w_mask = subsequent_mask(w_tgt.size(-1)) \n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    #pos = PositionalEncoding(d_model, drops, 2000)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    c_col_emb = nn.Embedding(100, d_model//2, 0)\n",
    "    w_col_emb = nn.Embedding(60, d_model//2, 0)\n",
    "    \n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8))\n",
    "            # .mul_(8) increases the gradients on these weights by 8!\n",
    "        ),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, c_col_emb),   \n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, w_col_emb),\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DistilBertMLM(nn.Module):\n",
    "    def __init__(self, vocab, name='distilbert-base-uncased'):\n",
    "        super().__init__()\n",
    "        self.tokenizer = DistilBertTokenizer.from_pretrained(name)\n",
    "        self.model = DistilBertForMaskedLM.from_pretrained(name)\n",
    "        self.vocab = vocab\n",
    "        self.model.train()\n",
    "        \n",
    "    def forward(self, t:tensor):\n",
    "        preds = torch.argmax(t, dim=-1)\n",
    "        lbls = []\n",
    "        for p in preds:\n",
    "            words = decode_spec_tokens(self.vocab.textify(p))\n",
    "            lbls.append(np.array(self.tokenizer.encode(words, add_special_tokens=True)))\n",
    "        input_ids = c_pad(np.array(lbls))\n",
    "        return self.model(input_ids.to(device=device))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, lm):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.img_enc(src)\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        bert_outs = self.lm(outs[1])\n",
    "        return (*outs, bert_outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CERMetric(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn):\n",
    "        super().__init__(learn)\n",
    "        self.itos = learn.data.vocab.itos\n",
    "        self.b_itos = learn.data.bert_vocab.itos\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word', 'bert'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        (c_out, w_out),b_out = last_output\n",
    "        c_targ, w_targ,b_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.itos)\n",
    "        w_error,_    = cer(w_out, w_targ, self.itos)\n",
    "        b_error,_    = cer(b_out, b_targ, self.b_itos)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.b_errors += b_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.b_errors, self.total = 0, 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        mets = [self.c_errors/self.total, self.w_errors/self.total, self.b_errors/self.total]\n",
    "        return add_metrics(last_metrics, mets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ, b_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        bl = loss(pred[2], b_targ)\n",
    "        print(f'char loss: {cl}  word_loss: {wl}  bert_loss:{bl}')\n",
    "        return cl + wl + bl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    vocab = data.vocab\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(vocab.itos), d_model, em_sz, N, drops, attn_type, heads)\n",
    "    lm = DistilBertMLM(vocab)\n",
    "    net = Img2Seq(img_encoder, transformer, lm)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing), callback_fns=[TeacherForce])\n",
    "    learn.callbacks.append(CERMetric(learn))\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Modified Feature Pyramid Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## word only modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordTransformer(nn.Module):\n",
    "    def __init__(self, encoder, decoder, embed, generator):\n",
    "        super(WordTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        \n",
    "        self.w_decoder = decoder\n",
    "        self.w_embed = embed\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt, 101).long()\n",
    "        mask = subsequent_mask(tgt.size(-1)) \n",
    "\n",
    "        feats = self.encoder(src)\n",
    "        return self.w_decoder(self.w_embed(tgt), feats, mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_word_model(vocab, d_model, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "    \n",
    "    model = WordTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), 2),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(60,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# word only\n",
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.adaptor(self.img_enc(src))\n",
    "        word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# word model only\n",
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    voc_len = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase()\n",
    "    load_pre(img_encoder)\n",
    "    adaptor = FPN(em_sz, d_model)\n",
    "    transformer = make_word_model(voc_len, d_model, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, adaptor, transformer)\n",
    "    learn = Learner(data, net, loss_func=LabelSmoothing(smoothing),\n",
    "                    metrics=[CER(data.y.reconstruct)], callback_fns=[TeacherForce])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Combo modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LearnedPositionalEmbeddings(nn.Module):\n",
    "    def __init__(self, d_model, v_embed, row_embed, col_embed, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = v_embed #nn.Embedding(vocab, d_model, 0)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        self.rows = row_embed #nn.Embedding(15, d_model//2, 0)\n",
    "        self.cols = col_embed #nn.Embedding(num_cols, d_model//2, 0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        rows,cols = self.encode_spatial_positions(x)      \n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = self.cols(torch.clamp(cols, max=self.cols.num_embeddings-1))  # clamp to max column value\n",
    "        pos_enc = torch.cat((row_t, col_t), dim=-1)\n",
    "                \n",
    "        x = self.embed(x)\n",
    "        x = (x + pos_enc) * math.sqrt(self.d_model)\n",
    "        return self.dropout(x)\n",
    "    \n",
    "    def encode_spatial_positions(self, x):\n",
    "        rows,cols = torch.zeros_like(x),torch.zeros_like(x)\n",
    "        for ii,batch in enumerate(x.unbind()):\n",
    "            nls = torch.nonzero(batch==self.nl_tok).flatten()\n",
    "            last = torch.nonzero(batch).flatten()[-1][None]\n",
    "            splits = torch.cat([nls,last])\n",
    "\n",
    "            p=0\n",
    "            for i,n in enumerate(splits, start=1):\n",
    "                rows[ii,p:n+1] = i\n",
    "                cols[ii,p:n+1] = torch.arange(1,n-p+2)\n",
    "                p = n+1\n",
    "        return rows,cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from torchvision.models.resnet import BasicBlock#, Bottleneck\n",
    "from torchvision.models.utils import load_state_dict_from_url"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(models.ResNet):\n",
    "    def __init__(self, block=BasicBlock, layers=[3,4,6,0]):\n",
    "        super().__init__(block, layers)\n",
    "        self.avgpool = None\n",
    "        self.fc = None\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Bottom-up - resnet backbone\n",
    "        c1 = self.maxpool(self.relu(self.bn1(self.conv1(x))))  #256 - 32\n",
    "        c2 = self.layer1(c1)  #128 - 64  # more spatial integrity (resolution)\n",
    "        c3 = self.layer2(c2)  #64 - 128\n",
    "        c4 = self.layer3(c3)  #32 - 256\n",
    "        c5 = self.layer4(c4)  #16 - 512  # more semantic information\n",
    "        return (c5,c4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FPN(nn.Module):\n",
    "    def __init__(self, em_sz=256, d_model=512):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.smooth = nn.Conv2d(em_sz, em_sz, kernel_size=3, stride=1, padding=1)\n",
    "        self.latlayer = nn.Conv2d(512, em_sz, kernel_size=1, stride=1, padding=0)\n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "    \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_uniform_(m.weight)  #not followed by activation functions\n",
    "                        \n",
    "    def _upsample_add(self, x, y):\n",
    "        _,_,H,W = y.size()\n",
    "        return F.interpolate(x, size=(H,W), mode='bilinear', align_corners=False) + y\n",
    "        #bilinear - supports arbitrary output sizes\n",
    "\n",
    "    def forward(self, x):   \n",
    "        out = self.smooth(self._upsample_add(self.latlayer(x[0]), x[1]))\n",
    "        out = out.flatten(2,3).mul(2)\n",
    "        out = self.linear(out.permute(0,2,1))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class WordCharTransformer(nn.Module):\n",
    "    def __init__(self, encoder, c_dec, w_dec, c_embed, w_embed, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        \n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        \n",
    "        self.c_embed = c_embed\n",
    "        self.w_embed = w_embed\n",
    "        \n",
    "        self.generator = generator\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "\n",
    "        feats = self.encoder(src)\n",
    "        char_outs = self.c_decoder(self.c_embed(c_tgt), feats, c_mask)\n",
    "        word_outs = self.w_decoder(self.w_embed(w_tgt), feats, w_mask)\n",
    "\n",
    "        return char_outs, word_outs\n",
    "    \n",
    "    def generate(self, c_outs, w_outs):\n",
    "        return self.generator(c_outs), self.generator(w_outs)\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt, 101).long()\n",
    "        w_tgt = rshift(w_tgt, 101).long()\n",
    "\n",
    "        c_mask = subsequent_mask(c_tgt.size(-1)) \n",
    "        w_mask = subsequent_mask(w_tgt.size(-1)) \n",
    "        return c_tgt,w_tgt,c_mask,w_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    \n",
    "    v_embed = nn.Embedding(vocab, d_model, 0)\n",
    "    row_emb = nn.Embedding(15, d_model//2, 0)\n",
    "\n",
    "    model = WordCharTransformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), 2),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(100, d_model//2, 0)),  #char\n",
    "        LearnedPositionalEmbeddings(d_model, v_embed, row_emb, nn.Embedding(60,  d_model//2, 0)),  #word\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        feats = self.adaptor(self.img_enc(src))\n",
    "        char_outs, word_outs = self.transformer(feats, tgt)\n",
    "        outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiCER(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn):\n",
    "        super().__init__(learn)\n",
    "        self.recon = learn.data.y.reconstruct_one\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.learn.recorder.add_metric_names(['char', 'word'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out = last_output\n",
    "        c_targ, w_targ = last_target\n",
    "        c_error,size = cer(c_out, c_targ, self.recon)\n",
    "        w_error,_    = cer(w_out, w_targ, self.recon)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.total = 0, 0, 0\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, [self.c_errors/self.total, self.w_errors/self.total])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(MultiLabelSmoothing, self).__init__()\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, c_targ, w_targ):\n",
    "        loss = LabelSmoothing(self.smoothing)\n",
    "        cl = loss(pred[0], c_targ)\n",
    "        wl = loss(pred[1], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def load_pre(model):\n",
    "    url = 'https://download.pytorch.org/models/resnet34-333f7ec4.pth'\n",
    "    #url = 'https://download.pytorch.org/models/resnet50-19c8e357.pth'\n",
    "    model.load_state_dict(load_state_dict_from_url(url, progress=True), strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    voc_len = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase()\n",
    "    load_pre(img_encoder)\n",
    "    adaptor = FPN(em_sz, d_model)\n",
    "    transformer = make_full_model(voc_len, d_model, N, drops, heads)\n",
    "    net = Img2Seq(img_encoder, adaptor, transformer)\n",
    "    learn = Learner(data, net, loss_func=MultiLabelSmoothing(smoothing), callback_fns=[TeacherForce, MultiCER])\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0.1, heads=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# number of trainable parameters\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split(lambda m: [m.adaptor, m.transformer]); None\n",
    "len(learn.layer_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def greedy_decode(src, model, seq_len, kind='char', bos_tok=1):\n",
    "    model.eval()\n",
    "    tfmr = model.transformer\n",
    "    adaptor = model.adaptor\n",
    "    img_enc = model.img_enc\n",
    "    \n",
    "    decoder = tfmr.c_decoder if kind=='char' else tfmr.w_decoder\n",
    "    embed = tfmr.c_embed if kind=='char' else tfmr.w_embed\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        feats = tfmr.encoder(adaptor(img_enc(src)))\n",
    "        bs = src.size(0)\n",
    "        tgt = torch.zeros((bs,1), dtype=torch.long, device=device) + bos_tok\n",
    "\n",
    "        res = []\n",
    "        for i in progress_bar(range(seq_len)):\n",
    "            mask = subsequent_mask(tgt.size(-1))            \n",
    "            dec_outs = decoder(embed(tgt), feats, mask)\n",
    "            prob = tfmr.generator(dec_outs[:,-1])\n",
    "            res.append(prob)\n",
    "            pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "            if (pred==0).all(): break\n",
    "            tgt = torch.cat([tgt,pred], dim=-1)\n",
    "        out = torch.stack(res).transpose(1,0).contiguous()\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## View Model Telemetry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FullStats(HookCallback):\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        self.modules = [m for m in flatten_model(self.learn.model) if hasattr(m, 'weight')]\n",
    "        self.g_hooks = Hooks(self.modules, self.g_hook, is_forward=False)\n",
    "        self.a_hooks = Hooks(self.modules, self.a_hook)\n",
    "        self.grads,self.acts = [],[]\n",
    "\n",
    "    def g_hook(self, m:nn.Module, i:Tensors, o:Tensors)->Tuple[Rank0Tensor,Rank0Tensor]:\n",
    "        oo = next(o)\n",
    "        return oo.mean().item(),oo.std().item()\n",
    "    \n",
    "    def a_hook(self, m:nn.Module, i:Tensors, o:Tensors)->Tuple[Rank0Tensor,Rank0Tensor]:\n",
    "        return o.mean().item(),o.std().item()\n",
    "\n",
    "    def on_batch_end(self, train, **kwargs):\n",
    "        if train:\n",
    "            self.acts.append(self.a_hooks.stored)\n",
    "            self.grads.append(self.g_hooks.stored)\n",
    "            \n",
    "    def on_train_end(self, **kwargs):\n",
    "        self.a_hooks.remove()\n",
    "        self.g_hooks.remove()\n",
    "        self.acts = tensor(self.acts).permute(2,1,0)\n",
    "        self.grads = tensor(self.grads).permute(2,1,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit(1, 1e-5, callbacks=[FullStats(learn)])#, StopAfterNBatches(n_batches=2)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "acts,grads = learn.full_stats.acts, learn.full_stats.grads\n",
    "acts.shape,grads.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "names=[]\n",
    "for name, param in learn.model.named_parameters():\n",
    "    if name.endswith('weight'):\n",
    "        names.append(name)\n",
    "\n",
    "names.insert(193, 'transformer.w_embed.embed.weight')\n",
    "names.insert(194, 'transformer.w_embed.rows.weight')\n",
    "len(names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# :64     img_enc\n",
    "# 64:67   adaptor\n",
    "# 67:84   encoder\n",
    "# 84:137  c_decoder\n",
    "# 137:190 w_decoder\n",
    "# 190:    embeddings/generator\n",
    "\n",
    "plt.figure(figsize=(20,10))\n",
    "for l in acts[1,137:190]:\n",
    "    plt.plot(l)\n",
    "plt.legend(names[137:190])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "avg_act_stds_by_layer = acts[1,:].mean(-1)\n",
    "avg_grad_stds_by_layer = grads[1,:].mean(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.plot(avg_act_stds_by_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.plot(avg_grad_stds_by_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for (i,mod),a,g in zip(enumerate(names), avg_act_stds_by_layer, avg_grad_stds_by_layer):\n",
    "    mod_name = str(mod).split('(')[0]\n",
    "    print(f\"{str(i).ljust(3)} {mod_name.ljust(60)} \\\n",
    "            {str(round(a.item(),5)).ljust(6)} {str(round(g.item(),5)).ljust(6)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Last batch activations by layer\n",
    "\n",
    "for (i,mod),m,s in zip(enumerate(names), acts[0,:,-1], acts[1,:,-1]):\n",
    "    mod_name = str(mod).split('(')[0]\n",
    "    print(f\"{str(i).ljust(3)} {mod_name.ljust(50)} \\\n",
    "            {str(round(m.item(),5)).ljust(6)}  {str(round(s.item(),5)).ljust(6)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Resolution Arch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Image Encoders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Resnet34"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageEncoder(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Xception"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from fastai.vision.models.cadene_models import xception_cadene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SeparableConv2d(nn.Module):\n",
    "    def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False):\n",
    "        super(SeparableConv2d,self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(in_channels,in_channels,kernel_size,stride,\n",
    "                               padding,dilation,groups=in_channels,bias=bias)\n",
    "        self.pointwise = nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias)\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.pointwise(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self,in_filters,out_filters,reps,strides=1,start_with_relu=True,grow_first=True):\n",
    "        super(Block, self).__init__()\n",
    "\n",
    "        if out_filters != in_filters or strides!=1:\n",
    "            self.skip = nn.Conv2d(in_filters,out_filters,1,stride=strides, bias=False)\n",
    "            self.skipbn = nn.BatchNorm2d(out_filters)\n",
    "        else:\n",
    "            self.skip=None\n",
    "\n",
    "        rep=[]\n",
    "\n",
    "        filters=in_filters\n",
    "        if grow_first:\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(out_filters))\n",
    "            filters = out_filters\n",
    "\n",
    "        for i in range(reps-1):\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(filters,filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(filters))\n",
    "\n",
    "        if not grow_first:\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(out_filters))\n",
    "\n",
    "        if not start_with_relu:\n",
    "            rep = rep[1:]\n",
    "        else:\n",
    "            rep[0] = nn.ReLU(inplace=False)\n",
    "\n",
    "        if strides != 1:\n",
    "            rep.append(nn.MaxPool2d(3,strides,1))\n",
    "        self.rep = nn.Sequential(*rep)\n",
    "\n",
    "    def forward(self,inp):\n",
    "        x = self.rep(inp)\n",
    "\n",
    "        if self.skip is not None:\n",
    "            skip = self.skip(inp)\n",
    "            skip = self.skipbn(skip)\n",
    "        else:\n",
    "            skip = inp\n",
    "\n",
    "        x+=skip\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageEncoder(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {728:-6, 1024:-5, 2048:-1}\n",
    "        #728, 4, 32\n",
    "        #1024, 2, 16\n",
    "        #2048, 2, 16\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = xception_cadene(True)\n",
    "        self.base = net[:s]\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### SEResnet50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from fastai.vision.models.cadene_models import se_resnet50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageEncoder(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {512:-4, 1024:-3, 2048:-2}\n",
    "        #512, 8, 64\n",
    "        #1024, 4, 32\n",
    "        #2048, 2, 16\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = se_resnet50(True)\n",
    "        modules = list(net.children())[:s]    #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Previous"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)                  #16x16\n",
    "        self.c1 = conv_layer(d_model, d_model, stride=2)     #8x8\n",
    "        self.c2 = conv_layer(d_model, d_model, stride=2)     #4x4\n",
    "        self.c3 = conv_layer(d_model, d_model, stride=2)     #2x2\n",
    "                \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x1 = self.c1(x)\n",
    "        x2 = self.c2(x1)\n",
    "        x3 = self.c3(x2)\n",
    "        cat = torch.cat([x.flatten(2,3),x1.flatten(2,3),x2.flatten(2,3),x3.flatten(2,3)], dim=-1)\n",
    "        return cat.permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):    \n",
    "    def __init__(self, encoder, decoder, src_adapt, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_adapt = src_adapt\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), self.embed(tgt), tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src, tgt_mask)\n",
    "    \n",
    "    def embed(self, tgt):\n",
    "        return self.tgt_embed(tgt)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0.2, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #, pos\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "            out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    voc_len = len(data.train_ds.y.vocab.itos)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER(data.y.reconstruct)], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=4, drops=0, attn_type='multi', heads=8)\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_mix_res9')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## SentencePiece Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)                  #32x32 : 256\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.base(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Adaptor(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.c1 = conv_layer(em_sz,   d_model, stride=2)     #16x16\n",
    "        self.c2 = conv_layer(d_model, d_model, stride=2)     #8x8\n",
    "        self.c3 = conv_layer(d_model, d_model, stride=2)     #4x4\n",
    "\n",
    "    def forward(self, x):   \n",
    "        x1 = self.c1(x)\n",
    "        x2 = self.c2(x1)\n",
    "        x3 = self.c3(x2)\n",
    "        cat = torch.cat([x1.flatten(2,3),x2.flatten(2,3),x3.flatten(2,3)], dim=-1)\n",
    "        return cat.permute(0,2,1).mul(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.embed = embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt)  # add bos token\n",
    "        return self.decode(self.encoder(src), tgt)\n",
    "        \n",
    "    def decode(self, src, tgt):\n",
    "        return self.decoder(self.embed(tgt), src, subsequent_mask(tgt.size(-1)))\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.1, heads=8):\n",
    "    c = deepcopy\n",
    "    attn = MultiHeadedAttention(d_model, heads)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "    \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "            \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, adaptor, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.adaptor = adaptor\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encoder(self.adaptor(self.img_enc(src)))\n",
    "                tgt = torch.ones((src.size(0), 1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    dec_outs = self.transformer.decode(feats, tgt)\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.adaptor(self.img_enc(src))\n",
    "            dec_outs = self.transformer(feats, tgt)\n",
    "            out = self.transformer.generate(dec_outs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, em_sz, d_model, N=4, drops=0.1, heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor(em_sz, d_model)\n",
    "    transformer = make_full_model(len(data.vocab), d_model, N=N, drops=drops, heads=heads)  # len(data.vocab)\n",
    "    #transformer = make_full_model(len(data.y.vocab.itos), d_model, N=N, drops=drops, heads=heads)  # len(data.vocab)\n",
    "    net = Img2Seq(img_encoder, adaptor, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER(data.y.textify)], callback_fns=[TeacherForce])  # data.y(.vocab).textify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 256, 512, N=6, drops=0.1, heads=8)\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# number of trainable parameters\n",
    "sum(p.numel() for p in learn.model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Adjust model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_char.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "model_sd = OrderedDict()\n",
    "\n",
    "for k,v in sd['model'].items():\n",
    "    if k.startswith('transformer.decoder'): pass\n",
    "    elif k.startswith('transformer.embed'): pass\n",
    "    elif k.startswith('transformer.generator'): pass\n",
    "    else: model_sd[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(model_sd, strict=False)"
   ]
  },
  {
   "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('combo_spm30k_2'); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split(lambda m: [m.adaptor, m.transformer, m.transformer.decoder]); None\n",
    "len(learn.layer_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.freeze()\n",
    "learn.unfreeze()\n",
    "\n",
    "# learn.model.img_enc.base[-1][-1].conv1.weight.requires_grad\n",
    "# learn.model.adaptor.c1[0].weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrs = slice(1e-6,1e-4)\n",
    "# lrs = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(5, max_lr=lrs, callbacks=[SaveModelCallback(learn, name='test_spm30k')])\n",
    "### SMALL: sz:256, bs:100, sl:100\n",
    "\n",
    "# char itos, N:6, 58M params\n",
    "    # 3.376897\t3.731502\t0.006678\t06:19   5cycle(1e-3)   'sm_char'\n",
    "    # greedy:    0.39514   .04035\n",
    "\n",
    "# 1k, N:6, 59M params\n",
    "    # 4.778663\t5.391022\t0.098169\t04:58   5cycle(1e-3)   'sm_spm'\n",
    "\n",
    "    # 39.337376\t36.469913\t0.707203\t04:24   freeze(1)  1cycle(1e-3)\n",
    "    # 8.762869\t8.898160\t0.184792\t04:59   unfreeze() 4cycle(1e-4,1e-3)  'sm_spm2'\n",
    "\n",
    "    # 24.711512\t22.952353\t0.486185\t04:27   freeze(1) 2cycle(1e-3)\n",
    "    # 6.491882\t6.510806\t0.121413\t04:59   unfreeze() 3cycle(1e-3)  'sm_spm3'\n",
    "\n",
    "    # 4.708159\t5.320672\t0.097103\t04:34   EncoderN:2, DecoderN:8, 55M params, 5cycle(1e-3)  'sm_spm4'\n",
    "\n",
    "# 5k, N:6, 63M params\n",
    "    # 9.057865\t9.884339\t0.249684\t04:55   5cycle(1e-3)  'sm_spm5'\n",
    "\n",
    "# 10k, N:6, 68M params\n",
    "    # 13.938698\t14.683375\t0.395669\t05:05   5cycle(1e-3)  'sm_spm_10k'\n",
    "\n",
    "### COMBO <6: sz:512, bs:32, sl:600\n",
    "\n",
    "# 5k, N:6, preload sm_spm5\n",
    "    # 17.112537\t16.310976\t0.063752\t20:37   'combo6_spm'\n",
    "    # greedy:    35.4337   .07501\n",
    "    # upload:    81.0504   .54935\n",
    "    #     pg:    79.9827   .25637\n",
    "\n",
    "# 10k, N:6, preload sm_spm_10k\n",
    "    # 76.443794\t70.586143\t0.384218\t21:58    'combo6_spm_10k'\n",
    "    # 74.961151\t70.133759\t0.382121\t21:45    'combo6_spm_10k_2'\n",
    "    \n",
    "# 30k, N:6, bs:20,  *preload char: combo_char\n",
    "    # 52.167412\t47.530663\t0.240256\t20:38   1cycle(1e-3) freeze(), 4 layer groups   'combo6_spm30k'\n",
    "    # 32.758427\t30.415497\t0.148593\t25:24   2nd run, 1cycle(1e-5,1e-3), unfreeze(), 'combo6_spm30k_2'\n",
    "    # greedy:    27.8843   .55172\n",
    "    \n",
    "### COMBO FULL: sz:512, bs:20, sl:750\n",
    "\n",
    "# char itos, N:6, 58M params\n",
    "    # 27.326603\t33.271706\t0.012635\t1:48:05   5cycle(1e-3), preload sm_char    'combo_char'\n",
    "    \n",
    "# 30k, N:6, 89M params, bs: 16, sl:750\n",
    "    # 64.151222\t51.847694\t0.083549\t1:19:01   2cycle(1e-6,1e-4), preload combo6_spm30k_2   'combo_spm30k'\n",
    "    # greedy:    56.6784   .08911\n",
    "    # 39.950447\t34.363934\t0.055618\t1:13:36   2nd run, 3cycle(1e-6,1e-4), \"\",    'combo_spm30k_2'\n",
    "    \n",
    "### COMBO TEST: sz:512, bs:16, sl:750\n",
    "    # 13.693297\t10.992985\t0.223984\t30:33    5cycle(1e-6,1e-4), preload above    'test_spm30k'\n",
    "    # greedy:    12.1767   .17320\n",
    "    #     pg:    70.9137   .07789\n",
    "    # upload:    49.3236   .40198"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Greedy Decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vdl = iter(learn.data.valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(vdl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Greedy Average"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "textify_fn = data.y.textify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def full_test(learn, sl, dl=data.valid_dl, batches=5):\n",
    "    learn.model.eval()\n",
    "    iterable = iter(dl)\n",
    "    g_loss,g_cer=0,0\n",
    "    if batches is None:\n",
    "        batches = len(dl.dl.dataset)//bs\n",
    "    for i in progress_bar(range(batches)):\n",
    "        x,y = next(iterable)\n",
    "        g_preds = learn.model(x, seq_len=sl)\n",
    "        g_res = torch.argmax(g_preds, dim=-1)\n",
    "        g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, textify_fn)[0]/bs]\n",
    "        g_loss+=g[0]\n",
    "        g_cer+=g[1]  \n",
    "    print(f'greedy:    {str(g_loss/batches)[:7]}   {str(g_cer/batches)[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "full_test(learn, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Single"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, textify_fn)[0]/bs]\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": [
    "#greedy\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    #i += 10\n",
    "    p = data.y.reconstruct(g_res[i])\n",
    "    ax= show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adjust State Dict and Split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/hw_combo_bt.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load(PATH/'models/distilbert_mlm_xtra.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "model_sd = OrderedDict()\n",
    "\n",
    "for k,v in sd['model'].items():\n",
    "    if k.startswith('transformer.encoder'):\n",
    "        c_k = k.replace('encoder', 'c_encoder')\n",
    "        w_k = k.replace('encoder', 'w_encoder')\n",
    "        model_sd[c_k] = v\n",
    "        model_sd[w_k] = v\n",
    "    else:\n",
    "        model_sd[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model']['transformer.generator.weight'] = sd['model']['transformer.c_generator.weight']\n",
    "sd['model']['transformer.generator.bias']  = sd['model']['transformer.c_generator.bias']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## load and split learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split(lambda m: [m.transformer]); None\n",
    "len(learn.layer_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# psutil.virtual_memory()  #41.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_to(1)\n",
    "learn.model.transformer.generator.weight.requires_grad #.conv[0].weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.unfreeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Experimentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_full_bert_tok_word'); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# lrs = slice(1e-5,1e-4)\n",
    "lrs = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(5, max_lr=lrs, callbacks=[SaveModelCallback(learn, name='test_combo_bert_tok_word')])\n",
    "# CHARS / WORDS\n",
    "# d_model:512, em_sz:256, resnet34, 5cycle(1e-3), N:4, drops:0.1, multi:8\n",
    "\n",
    "# SMALL DATASET: sz:256, sl:(100,50)\n",
    "\n",
    "# 4.805066\t4.812465\t0.051946\t07:56   bert_tok char only, p-mask(25)   'sm_char_pmask2'\n",
    "# chars:     0.56753   .07858\n",
    "\n",
    "# fit (not one_cycle); bert_tok char only\n",
    "# 36.014137\t29.500330\t0.354282\t07:09   freeze_to(-1); only tfmr 1e-3\n",
    "# 6.228029\t5.532068\t0.062850\t07:43   2nd - unfreeze; 4(1e-4,1e-3)\n",
    "# 3.253963\t3.492219\t0.039213\t07:44   3rd - 2(1e-5,1e-4)    'sm_char_pmask'\n",
    "\n",
    "# 2.858166\t3.149298\t0.034373\t07:34   char only (itos), p-mask(25), bs:100   'sm_char_itos'\n",
    "# greedy:    0.22749   .06057\n",
    "\n",
    "# bs:80, char_itos/word_itos(60k)\n",
    "# 4.847745\t5.941104\t0.027651\t0.029362\t08:52   'sm_char_word'\n",
    "# chars:    0.346770    0.03947     3.0s\n",
    "# words:    0.903070    0.03809\n",
    "# 4.197776\t5.396721\t0.024970\t0.022618\t14:57   2d pos, parallel_mask(25,10)  'sm_char_word2'\n",
    "# chars:    0.336420    0.03665     2.8s\n",
    "# words:    0.019130    0.02098\n",
    "\n",
    "# bs:60, 2d pos, parallel_mask(25,10)\n",
    "# 6.224502\t7.240156\t0.025720\t0.042237\t12:05   combo: itos  'sm_combo'\n",
    "# chars:    2.10703   .05335\n",
    "# words:    0.66827   .02541\n",
    "# 6.176572\t7.659454\t0.038917\t0.039825\t11:38   combo: bert_tok  'sm_combo_bert_tok'\n",
    "# chars:    0.89475   .04761\n",
    "# words:    0.98421   .02428\n",
    "# 6.128839\t7.527630\t0.042009\t0.048798\t14:26   combo: bert_tok; w/ STN before ResnetBase\n",
    "# chars:    0.64202   .06120\n",
    "# words:    0.75339   .02804\n",
    "\n",
    "# bert_tok(words), bs: 80\n",
    "# 12.985831\t13.813116\t0.381876\t07:27    2d pos, parallel_mask(10)  'sm_word_bert'   \n",
    "# words:    1.171370    0.54847\n",
    "\n",
    "# bert_tok(combo), bs: 40, 5cycle(1e-4)\n",
    "# 42.240177\t39.737682\t0.263751\t0.505153\t12:37   FPN, s_mask   'sm_bert_itos_fpn'\n",
    "\n",
    "# bert_tok(combo), bs: 60, 1cycle(1e-4,1e-3)->3fit(1e-4)\n",
    "# 24.818815\t23.533222\t0.137426\t0.344319\t10:53   simplified FPN (nn.tfmr)   'sm_bert_itos_fpn3'\n",
    "\n",
    "# bert_tok(combo), bs: 60, 5cycle(1e-3)\n",
    "# 13.043773\t13.111170\t0.072459\t0.136015\t10:40   simplified FPN (manual tfmr)  'sm_bert_itos_fpn4'\n",
    "\n",
    "# bert_tok(combo), bs: 60, 3cycle(1e-4,1e-3,1e-3)\n",
    "# 65.371353\t60.531727\t0.505481\t0.715806\t12:21   \"\", learned pos_enc  'sm_bert_itos_fpn5'\n",
    "# 14.045170\t14.263741\t0.098271\t0.147632\t12:11   \"\", 3 more cycles   'sm_bert_itos_fpn6'\n",
    "\n",
    "# COMBO_145<=6: sz:512, bs:15, seq_len:600/200\n",
    "\n",
    "# bert_tok(combo)\n",
    "# 26.950352\t27.504612\t0.025680\t0.074958\t51:41   FPN(simple), learned pos_enc  'combo<6_bert_itos_fpn'\n",
    "# chars:    29.5073     0.03396\n",
    "# words:    24.1089     0.04350\n",
    "\n",
    "# 40.977287\t36.501549\t0.039365\t0.133388\t34:21   \"\", train from scratch no preloading  'combo<6_char_word'\n",
    "# 25.210848\t24.398312\t0.024259\t0.088754\t38:35   \"\", preload 'sm_char_word2'    'combo<6_char_word2'\n",
    "# chars:    18.64750    0.01549\n",
    "# words:    7.597820    0.07416\n",
    "\n",
    "# 50.447319\t46.204708\t0.045057\t0.114878\t40:11   preload sm_combo_bert_tok, 5(1e-4),  'combo<6_bert_tok'\n",
    "# chars:    32.05150    0.03646\n",
    "# words:    30.36220    0.11244\n",
    "# --no improvement--   2nd run; 1(1e-5); Freeze img_enc, BnFreeze, AccumulateScheduler(3)\n",
    "# --no improvement--   2nd run; slice(3e-6,2e-5)\n",
    "\n",
    "# COMBO_145>=6\n",
    "\n",
    "# sz:512, bs:5, seq_len:750/250\n",
    "#  ????     178.2133    0.03660     0.113363            preload combo<6_bert_tok, 3(5e-5)    'combo>6_bert_tok'\n",
    "# chars:    716.0820    0.06242\n",
    "# words:    324.4240    0.11281\n",
    "#  ????     171.1198    0.03438     0.108465           'combo>6_bert_tok2'\n",
    "# chars:    494.2160    0.05982\n",
    "# words:    330.0630    0.20758\n",
    "\n",
    "# word decoder only; bs:15; lrs: slice(6e-6,2e-4)\n",
    "# 78.159187\t66.368530\t0.063820\t59:51   preload combo>6_bert_tok\n",
    "# words:    115.2590    0.09386\n",
    "# 63.500256\t54.264275\t0.052015\t59:35   2nd run   'combo>6_bert_tok_word'\n",
    "# words:    120.4400    0.08898\n",
    "\n",
    "\n",
    "# COMBO FULL\n",
    "\n",
    "# bert_tok(combo), bs: 8\n",
    "# 114.737350\t86.396645\t0.053524\t0.126384\t3:41:44   FPN, preload...   'combo_full_bert_itos_fpn'\n",
    "# chars:    84.9115     0.07217\n",
    "# words:    51.8369     0.07182\n",
    "\n",
    "# bert_tok(word), bs: 10\n",
    "# 56.488300\t 45.808746\t0.057084\t1:40:46  preload combo_full_bert_itos_fpn -> word   'word_full_bert_itos_fpn'\n",
    "# greedy:    77.2583    0.05492\n",
    "\n",
    "# char itos only, sz:512, bs:20, sl: 700, 3cycle(1e-4,1e-3)\n",
    "# 17.519344\t12.332630\t0.012053\t1:33:27     preload combo<6_char_word2;  'full_combo_char_itos'\n",
    "#  chars:    14.1142   .03507\n",
    "#     pg:    109.775   .06105\n",
    "# upload:    81.1906   .43663\n",
    "\n",
    "# word decoder only, sz:512, bs:15, word_len:300, lrs: slice(1e-6,2e-5)\n",
    "# 48.172344\t39.704689\t0.058319\t1:31:44  preload combo>6_bert_tok_word    'combo_full_bert_tok_word'\n",
    "# words:    49.00830    0.05337\n",
    "# 50.854053\t39.678852\t0.058055\t1:24:53  2nd run    'combo_full_bert_tok_word2'\n",
    "# words:    50.46420    0.05224\n",
    "#    pg:    73.38090    0.06796\n",
    "# 51.945342 40.016457   0.058763    1:27:13   add 2 additional layers; 1e-6  --failure!!\n",
    "\n",
    "# word/char, sz:512, bs:8, 1cycle(2e-5)\n",
    "# 134.963303\t110.691048\t0.047529\t0.128654\t3:29:36   preload combo>6_bert_tok2   'combo_full_bert_tok'\n",
    "# chars:    96.1589   .07463\n",
    "# words:    41.3268   .08943\n",
    "\n",
    "# COMBO TEST\n",
    "\n",
    "# word decoder only, sz:512, bs:20, sl:300, 5cycle(1e-5,1e-4)\n",
    "# 6.876280\t5.052646\t0.029255\t34:44      'test_combo_bert_tok_word'\n",
    "# greedy:    3.25465   .01724\n",
    "#     pg:    81.2414   .06737\n",
    "# upload:    52.4035   .29801\n",
    "\n",
    "# char itos only, sz:512, bs:20, sl:700, 5cycle(1e-5,1e-4)\n",
    "# 11.622255\t8.140586\t0.049724\t38:43      'test_combo_char_itos'\n",
    "#  chars:    10.6597   .10455\n",
    "#     pg:    118.629   .05764\n",
    "# upload:    83.0442   .28538\n",
    "\n",
    "# word/char, sz:512, bs:10, 1cycle(1e-4)\n",
    "# 52.687794\t39.834663\t0.137030\t0.224802\t1:05:14   preload combo_full_bert_tok    'test_combo_bert_tok'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# bs:60, 2d pos, parallel_mask(25,10)\n",
    "# 6.176572\t7.659454\t0.038917\t0.039825\t11:38   combo: bert_tok  'sm_combo_bert_tok'\n",
    "# chars:    0.89475   .04761\n",
    "# words:    0.98421   .02428\n",
    "\n",
    "# COMBO_145<=6: sz:512, bs:15, seq_len:600/200\n",
    "# 50.447319\t46.204708\t0.045057\t0.114878\t40:11   preload sm_combo_bert_tok, 5(1e-4),  'combo<6_bert_tok'\n",
    "# chars:    32.05150    0.03646\n",
    "# words:    30.36220    0.11244\n",
    "\n",
    "# COMBO_145>=6, sz:512, bs:5, seq_len:750/250\n",
    "#  ????     171.1198    0.03438     0.108465           preload combo<6_bert_tok, 3(5e-5)    'combo>6_bert_tok2'\n",
    "# chars:    494.2160    0.05982\n",
    "# words:    330.0630    0.20758\n",
    "\n",
    "# COMBO FULL, sz:512, bs:8, 1cycle(2e-5)\n",
    "# 134.963303\t110.691048\t0.047529\t0.128654\t3:29:36   preload combo>6_bert_tok2   'combo_full_bert_tok'\n",
    "# chars:    96.1589   .07463\n",
    "# words:    41.3268   .08943\n",
    "\n",
    "# COMBO TEST, sz:512, bs:20, sl:300, 5cycle(1e-5,1e-4)\n",
    "# 6.876280\t5.052646\t0.029255\t34:44      word decoder only   'test_combo_bert_tok_word'\n",
    "# greedy:    3.25465   .01724\n",
    "#     pg:    81.2414   .06737\n",
    "# upload:    52.4035   .29801"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Save Outputs for LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SaveOutput(LearnerCallback):\n",
    "    def __init__(self, learn:Learner, filepath):\n",
    "        super().__init__(learn)\n",
    "        self.filepath = filepath\n",
    "        self.recon_fn = learn.data.y.reconstruct\n",
    "\n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.res = []\n",
    "\n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        preds = torch.argmax(last_output, dim=-1)\n",
    "        for p,t in zip(preds,last_target):\n",
    "            p_str = str(self.recon_fn(p))\n",
    "            t_str = str(self.recon_fn(t))\n",
    "            self.res.append([t_str,p_str])\n",
    "\n",
    "    def on_epoch_end(self, **kwargs):\n",
    "        df = pd.DataFrame(self.res, columns=['targets', 'preds'])\n",
    "        df.to_csv(self.filepath, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.validate(dl=data.train_dl,\n",
    "               callbacks=[TeacherForce(learn), SaveOutput(learn, filepath=PATH/'output.csv')],\n",
    "               metrics=[CER(data.y.reconstruct)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv(PATH/'output.csv'); df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Char/Word Greedy results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def greedy_decode(src, model, seq_len, kind='char', bos_tok=1):\n",
    "    model.eval()\n",
    "    tfmr = model.transformer\n",
    "    img_enc = model.img_enc\n",
    "    adaptor = model.adaptor\n",
    "    \n",
    "    decoder = tfmr.c_decoder if kind=='char' else tfmr.w_decoder\n",
    "    embed = tfmr.c_embed if kind=='char' else tfmr.w_embed\n",
    "    p_num = 20 #25 if kind=='char' else 20\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        feats = tfmr.encode(adaptor(img_enc(src)))\n",
    "        #feats = tfmr.encoder(img_enc(src))\n",
    "        bs = src.size(0)\n",
    "        tgt = torch.zeros((bs,1), dtype=torch.long, device=device) + bos_tok\n",
    "\n",
    "        res = []\n",
    "        for i in progress_bar(range(seq_len)):\n",
    "#             mask = subsequent_mask(tgt.size(-1))\n",
    "            mask = parallelogram_mask(tgt.size(-1), p_num)\n",
    "            \n",
    "            dec_outs = decoder(embed(tgt), feats, mask)\n",
    "            prob = tfmr.generator(dec_outs[:,-1])\n",
    "            res.append(prob)\n",
    "            pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "            if (pred==0).all(): break\n",
    "            tgt = torch.cat([tgt,pred], dim=-1)\n",
    "        out = torch.stack(res).transpose(1,0).contiguous()\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vdl = iter(learn.data.valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(vdl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Single Word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = greedy_decode(x, learn.model, word_len, 'word', 101)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, data.y.reconstruct)[0]/bs]\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": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 50))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    #i += 8\n",
    "    p = data.y.reconstruct(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Single Char"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = greedy_decode(x, learn.model, seq_len, 'char')\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, data.y.reconstruct)[0]/bs]\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": [
    "#greedy\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = data.y.reconstruct(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-11-01T17:02:56.281776Z",
     "start_time": "2019-11-01T17:02:53.656973Z"
    },
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Combo Chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = greedy_decode(x, learn.model, seq_len, 'char', 101)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "loss_func = LabelSmoothing()\n",
    "g = [loss_func(g_preds, y[0]).item()/bs, cer(g_preds, y[0], data.y.reconstruct_one)[0]/bs]\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,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 50))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = data.y.reconstruct_one(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-11-01T17:02:56.281776Z",
     "start_time": "2019-11-01T17:02:53.656973Z"
    },
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Combo Words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = greedy_decode(x, learn.model, word_len, 'word', 101)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "loss_func = LabelSmoothing()\n",
    "g = [loss_func(g_preds, y[1]).item()/bs, cer(g_preds, y[1], data.y.reconstruct_one)[0]/bs]\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,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 50))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = data.y.reconstruct_one(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Full Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return self.linear(x) * 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "            out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    #transformer.generator.weight = transformer.tgt_embed[0].lut.weight\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# TransformerXL"
   ]
  },
  {
   "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)  #(outer) matrix product of 2 vectors\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",
    "        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 _left_shift(x:Tensor):\n",
    "    \"Shift the line i of `x` by p-i elements to the left.\"\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 DecoderLayerXL(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.emb = nn.Embedding(vocab_sz, d_model)\n",
    "        self.drop_emb = nn.Dropout(drops)\n",
    "        self.pos_enc = PositionalEncoding(d_model)\n",
    "        self.u = nn.Parameter(torch.Tensor(n_heads, 1, d_head))\n",
    "        self.v = nn.Parameter(torch.Tensor(n_heads, 1, d_head))\n",
    "        self.layers = nn.ModuleList([DecoderLayerXL(n_heads, d_model, drops) for k in range(n_layers)])\n",
    "        self.init = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        if not self.init:\n",
    "            self.reset()\n",
    "            self.init = True\n",
    "        bs,x_len = x.size()\n",
    "        inp = self.drop_emb(self.emb(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",
    "        mask = torch.triu(x.new_ones(x_len, seq_len), diagonal=1+m_len).bool()[None,None]\n",
    "        pos = torch.arange(seq_len-1, -1, -1, device=inp.device, dtype=inp.dtype)\n",
    "        pos_enc = self.pos_enc(pos)\n",
    "        hids = []\n",
    "        hids.append(inp)\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            mem = self.hidden[i]\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",
    "        self._update_mems(hids)\n",
    "        return self.hidden,[core_out]\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 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 select_hidden(self, idxs):\n",
    "        # used in beam search only...\n",
    "        self.hidden = [h[idxs] for h in self.hidden]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, lm, layer, N):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.lm_layer = lm\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, src):\n",
    "        x = self.lm_layer(x)\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    \"Decoder: self-attn, src-attn, and feed forward\"\n",
    "    def __init__(self, size, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src):\n",
    "        x = self.sublayer[0](x, lambda x: self.src_attn(x, src, src))\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 make_full_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    lm = TransformerXL(vocab, d_model, n_layers=10, n_heads=8, drops=0.1)\n",
    "    attn = MultiHeadedAttention(d_model, attn_heads) \n",
    "    ff   = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(lm, DecoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    dec_outs = self.transformer.decode(feats, tgt)\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt)\n",
    "            out = self.transformer.generate(dec_outs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ Transformer LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayerWithLM(nn.Module):\n",
    "    \"Decoder: self-attn, src-attn, and feed forward\"\n",
    "    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = clones(feed_forward, 3)\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 5)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        # lm\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))  # shared btw lm and decoder\n",
    "        lm = self.sublayer[1](x, self.feed_forward[0])\n",
    "        \n",
    "        # decoder layer\n",
    "        dec = self.sublayer[2](x, lambda x: self.src_attn(x, src, src))\n",
    "        dec = self.sublayer[3](dec, self.feed_forward[1])\n",
    "        return self.sublayer[4](dec+lm, self.feed_forward[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', attn_heads=8, weight_tying=False):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayerWithLM(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    if weight_tying:\n",
    "        model.generator.weight = model.tgt_embed[0].lut.weight\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ AWDLSTM LM integrated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, pos_enc):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.pos_enc = pos_enc\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.pos_enc(tgt), src, tgt_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8   #(cube root of d_model?)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def lm_mod_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        PositionalEncoding(d_model, drops, 2000),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LM(nn.Module):\n",
    "    def __init__(self, vocab, d_model, n_hidden, n_layers):\n",
    "        super(LM, self).__init__()\n",
    "        self.lm = AWD_LSTM(vocab, d_model, n_hidden, n_layers, pad_token=0,\n",
    "                           hidden_p=0.1, input_p=0.25, embed_p=0.05, weight_p=0.2)\n",
    "                \n",
    "    def forward(self, tgts):\n",
    "        res = self.lm(tgts, from_embeddings=True)\n",
    "        pdb.set_trace()\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Mixer(nn.Module):\n",
    "    def __init__(self, d_model, vocab, drops=0.2):\n",
    "        super(Mixer, self).__init__()\n",
    "        self.mixer = nn.Sequential(PositionwiseFeedForward(d_model, drops), LayerNorm(d_model))\n",
    "        self.generator = nn.Linear(d_model, vocab)\n",
    "        \n",
    "    def forward(self, dec_outs, lm_outs):\n",
    "        mix = self.mixer((lm_outs+outs)/2)\n",
    "        return self.generator(mix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    def __init__(self, d_model, vocab, drops, pad_tok=0):\n",
    "        super(Embedding, self).__init__()\n",
    "        self.emb = nn.Embedding(vocab, d_model, padding_idx=pad_tok)\n",
    "        self.emb_drop = EmbeddingDropout(self.emb, drops)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.emb_drop(x)  #self.lut(x) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, embedding, lm, mixer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.embedding = embedding\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        self.mixer = mixer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    tgt_emb = self.embedding(tgt)\n",
    "                    dec_outs = self.transformer.decode(feats, tgt_emb, mask)\n",
    "                    lm_outs = self.lm(tgt_emb)\n",
    "                    prob = self.mixer(dec_outs[:,-1], lm_outs[:,-1])\n",
    "                    #outs = self.lm(dec_outs[:,-1], self.transformer.embed, decode=True)\n",
    "                    #prob = self.transformer.generate(outs)\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            tgt_emb = self.embedding(tgt)\n",
    "            dec_outs = self.transformer(feats, tgt_emb, tgt_mask)    # ([bs, sl, d_model])\n",
    "            lm_outs = self.lm(tgt_emb)\n",
    "            out = self.mixer(dec_outs, lm_outs)\n",
    "            #out = self.transformer.generate(outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    vocab = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = lm_mod_model(vocab, d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    embedding = Embedding(d_model, vocab, drops, pad_tok=0)\n",
    "    lm = LM(vocab, d_model, 1400, 3)\n",
    "    mixer = Mixer(d_model, vocab, drops)\n",
    "    net = Img2Seq(img_encoder, transformer, embedding, lm, mixer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# update existing model according to LM modifications\n",
    "sd['model']['embedding.emb.weight'] = sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "sd['model']['embedding.emb_drop.emb.weight'] = sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "sd['model'][\"transformer.pos_enc.pe\"] = sd['model']['transformer.tgt_embed.1.pe']\n",
    "sd['model']['mixer.generator.weight'] = sd['model']['transformer.generator.weight']\n",
    "sd['model']['mixer.generator.bias'] = sd['model']['transformer.generator.bias']\n",
    "\n",
    "del sd['model']['transformer.tgt_embed.1.pe']\n",
    "del sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "del sd['model']['transformer.generator.weight']\n",
    "del sd['model']['transformer.generator.bias']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load('data/wikitext/models/wiki103_lm_enc.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_lm_sd = OrderedDict()\n",
    "for k, v in lm_sd.items():\n",
    "    name = 'lm.lm.'+k\n",
    "    new_lm_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].update(new_lm_sd)\n",
    "\n",
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tie weights of transformer embedding and generator w/ lm encodings\n",
    "learn.model.embedding.emb.weight = learn.model.lm.lm.encoder.weight\n",
    "learn.model.embedding.emb_drop.emb.weight = learn.model.lm.lm.encoder_dp.emb.weight\n",
    "learn.model.mixer.generator.weight = learn.model.lm.lm.encoder.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('combo_512_9_wiki103_base_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## load and split learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_9_wiki103_base_lm')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split([learn.model.img_enc, learn.model.embedding, learn.model.transformer, learn.model.lm, learn.model.mixer])\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.layer_groups[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_to(-3)\n",
    "learn.model.img_enc.linear.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrs = slice(2e-5, 2e-4, 2e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ AWDLSTM added"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', attn_heads=8, weight_tying=False):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    if weight_tying:\n",
    "        model.generator.weight = model.tgt_embed[0].lut.weight\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, lm):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    dec_prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    lm_prob,_,_ = self.lm(tgt)\n",
    "                    prob = (dec_prob + lm_prob[:,-1])/2\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "            out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    vocab = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(vocab, d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    lm = get_language_model(AWD_LSTM, vocab, lm_config, drop_mult=0.5)\n",
    "    net = Img2Seq(img_encoder, transformer, lm)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.lm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load(PATH/'models/wiki2_lm.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_lm_sd = OrderedDict()\n",
    "for k, v in lm_sd['model'].items():\n",
    "    name = 'lm.'+k\n",
    "    new_lm_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd['model'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].update(new_lm_sd)\n",
    "\n",
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('combo_512_9_wiki2_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Num Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "                .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "                .label_from_df(cols=2)\n",
    "                .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "                .databunch(bs=bs, val_bs=bs*2, device=device)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2 = cnn_learner(data, models.resnet34, metrics=accuracy)\n",
    "# Note: already split and frozen:  m[0][6], m[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# load saved sd into l2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_sd = OrderedDict()\n",
    "\n",
    "sd = torch.load(PATH/'models/combo_512_mix_res9.pth', map_location=device)\n",
    "for k, v in sd['model'].items():\n",
    "    if k.startswith('img_enc.base'):\n",
    "        name = k.replace('img_enc.base', '0')\n",
    "        new_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2_sd = torch.load(PATH/'models/line_97acc.pth', map_location=device)\n",
    "for k, v in l2_sd['model'].items():\n",
    "    if k.startswith('1.'):\n",
    "        new_sd[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.model.load_state_dict(new_sd, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.layer_groups[1][47].weight.requires_grad, l2.layer_groups[2][9].weight.requires_grad"
   ]
  },
  {
   "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,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "l2.lr_find()\n",
    "l2.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.fit_one_cycle(1,1e-3)\n",
    "# 0.171519\t0.069522\t0.976006\t12:57     'line_97acc'\n",
    "# 0.094188\t0.050964\t0.988564\t12:01  preload img_enc: combo_512_mix_res9 and head: line_97acc   'line_98acc'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.save('line_98acc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn = cnn_learner(data, models.resnet34, custom_head=create_head(1024, 1), metrics=accuracy, loss_func=MSELossFlat())\n",
    "# learn.fit_one_cycle(1,1e-3)\n",
    "# # 1.229723\t0.113453\t0.9057\t12:43\n",
    "\n",
    "# preds,y,losses = learn.get_preds(with_loss=True)\n",
    "\n",
    "# # actual accuracy\n",
    "# preds.apply_(round)\n",
    "# (preds.long().squeeze(-1) == y).float().mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=4, drops=0, attn_type='multi', heads=8)\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_mix_res9')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.data = data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(5, max_lr=1e-5, callbacks=[SaveModelCallback(learn, name='sm_256_lines2')])\n",
    "#sm, 5cycle, 1e-3\n",
    "\n",
    "# Image Preprocessing by num_lines\n",
    "# 11.384931\t10.208405\t0.122443   'sm_256_lines'\n",
    "# 10.548290\t9.845284\t0.116961   2nd run; 5cycle(1e-5); 'sm_256_lines2'\n",
    "\n",
    "# \n",
    "\n",
    "# 2.474220\t4.081600\t0.036482  N:4,F.gelu,sz:256,em_sz:512,single,drop:0  'sm_256_1'\n",
    "# greedy:    1.36697   .066\n",
    "# 4.430638\t4.402001\t0.034730  2nd run, lr:1.5e-5, add tfms, drop:0.2   'sm_256_2'\n",
    "# greedy:    0.37206   .04145\n",
    "\n",
    "# 5.158808\t4.964441\t0.042544  \"\", w/ tfms, drop:0.2   'sm_256_3'\n",
    "# greedy:    0.50646   .04103\n",
    "# 4.266788\t4.694098\t0.039822  2nd run, lr:1.5e-5   'sm_256_4'\n",
    "# greedy:    0.38972   .03879\n",
    "\n",
    "# 3.678168\t3.962710\t0.035466  N:8,tfms   'sm_256_5'\n",
    "# greedy:    0.38097   .04405\n",
    "\n",
    "# 2.840161\t3.429177\t0.030243  N:4,tfms,multi(8)  'sm_256_6'\n",
    "# greedy:    0.32775   .04201\n",
    "# greedy:    1.38865   .06943\n",
    "\n",
    "# 3.377438\t3.786088\t0.033931   N:6,tfms,drop:0.1,multi(8)   'sm_256_7'\n",
    "# greedy:    0.36201   .04266  \n",
    "# greedy:    1.40557   .06902\n",
    "\n",
    "# 3.989137\t4.353169\t0.037717   N:4,weight-tying,multi(16)   'sm_256_8'\n",
    "# greedy:    1.44291   .07299\n",
    "\n",
    "# 4.806501\t5.007887\t0.042339   N:4,weight-tying,multi(8),drops:0.2   'sm_256_9'\n",
    "# greedy:    1.46753   .07424\n",
    "\n",
    "# 12.817862\t11.842313\t0.105486   N:4,no init/weight tying,multi(8),drops:0.1   'sm_256_10'\n",
    "\n",
    "# 3.515456\t3.880955\t0.033763   N:4,multi(8),drops:0.1   'sm_256_11'\n",
    "# greedy:    1.40945   .06964\n",
    "\n",
    "# combo_cat6 - preload 'sm_256_7'\n",
    "# 7.570516\t6.670641\t0.015070    'combo_512_7'\n",
    "# greedy:    6.29406   .02186\n",
    "\n",
    "\n",
    "# combo 6lines and under, preload 'sm_256_6', 10cycle(1e-3)\n",
    "# 10.784849\t9.144131\t0.024880     'combo_512_13'\n",
    "# greedy:    9.21715   .02788\n",
    "# upload:    103.456   .44797\n",
    "#     pg:    148.473   .69869\n",
    "\n",
    "# combo 6lines and over, preload 'combo_512_13', 5cycle(2e-5)\n",
    "# 56.732105\t38.604137\t0.026153    'combo_512_14'\n",
    "# greedy:    141.022   .02933\n",
    "# upload:    179.606   .81724\n",
    "\n",
    "# combo_cat_pg - preload 'combo_512_7'\n",
    "# 25.674347\t21.161726\t0.015869    'combo_512_8'\n",
    "# greedy:    114.878   .03654\n",
    "#   test:    115.247   .05014\n",
    "\n",
    "# combo_cat_pg_dl_sorted - preload 'combo_512_8', 5cycle, 1.5e-4\n",
    "# 9.800321\t5.676855\t0.006939    'combo_512_9'\n",
    "# greedy:    34.8757   .01523\n",
    "\n",
    "#   test:    91.3226   .05102\n",
    "#   test:    104.685   .05483\n",
    "#   test:    116.467   .05315\n",
    "\n",
    "# upload:    115.149   .66801\n",
    "\n",
    "# imdb_wiki_combo - preload 'combo_512_9'; split: img_enc, transformer; 3cycle, slice(5e-5, 1e-3)\n",
    "# 20.390352\t16.698141\t0.018726    'combo_512_12'\n",
    "# greedy:    60.8384   .02474\n",
    "#     pg:    112.755   .05204\n",
    "\n",
    "# combo_145k - fit:1cycle(1e-3); preload combo_512_9; drops:0.2\n",
    "# 30.075714\t31.024599\t0.031703\t1:50:33\n",
    "\n",
    "# combo_145k <= 6 - 5cycle(1e-4); preload sm_256_mix_res\n",
    "# 19.502026\t17.365313\t0.044436\n",
    "# greedy:    14.6870   .04744\n",
    "# upload:    75.1357   .47955\n",
    "#     pg:    142.710   .66233\n",
    "\n",
    "# test_combo 65k (sm,cat,pg,dl)\n",
    "# 18.262995\t15.091379\t0.039351   5cycle(1e-5) \n",
    "# greedy:    129.968   .03932\n",
    "# upload:    81.6835   .21794\n",
    "#     pg:    127.317   .05613\n",
    "\n",
    "# 16.678631\t12.867742\t0.035985   1cycle(1e-5), drops:0.1   \n",
    "# 14.878900\t12.797726\t0.035834   1cycle(1e-6), drops:0.1  \n",
    "# greedy:    34.3988   .05552\n",
    "# upload:    75.5511   .21010\n",
    "#     pg:    118.963   .05287\n",
    "\n",
    "# 16.086544\t12.320695\t0.034263   1cycle(1e-5), drops:0.1   \n",
    "# 11.640953\t11.339241\t0.032201   1cycle(1e-5), drops:0     'combo_512_mix_res9' ***\n",
    "# greedy:    31.5506   .05410\n",
    "#   test:    74.4820   .22067\n",
    "#     pg:    105.182   .05076"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Greedy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "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": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'][\"transformer.src_adapt.0.weight\"] = sd['model']['img_enc.linear.weight']\n",
    "sd['model'][\"transformer.src_adapt.0.bias\"] = sd['model']['img_enc.linear.bias']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "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": [
    "# learn.load('combo_512_9_wiki2_lm')\n",
    "learn.load('combo_512_mix_res')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def full_test(learn, sl, dl=data.valid_dl, batches=10):\n",
    "    learn.model.eval()\n",
    "    iterable = iter(dl)\n",
    "    g_loss,g_cer=0,0\n",
    "    if batches is None:\n",
    "        batches = len(dl.dl.dataset)//bs\n",
    "    for i in progress_bar(range(batches)):\n",
    "        x,y = next(iterable)\n",
    "        g_preds = learn.model(x, seq_len=sl)\n",
    "        g_res = torch.argmax(g_preds, dim=-1)\n",
    "        g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, itos)[0]/bs]\n",
    "        g_loss+=g[0]\n",
    "        g_cer+=g[1]\n",
    "    return [g_loss/batches, g_cer/batches]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "g = full_test(learn, seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": [
    "# losses = np.array([learn.loss_func(g_preds[i:i+1],y[i:i+1]).item() for i in range(bs)])\n",
    "# cers = np.array([cer(g_preds[i:i+1],y[i:i+1])[0] for i in range(bs)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "x,y = next(iter(learn.data.valid_dl))\n",
    "\n",
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, word_itos)[0]/bs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(2,3, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i], word_itos, sep='')\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'uploads'\n",
    "df = pd.read_csv(PATH/'uploads.csv')\n",
    "len(df)\n",
    "\n",
    "sz,bs = 512,14\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'paragraphs'\n",
    "df = pd.read_csv(PATH/'test_pg.csv')\n",
    "len(df)\n",
    "\n",
    "sz,bs = 512,15\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "## Test dataset only!!!\n",
    "test_data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=SequenceList, vocab=CharVocab(itos), tokenizer=CharTokenizer)\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\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": [
    "# combo(bert_tok)\n",
    "test_data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=MultiSequenceList, vocab=BertVocab(), pad_idx=0)\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# words(bert_tok)\n",
    "test_data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=MultiSequenceList, vocab=BertVocab(), pad_idx=0)\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\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": [
    "# sentencepiece model\n",
    "test_data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=SPMList, sp=sp)\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\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": [
    "# learn.data = test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "dl = iter(test_data.train_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=4, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.load('combo_512_9_wiki2_lm')\n",
    "# learn.load('combo_512_mix_res')\n",
    "learn.load('test_combo_char_itos')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.data = data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Bert tok"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = greedy_decode(x, learn.model, word_len, 'word', 101)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, data.y.reconstruct)[0]/bs]\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": [
    "#greedy\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = data.y.reconstruct(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## GPU batch testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y)[0]/bs]\n",
    "\n",
    "print(f'  test:    {str(g[0])[:7]}   {str(g[1])[1:7]}')\n",
    "\n",
    "# test:    397.121   .29261  ~3:55   combo_512_9_wiki2_lm  (LM isn't aware of '\\n'!!)\n",
    "# test:    101.374   .05139  ~27s    combo_512_9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#test\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    i +=4\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Preprocessing Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# test data\n",
    "x,y = next(iter(data.train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def lines_into_paragraph(res):\n",
    "    out = []\n",
    "    for r in res:\n",
    "        s = np.split(r, np.where(r == 2)[0])[0].numpy()\n",
    "        out.append(np.append(s,4))\n",
    "    out = np.concatenate(out)\n",
    "    out[-1] = 2  # replace final '\\n' with 'eos'\n",
    "    return torch.from_numpy(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def target_into_lines(targ, split_idx=4):\n",
    "    nonzero = targ[targ.nonzero()].flatten()\n",
    "    lines = np.split(nonzero, np.where(nonzero == split_idx)[0])\n",
    "    maxlen = len(max(lines,key=len))\n",
    "    res = torch.zeros((len(lines),maxlen))\n",
    "    for i,arr in enumerate(lines):\n",
    "        res[i,:len(arr)] = arr\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from scipy import signal,ndimage\n",
    "import statistics\n",
    "\n",
    "def img_into_line_tensor(img):\n",
    "    arr = img[0].numpy()\n",
    "    w = arr.shape[-1]\n",
    "    heights = []\n",
    "    lengths = []\n",
    "    stds = arr.std(axis=1)\n",
    "    g_stds = scipy.ndimage.gaussian_filter1d(stds, 5)\n",
    "    peaks,_ = scipy.signal.find_peaks(g_stds, prominence=stds.std()//3, distance=20)\n",
    "    mins = scipy.signal.argrelextrema(g_stds, np.less_equal)[0]  # np.less_equal critical for flat minima, edges\n",
    "    for p in peaks:\n",
    "        rows = range(mins[mins < p][-1], mins[mins > p][0])\n",
    "        lengths.append(len(rows))\n",
    "        heights.append(arr[rows])\n",
    "    max_len = max(lengths)\n",
    "    outs = torch.ones((len(heights),3 , max_len, w))\n",
    "    for i,h in enumerate(heights):\n",
    "        outs[i,:,:len(h)] = torch.from_numpy(h)\n",
    "    return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "inps = img_into_line_tensor(x[idx])\n",
    "targs = target_into_lines(y[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(inps, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(f'test cer:   {str(cer(g_preds,targs)[0]/g_preds.size(0))[1:7]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true
   },
   "source": [
    "## Single Image (cpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = test_data.train_ds[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "%time pred = learn.predict(x)\n",
    "# CPU total time: 4min 24s\n",
    "\n",
    "# tx,_ = test_data.one_item(x)\n",
    "# %time pred = greedy_decode(tx, learn.model, seq_len, 'char')    p_mask(25)\n",
    "# CPU total time: 9min 42s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "show_img(x, title=str(pred[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "probs = F.softmax(pred[2], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "scores, idxs = torch.topk(probs, 3, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "scores[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for s,i in zip(scores,idxs):\n",
    "    new = {}\n",
    "    for score, idx in zip(s,i):\n",
    "        new[itos[idx]] = score.data\n",
    "    print(new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chars = idxs.numpy()\n",
    "chars.appl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idxs.apply_(lambda x: itos[x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for i in range(30):\n",
    "    print(char_label_text(idxs[i], sep=' '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.show_results(ds_type=data.train_ds, rows=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "i = 0\n",
    "x = xs[i][None]\n",
    "y = ys[i][None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "batch = data.one_item(xs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xs,ys = data.one_batch(detach=False, denorm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item(), cer(g_preds, y)[0]]\n",
    "\n",
    "print(f'  test:    {str(g[0])[:7]}   {str(g[1])[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "p = char_label_text(g_res[0])\n",
    "show_img(x[0], figsize=(18,10), title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "im = PATH/'uploads'/'test1.png'\n",
    "img = open_image(im)\n",
    "prediction = learn.predict(img)[0]\n",
    "show_img(img, title=str(prediction))"
   ]
  }
 ],
 "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"
  },
  "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
}