1--DistilBert_Arch.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": [
    "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": {},
   "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,15\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=2, max_warp=0.1, max_lighting=0.5)\n",
    "\n",
    "def force_gray(image): return image.convert('L').convert('RGB')"
   ]
  },
  {
   "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": [
    "from transformers import DistilBertTokenizer, DistilBertForMaskedLM"
   ]
  },
  {
   "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 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": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Individual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bs = 16"
   ]
  },
  {
   "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 BertTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[str]: return bert_tok.tokenize(t) + [\"[SEP]\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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])    "
   ]
  },
  {
   "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": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": [
    "### Combo"
   ]
  },
  {
   "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 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 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": [
    "## SentencePiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "\n",
    "# sp = spm.SentencePieceProcessor()\n",
    "# sp.Load(str(PATH/'spm_full_10k.model'))\n",
    "# sp.SetEncodeExtraOptions(\"eos\")\n",
    "# sp.SetDecodeExtraOptions(\"bos:eos\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": "markdown",
   "metadata": {},
   "source": [
    "### Word only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bs=20"
   ]
  },
  {
   "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": [
    "# # old\n",
    "# class SPMTokenizer(BaseTokenizer):\n",
    "#     def tokenizer(self, t:str) -> List[str]: return sp.EncodeAsIds(t)[1:]\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.pad_idx = 0\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": [
    "# #old\n",
    "# odata = (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
   },
   "outputs": [],
   "source": [
    "class SPTokenizer(BaseTokenizer):\n",
    "    \"Wrapper around a SentncePiece tokenizer to make it a `BaseTokenizer`.\"\n",
    "    def __init__(self, model_prefix:str):\n",
    "        self.tok = spm.SentencePieceProcessor()\n",
    "        self.tok.Load(f'{model_prefix}.model')\n",
    "        self.tok.SetEncodeExtraOptions(\"eos\")\n",
    "\n",
    "    def tokenizer(self, t:str) -> List[str]:\n",
    "        return self.tok.EncodeAsIds(t)[1:]\n",
    "      \n",
    "class CustomTokenizer():\n",
    "    def __init__(self, tok_func:Callable, model_prefix:str):\n",
    "        self.tok_func, self.model_prefix = tok_func,model_prefix\n",
    "        self.pre_rules = [rm_useless_spaces, add_cap_tokens]\n",
    "        \n",
    "    def __repr__(self) -> str:\n",
    "        res = f'Tokenizer {self.tok_func.__name__} using `{self.model_prefix}` model with the following rules:\\n'\n",
    "        for rule in self.pre_rules: res += f' - {rule.__name__}\\n'\n",
    "        return res        \n",
    "\n",
    "    def process_one(self, t:str, tok:BaseTokenizer) -> List[str]:\n",
    "        \"Processe one text `t` with tokenizer `tok`.\"\n",
    "        for rule in self.pre_rules: t = rule(t)  \n",
    "        toks = tok.tokenizer(t) \n",
    "        return toks \n",
    "                                                                         \n",
    "    def process_all(self, texts:Collection[str]) -> List[List[str]]: \n",
    "        \"Process a list of `texts`.\" \n",
    "        tok = self.tok_func(self.model_prefix)\n",
    "        return [self.process_one(t, tok) for t in texts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SPList(ItemList):\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        model_prefix = self.path/'spm_full_10k'\n",
    "        cust_tok = CustomTokenizer(SPTokenizer, model_prefix)\n",
    "        self.processor = TokenizeProcessor(tokenizer=cust_tok, include_bos=False)\n",
    "        self.sp = spm.SentencePieceProcessor()\n",
    "        self.sp.Load(str(model_prefix)+'.model')\n",
    "        self.sp.SetDecodeExtraOptions(\"bos:eos\")\n",
    "        \n",
    "        self.pad_idx = 0\n",
    "        self.copy_new += ['sp']\n",
    "    \n",
    "    def get(self, i):\n",
    "        o = self.items[i]\n",
    "        return Text(o, self.textify(o))\n",
    "    \n",
    "    def reconstruct(self, t:Tensor):\n",
    "        nonzero_idxs = (t != self.pad_idx).nonzero()\n",
    "        idx_min = 0\n",
    "        idx_max = nonzero_idxs.max() if len(nonzero_idxs) > 0 else 0\n",
    "        return Text(t[idx_min:idx_max+1], self.textify(t[idx_min:idx_max+1]))\n",
    "\n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)        \n",
    "    \n",
    "    def textify(self, ids):\n",
    "        if isinstance(ids, torch.Tensor): ids = ids.tolist()\n",
    "        st = self.sp.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=SPList)\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
   },
   "outputs": [],
   "source": [
    "# data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# data.export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# x,y = next(iter(data.valid_dl))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Combo"
   ]
  },
  {
   "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": [
    "class SPMTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[int]:\n",
    "        w_toks = sp.EncodeAsIds(t)[1:]\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
   },
   "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": [
    "# Bert_tok architectures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Word 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, 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, spatial_postions=None):\n",
    "#         if spatial_postions:\n",
    "#             rows,cols = spatial_positions\n",
    "#         else:\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 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.w_decoder = decoder\n",
    "        self.w_embed = embed\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",
    "    \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",
    "        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",
    "        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",
    "    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=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,
    "hidden": true
   },
   "source": [
    "## Word Arch w/ MLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Load above Word Arch modules plus the following:"
   ]
  },
  {
   "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",
    "            preds = torch.argmax(outs, dim=-1)\n",
    "        return self.lm(preds)[0]"
   ]
  },
  {
   "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(PATH/'distilbert')\n",
    "    lm.resize_token_embeddings(len(itos))\n",
    "    lm.train()\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": "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",
    "    lm = model.lm\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\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        feats = tfmr.encode(adaptor(img_enc(src)))\n",
    "\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)   #[bs,sl]\n",
    "            if (pred==0).all(): break\n",
    "            tgt = torch.cat([tgt,pred], dim=-1)\n",
    "        #out = torch.stack(res).transpose(1,0).contiguous()\n",
    "        out = lm(tgt)[0]\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Char 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, 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,
    "hidden": true
   },
   "source": [
    "## Combo 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, rows, cols, dropout=0.1):\n",
    "        super(LearnedPositionalEmbeddings, self).__init__()\n",
    "        self.nl_tok  = 30522 #4\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        self.embed = nn.Embedding(vocab, d_model, 0)\n",
    "        self.rows = nn.Embedding(15, d_model//2, 0)\n",
    "        self.c_cols = nn.Embedding(cols[0], d_model//2, 0)\n",
    "        self.w_cols = nn.Embedding(cols[1], d_model//2, 0)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, char, word):\n",
    "        return self.encode_one(char, 'char'), self.encode_one(word, 'word')\n",
    "    \n",
    "    def encode_one(self, x, kind):\n",
    "        rows,cols = self.encode_spatial_positions(x)\n",
    "        \n",
    "        x_cols = self.c_cols if kind=='char' else self.w_cols\n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = x_cols(torch.clamp(cols, max=x_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.resnet18(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 forward(self, x):\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return x.mul(8)"
   ]
  },
  {
   "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, embeddings, 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.embed = embeddings\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",
    "        #LPE\n",
    "        #c_emb,w_emb = self.embed(c_tgt, w_tgt)\n",
    "        #pos_enc\n",
    "        c_emb = self.embed(c_tgt)\n",
    "        w_emb = self.embed(w_tgt)\n",
    "\n",
    "        feats = self.encoder(src)\n",
    "        char_outs = self.c_decoder(c_emb, feats, c_mask)\n",
    "        word_outs = self.w_decoder(w_emb, feats, w_mask)\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 = 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, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\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",
    "        #LearnedPositionalEmbeddings(d_model, vocab, rows=15, cols=[100,60]),   #LPE\n",
    "        nn.Sequential( Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000) ),  #pos_enc\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 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()\n",
    "    transformer = make_full_model(len(itos), d_model, 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, 512, 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": {},
   "source": [
    "# SentencePiece Architectures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Word 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, 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.rows = nn.Embedding(15, d_model//2, 0)\n",
    "        self.w_cols = nn.Embedding(60, d_model//2, 0)\n",
    "        \n",
    "        self.dropout = nn.Dropout(p=dropout)\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.w_cols(torch.clamp(cols, max=self.w_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.resnet18(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 forward(self, x):\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return x.mul(8)"
   ]
  },
  {
   "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, embeddings, generator):\n",
    "        super(WordTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.w_decoder = decoder\n",
    "        self.embed = embeddings\n",
    "        self.generator = generator\n",
    "            \n",
    "    def forward(self, src, tgt):\n",
    "        tgt = rshift(tgt, 1).long()\n",
    "        mask = subsequent_mask(tgt.size(-1))\n",
    "        return self.w_decoder(self.embed(tgt), self.encoder(src), 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_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",
    "    model = WordTransformer(\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",
    "        LearnedPositionalEmbeddings(d_model, vocab, drops),  #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=None, seq_len=300):\n",
    "        if tgt is not None:   #train\n",
    "            feats = self.adaptor(self.img_enc(src))\n",
    "            outs = self.transformer(feats, tgt)\n",
    "            return self.transformer.generate(outs)\n",
    "        else:                 #predict\n",
    "            self.eval()\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",
    "                    emb = self.transformer.embed(tgt)\n",
    "                    #mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.w_decoder(emb, feats)#, 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",
    "                return torch.stack(res).transpose(1,0).contiguous()"
   ]
  },
  {
   "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",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor()\n",
    "    transformer = make_full_model(10000, 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])  #remove for export\n",
    "    return learn"
   ]
  },
  {
   "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, 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)\n",
    "# 50,873,424"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# sd = torch.load(PATH/'models/font_sp10k.pth', map_location=device)\n",
    "# learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Combo 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, rows, 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.rows = nn.Embedding(15, d_model//2, 0)\n",
    "        self.c_cols = nn.Embedding(cols[0], d_model//2, 0)\n",
    "        self.w_cols = nn.Embedding(cols[1], d_model//2, 0)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, char, word):\n",
    "        return self.encode_one(char, 'char'), self.encode_one(word, 'word')\n",
    "    \n",
    "    def encode_one(self, x, kind):\n",
    "        rows,cols = self.encode_spatial_positions(x)\n",
    "        \n",
    "        x_cols = self.c_cols if kind=='char' else self.w_cols\n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = x_cols(torch.clamp(cols, max=x_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.resnet18(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 forward(self, x):\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return x.mul(8)"
   ]
  },
  {
   "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, embeddings, generator):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        self.embed = embeddings\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",
    "        c_emb,w_emb = self.embed(c_tgt, w_tgt)\n",
    "        feats = self.encoder(src)\n",
    "        char_outs = self.c_decoder(c_emb, feats, c_mask)\n",
    "        word_outs = self.w_decoder(w_emb, feats, w_mask)\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, 1).long()\n",
    "        w_tgt = rshift(w_tgt, 1).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, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\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",
    "        LearnedPositionalEmbeddings(d_model, vocab, rows=15, cols=[100,60]),\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",
    "        return self.transformer.generate(char_outs, word_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",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor()\n",
    "    transformer = make_full_model(len(data.vocab), d_model, 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, 512, 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": [
    "## Combo Arch w/ integrated LM"
   ]
  },
  {
   "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, rows, 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.rows = nn.Embedding(15, d_model//2, 0)\n",
    "        self.c_cols = nn.Embedding(cols[0], d_model//2, 0)\n",
    "        self.w_cols = nn.Embedding(cols[1], d_model//2, 0)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, char, word):\n",
    "        return self.encode_one(char, 'char'), self.encode_one(word, 'word')\n",
    "    \n",
    "    def encode_one(self, x, kind):\n",
    "        rows,cols = self.encode_spatial_positions(x)\n",
    "        \n",
    "        x_cols = self.c_cols if kind=='char' else self.w_cols\n",
    "\n",
    "        row_t = self.rows(rows)            \n",
    "        col_t = x_cols(torch.clamp(cols, max=x_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.resnet18(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 forward(self, x):\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return x.mul(8)"
   ]
  },
  {
   "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, embeddings, generator, lm_enc):\n",
    "        super(WordCharTransformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.c_decoder = c_dec\n",
    "        self.w_decoder = w_dec\n",
    "        self.embed = embeddings\n",
    "        self.generator = generator\n",
    "        self.lm_enc = lm_enc\n",
    "    \n",
    "    def forward(self, src, mix_tgt):\n",
    "        c_tgt,w_tgt,c_mask,w_mask = self.shift_with_masks(mix_tgt)\n",
    "        c_emb,w_emb = self.embed(c_tgt, w_tgt)\n",
    "        feats = self.encoder(src)\n",
    "        char_outs = self.c_decoder(c_emb, feats, c_mask)\n",
    "        word_outs = self.w_decoder(w_emb, feats, w_mask)\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 lm(self, x):\n",
    "        # x => [bs, sl, vocab] : generator output\n",
    "        with torch.no_grad():\n",
    "            x = torch.argmax(x, dim=-1)\n",
    "        x = self.embed.embed(x)\n",
    "        #x = torch.matmul(F.softmax(x, dim=-1), self.embed.embed.weight)  # reverse embedding from vocab\n",
    "        return self.generator(self.lm_enc(x))\n",
    "    \n",
    "    def shift_with_masks(self, mix_tgt):\n",
    "        c_tgt,w_tgt = mix_tgt\n",
    "        c_tgt = rshift(c_tgt, 1).long()\n",
    "        w_tgt = rshift(w_tgt, 1).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, N=4, drops=0, heads=8):\n",
    "    c = deepcopy\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",
    "        LearnedPositionalEmbeddings(d_model, vocab, rows=15, cols=[100,60]),\n",
    "        nn.Linear(d_model, vocab),\n",
    "        nn.Sequential(\n",
    "            Lambda(lambda x: x.mul_(math.sqrt(d_model))),\n",
    "            PositionalEncoding(d_model, drops),\n",
    "            Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        )\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",
    "        c_outs,w_outs = self.transformer.generate(char_outs, word_outs)\n",
    "        return c_outs,w_outs,self.transformer.lm(w_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', 'word_lm'])\n",
    "            \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        c_out, w_out, w_lm = 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",
    "        lm_error,_   = cer(w_lm, w_targ, self.recon)\n",
    "        self.c_errors += c_error\n",
    "        self.w_errors += w_error\n",
    "        self.lm_errors += lm_error\n",
    "        self.total += size\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.c_errors, self.w_errors, self.lm_errors, self.total = 0, 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, self.lm_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",
    "        lm = loss(pred[2], w_targ)\n",
    "        #print(f'char loss: {cl}  word_loss: {wl}')\n",
    "        return cl + wl + lm"
   ]
  },
  {
   "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",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    adaptor = Adaptor()\n",
    "    transformer = make_full_model(len(data.vocab), d_model, 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, 512, 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": {},
   "source": [
    "# Experiments"
   ]
  },
  {
   "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.load('word_hw_sp10k'); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.to_fp16(); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = load_learner(PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "img = open_image(PATH/'test/a03-014.png');\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "res = learn.predict(img)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(str(res))"
   ]
  },
  {
   "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 = 4e-5 #1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(3, max_lr=lrs, callbacks=[SaveModelCallback(learn, name='word_hw_sp10k_half')])\n",
    "# data: small dataset, sz:256, bs:50, sentence_piece\n",
    "# train: 5cycle(1e-3)\n",
    "# arch: (512, 512, N=4, drops=0.1, heads=8); subsequent_masks\n",
    "\n",
    "# 10k, lpe, 67.7M\n",
    "# 71.242447\t68.054993\t0.759380\t0.841816\t08:04   freeze(), 1cycle(1e-3)\n",
    "# 12.448058\t13.792800\t0.082833\t0.123208\t08:18   unfreeze(), 5cycle(1e-4,1e-3)  'sm_sp10k_lpe'\n",
    "\n",
    "# 10k, lpe, 67.7M, 5cycle(1e-3)   \n",
    "# 10.024531\t10.917119\t0.059539\t0.100420\t08:22   'sm_sp10k_lpe2'  **\n",
    "# chars:    1.71376   .10433\n",
    "# words:    1.56690   .06228\n",
    "\n",
    "# 30k, lpe, 88.2M, 5cycle(1e-3)\n",
    "# 11.141268\t11.952768\t0.070033\t0.114384\t09:21   'sm_sp30k_lpe'\n",
    "\n",
    "# 50k, lpe, 108.7M, 5cycle(1e-3)\n",
    "# 12.930821\t12.563131\t0.097857\t0.161395\t11:04   'sm_sp50k_lpe2'\n",
    "\n",
    "\n",
    "# 10k, lpe, parallelogram_mask, 5cycle(1e-3)\n",
    "# 10.259628\t11.054597\t0.059808\t0.101713\t09:08\n",
    "\n",
    "\n",
    "# w/ WORD LM\n",
    "# 10k, lpe, attached word_lm, 80.3M  --  no good for inference:(\n",
    "# 5.786808\t6.193685\t0.054422\t0.026740\t09:50   'sm_sp10k_lm'   **\n",
    "\n",
    "# word_lm after generator, 80.3M\n",
    "# 30.129230\t29.837000\t0.053065\t0.672288\t09:11\n",
    "\n",
    "# word_lm after generator from embedding\n",
    "# 11.316447\t12.298356\t0.057172\t0.115761\t0.006366\t09:27    'sm_sp10k_lm2'\n",
    "# chars:    1.79026   .10513\n",
    "# words:    1.46591   .07781\n",
    "#    lm:    1.57080   .07781\n",
    "\n",
    "# word_lm from generator w/ reverse embedding\n",
    "# 18.072691\t18.354738\t0.079624\t0.202494\t0.034407\t10:21    'sm_sp10k_lm3'\n",
    "# chars:    1.92827   .14636\n",
    "# words:    1.67687   .13378\n",
    "#    lm:    1.87302   .13378\n",
    "\n",
    "# word_lm from generator, argmax, embedding, pos_enc\n",
    "# 17.686855\t18.746574\t0.059707\t0.110811\t0.114993\t09:31   'sm_sp10k_lm4'\n",
    "# chars:    2.21205   .10050\n",
    "# words:    1.32152   .07384\n",
    "#    lm:    1.37493   .08308\n",
    "\n",
    "\n",
    "# data: font generated, sz:512, bs:15, sentence_piece10k\n",
    "# train: 3cycle(1e-3)\n",
    "# arch: (512, 512, N=4, drops=0.1, heads=8); subsequent_masks\n",
    "\n",
    "# 10k, lpe, 67.7M, combo\n",
    "# 31.972498\t28.244387\t0.017639\t0.037072\t1:45:40   'font_sp10k'\n",
    "# chars:    32.2701   .01937\n",
    "# words:    41.7109   .02917\n",
    "\n",
    "# word_lm from generator w/ reverse embedding - preload sm_sp10k_lm3\n",
    "# 66.723541\t62.563965\t0.039113\t0.102822\t0.010772\t1:59:23   stopped after 2cycles, 'font_sp10k_lm'\n",
    "# chars:    49.6740   .05171\n",
    "# words:    51.9776   .03500\n",
    "#    lm:    59.4986   .03500\n",
    "\n",
    "# word_lm from generator, argmax, embedding, pos_enc - preload sm_sp10k_lm4\n",
    "# 84.748764\t66.507683\t0.034499\t0.074262\t0.069198\t1:58:18   stopped after 2cycles, 'font_sp10k_lm2'\n",
    "# chars:    73.3652   .10826\n",
    "# words:    65.1936   .06696\n",
    "#    lm:    65.5722   .06706\n",
    "\n",
    "\n",
    "# data: handwriting, sz:512, bs:15, sentence_piece10k\n",
    "# train: 5cycle(1e-3)\n",
    "# arch: (512, 512, N=4, drops=0.1, heads=8); subsequent_masks\n",
    "\n",
    "# 10k, lpe, 67.7M, combo\n",
    "# 14.077939\t11.943507\t0.006205\t0.003816\t21:38   'hw_sp10k'\n",
    "# chars:    50.1363   .00435\n",
    "# words:    5.17031   .00043\n",
    "# test pg:\n",
    "# chars:    144.698   .03923\n",
    "# words:    80.1309   .04630\n",
    "# test upload:\n",
    "# chars:    129.044   .26124\n",
    "# words:    61.2923   .30572\n",
    "\n",
    "# \"\", 5cycle(1e-6), preload hw_sp10k\n",
    "# 14.288121\t11.759130\t0.005498\t0.003153\t19:59    'hw_sp10k1'\n",
    "# test pg:\n",
    "# chars:    156.577   .03828\n",
    "# words:    79.5895   .04714\n",
    "# test upload:\n",
    "# chars:    130.981   .24922\n",
    "# words:    63.3666   .30366\n",
    "\n",
    "# \"\", 2cycle(1e-4), preload hw_sp10k\n",
    "# 12.602509\t11.120279\t0.005010\t0.002936\t20:04    'hw_sp10k5'\n",
    "# test pg:\n",
    "# chars:    152.536   .03506\n",
    "# words:    86.5458   .04597\n",
    "# test upload:\n",
    "# chars:    131.504   .35796\n",
    "# words:    59.7730   .27973\n",
    "\n",
    "# word only, 3cycle(1e-4), preload hw_sp10k\n",
    "# 6.081270\t6.041916\t0.000911\t14:51    'word_hw_sp10k'\n",
    "#     pg:    82.7050   .04591\n",
    "#     pg:    83.1463   .04556  no mask\n",
    "# upload:    64.2214   .31966\n",
    "# upload:    61.1524   .30548  no mask\n",
    "\n",
    "# word only, 3cycle(4e-5), preload word_hw_sp10k\n",
    "# 5.744105\t5.858494\t0.000798\t16:11    'word_hw_sp10k_half'\n",
    "#     pg:    83.2666   .04203\n",
    "# upload:    61.8928   .31536\n",
    "\n",
    "# \"\", 1cycle(1e-4), preload hw_sp10k1\n",
    "# 13.631779\t12.327535\t0.005439\t0.003256\t19:56    'hw_sp10k4'\n",
    "# test pg:\n",
    "# chars:    134.723   .03913\n",
    "# words:    72.2786   .04452\n",
    "\n",
    "# word-only, 50.8M, preload hw_sp10k\n",
    "# 6.726398\t6.525873\t0.001182\t15:27    2cycle(1e-6)   'hw_sp10k3'\n",
    "# greedy:    101.801   .01792\n",
    "#     pg:    96.8504   .05395\n",
    "# upload:    64.6677   .35998\n",
    "\n",
    "# data: handwriting+ 50k; sz:512, bs: 20, sp10k, 50.8M, word-only, preload font_sp10k\n",
    "# 4.633609\t4.357923\t0.010759\t19:25   'hw_sp10k2'\n",
    "# greedy:    42.3496   .06626\n",
    "#     pg:    95.0345   .06356\n",
    "# upload:    66.4512   .51736"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Previous Architecture testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# pg, bs: 10, N:4, d_model:512\n",
    "\n",
    "# alternate LearnedPositionalEmbeddings\n",
    "# 1767.251709\t1570.349365\t1.094300\t0.893096\t02:09   split(adaptor); freeze(1e-3)\n",
    "# 1018.636475\t1082.605713\t0.556753\t0.688827\t02:10   unfreeze(1e-4,1e-3) alt_pg_combo_bt\n",
    "\n",
    "# resnet18 (512, 512, N=4, drops=0.1, heads=8), 88M\n",
    "# 1819.735596\t1640.198608\t1.139695\t0.910257\t01:53\n",
    "# 978.301208\t1062.761963\t0.558057\t0.660668\t01:56   alt_pg_combo_bt2\n",
    "\n",
    "# resnet34 (512, 512, N=4, drops=0.1, heads=8), 98M\n",
    "# 1821.668213\t1645.444458\t1.102799\t0.891206\t01:59\n",
    "# 982.210693\t1061.083496\t0.554228\t0.664686\t02:03    alt_pg_combo_bt3\n",
    "\n",
    "\n",
    "# resnet18 (512, 512, N=4, drops=0.1, heads=16), 88M\n",
    "# 1810.988281\t1641.318848\t1.107212\t0.897255\t01:59\n",
    "# 965.121765\t1060.464966\t0.552585\t0.658893\t02:01   alt_pg_combo_bt4\n",
    "\n",
    "# resnet18 + lin adaptor (768, 512, N=4, drops=0.1, heads=12), 162M\n",
    "# 1725.002075\t1577.129517\t1.024268\t0.871293\t02:18\n",
    "# 977.914246\t1056.783325\t0.542184\t0.670891\t02:22   alt_pg_combo_bt5\n",
    "\n",
    "# resnet18 (512, 512, N=4, drops=0.1, heads=16), 88M, subsequent_masks\n",
    "# 1813.572388\t1627.146973\t1.082598\t0.888579\t01:58\n",
    "# 1042.711060\t1122.876343\t0.601146\t0.677406\t02:01   alt_pg_combo_bt6\n",
    "\n",
    "\n",
    "# resnet18 (512, 512, N=4, drops=0.1, heads=16), 88M [mod tfms]\n",
    "# 1825.002075\t1646.828247\t1.200644\t0.911882\t01:59\n",
    "# 986.249390\t1067.287598\t0.556860\t0.663619\t02:01   alt_pg_combo_bt7\n",
    "\n",
    "\n",
    "\n",
    "# SentencePiece\n",
    "# resnet18 (512, 512, N=4, drops=0.1, heads=16), 88M [mod tfms]\n",
    "# 1807.717773\t1669.616699\t1.265745\t0.939794\t01:35\n",
    "# 944.019165\t1058.646973\t0.548149\t0.669054\t01:38   alt_pg_combo_sp\n",
    "\n",
    "# N=6, 111M\n",
    "# 1814.066284\t1672.442993\t1.256078\t0.922056\t01:54\n",
    "# 961.308472\t1061.681641\t0.546759\t0.687847\t01:56   alt_pg_combo_sp2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# small dataset, sz:256, bs:50, 3cycle(1e-3)\n",
    "# arch: (512, 512, N=4, drops=0.1, heads=8); subsequent_masks\n",
    "\n",
    "# bert_tok, learned positional embedding\n",
    "# 23.695314\t22.344717\t0.163726\t0.301154\t09:50\n",
    "\n",
    "# bert_tok, pos_enc\n",
    "# 30.034079\t27.572166\t0.206749\t0.367589\t08:39\n",
    "\n",
    "# sentence_piece, learned positional embedding\n",
    "# 23.044443\t21.780771\t0.161438\t0.296342\t09:04\n",
    "\n",
    "# sentence_piece, pos_enc\n",
    "# 31.692009\t28.524887\t0.213417\t0.392681\t07:51\n",
    "\n",
    "# sentence_piece, lpe, 50k, 108.7M\n",
    "# 24.660820\t23.094442\t0.174975\t0.321934\t10:13\n",
    "\n",
    "# sentence_piece, lpe, 10k, 67.7M\n",
    "# 17.922606\t17.049728\t0.113365\t0.216181\t08:12"
   ]
  },
  {
   "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('font_combo_sp2', strict=False); None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split(lambda m: [m.adaptor]); None\n",
    "len(learn.layer_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze()\n",
    "lrs = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.unfreeze()\n",
    "lrs = slice(1e-5,1e-4)"
   ]
  },
  {
   "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": [
    "learn.fit_one_cycle(3, max_lr=lrs, callbacks=[SaveModelCallback(learn, name='font_combo_sp2')])\n",
    "\n",
    "# small, bs: 50\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",
    "# sentencepiece, modified tfms, resnet18 (512, 512, N=4, drops=0.1, heads=16), 88M\n",
    "# 74.434914\t70.777931\t0.816833\t0.876581\t09:21   freeze(), 1cycle(1e-3)\n",
    "# 12.942133\t15.643773\t0.080849\t0.165932\t09:28   unfreeze(), 5cycle(1e-4,1e-3)   *sm_combo_sp\n",
    "# chars:    2.22918   .17094\n",
    "# words:    1.13729   .17206\n",
    "\n",
    "# sentencepiece, \"\", preload alt_pg_combo_sp\n",
    "# 61.517582\t57.855293\t0.514289\t0.772304\t09:09   freeze()\n",
    "# 11.257002\t14.372008\t0.071675\t0.139939\t09:30   unfreeze(),    *sm_combo_sp2\n",
    "\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",
    "# 53.768700\t45.393501\t0.020888\t0.074038\t2:12:17   4epochs(1e-4,1e-3)    sentencepiece  *font_combo_sp\n",
    "# 50.356300\t44.293865\t0.016115\t0.061809\t2:18:10   1epoch(3e-6,2e-5)\n",
    "# 49.612297\t42.464523\t0.014771\t0.061766\t2:17:11   2epochs(1e-5,1e-4)    *font_combo_sp2\n",
    "# chars:    27.5662   .00351\n",
    "# words:    90.0695   .00650\n",
    "#TEST pg\n",
    "# chars:    216.829   .19460\n",
    "# words:    89.0265   .18346\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\n",
    "\n",
    "# w/ pretrained LM: 'distilbert_mlm_xtra', preload: 'hw_word_bt'\n",
    "# 2.109078\t1.383242\t0.000374\t15:17     *hw_word_bt_lm\n",
    "#     pg:    94.0378   .05911\n",
    "# upload:    64.4596   .30308\n",
    "\n",
    "# word_only, preload 'font_combo_bt2', bs: 16\n",
    "# 5.143809\t5.627481\t0.000719\t15:40   5cycle(1e-4,1e-3)    *hw_word_bt_alt\n",
    "#     pg:    64.6734   .06081\n",
    "# upload:    50.5408   .25785\n",
    "\n",
    "# word, preload 'hw_word_bt_alt', w/ pretrained bertLM, bs:5\n",
    "# 2.977787\t1.853431\t0.000390\t17:54    3cycle(1e-4)   *hw_word_bt_lm\n",
    "#     pg:    97.8368   .06604\n",
    "# upload:    62.4515   .27278"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": 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": {
    "heading_collapsed": true
   },
   "source": [
    "# Adjust State Dict and Split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "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.modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/hw_word_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,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "lm_sd['model'].keys()"
   ]
  },
  {
   "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 lm_sd['model'].items():\n",
    "    lm_k = k.replace('model', 'lm')\n",
    "    model_sd[lm_k] = v\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'].update(model_sd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].keys()"
   ]
  },
  {
   "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": {
    "hidden": 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": {},
   "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",
    "    decoder = tfmr.c_decoder if kind=='char' else tfmr.w_decoder\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",
    "            emb = tfmr.embed.encode_one(tgt, kind)\n",
    "#             emb = tfmr.embed(tgt)\n",
    "            \n",
    "            dec_outs = decoder(emb, 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",
    "        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": {},
   "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', 1)\n",
    "g_preds = learn.model(x.half())\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",
    "    #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
   },
   "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
   },
   "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', 1)\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": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(3,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\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
   },
   "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', 1, lm=False)\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(3,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\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": {},
   "source": [
    "# Test"
   ]
  },
  {
   "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": [
    "# 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 combo\n",
    "test_data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=SPMMultiList, sp=sp)\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": [
    "# sentencepiece word\n",
    "test_data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=SPList)\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": [
    "x,y = next(iter(test_data.train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "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
}