4_text_denoising.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text Denoising\n",
    "\n",
    "Inspired by \"Neural Networks for Text Correction and Completion in Keyboard Decoding\" by Shaona Ghosh and Per Ola Kristensson. https://arxiv.org/pdf/1709.06429.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gluoncv.data.batchify import Tuple, Stack, Append, Pad\n",
    "import gluonnlp as nlp\n",
    "import hnswlib # https://github.com/nmslib/hnswlib\n",
    "import mxboard\n",
    "import mxnet as mx\n",
    "from mxnet import gluon, autograd\n",
    "import numpy as np\n",
    "import re\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ocr.utils.encoder_decoder import get_transformer_encoder_decoder, Denoiser, encode_char, decode_char, LabelSmoothing, SoftmaxCEMaskedLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctx = mx.gpu() if mx.context.num_gpus() > 0 else mx.cpu()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See get_mode.py\n",
    "text_filepath = 'dataset/typo/all.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ALPHABET = ['<UNK>', '<PAD>', '<BOS>', '<EOS>']+list(' ' + string.ascii_letters + string.digits + string.punctuation)\n",
    "ALPHABET_INDEX = {letter: index for index, letter in enumerate(ALPHABET)} # { a: 0, b: 1, etc}\n",
    "FEATURE_LEN = 150 # max-length in characters for one document\n",
    "NUM_WORKERS = 8 # number of workers used in the data loading\n",
    "BATCH_SIZE = 64 # number of documents per batch\n",
    "MAX_LEN_SENTENCE = 150\n",
    "PAD = 1\n",
    "BOS = 2\n",
    "EOS = 3\n",
    "UNK = 0\n",
    "max_len_vocab = 500000\n",
    "\n",
    "moses_detokenizer = nlp.data.SacreMosesDetokenizer()\n",
    "moses_tokenizer = nlp.data.SacreMosesTokenizer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generic Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_knn_index():\n",
    "    model, vocab = nlp.model.big_rnn_lm_2048_512(dataset_name='gbw', pretrained=True, ctx=mx.cpu())\n",
    "\n",
    "    step = 1024\n",
    "    dim = 512\n",
    "    num_elements = max_len_vocab+step\n",
    "    data = np.zeros((num_elements, dim), dtype='float32')\n",
    "    data_labels = np.arange(max_len_vocab)\n",
    "    for i in tqdm(range(1, max_len_vocab, step)):\n",
    "        data[i:i+step,:] = model.embedding(mx.nd.arange(i,i+step)).asnumpy()\n",
    "    # Declaring index\n",
    "    p = hnswlib.Index(space = 'cosine', dim = dim) # possible options are l2, cosine or ip\n",
    "\n",
    "    # Initing index - the maximum number of elements should be known beforehand\n",
    "    p.init_index(max_elements = max_len_vocab, ef_construction = 200, M = 16)\n",
    "\n",
    "    # Element insertion (can be called several times):\n",
    "    p.add_items(data[:max_len_vocab], data_labels)\n",
    "    # Controlling the recall by setting ef:\n",
    "    p.set_ef(50) # ef should always be > k\n",
    "    return p, data, vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 489/489 [00:01<00:00, 382.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 45min 3s, sys: 28.7 s, total: 45min 32s\n",
      "Wall time: 1min 33s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# We get a knn index to substitute words\n",
    "knn_index, knn_data, knn_vocab  = get_knn_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NoisyTextDataset(mx.gluon.data.Dataset):\n",
    "    def __init__(self, \n",
    "                 text_filepath=None, \n",
    "                 substitute_costs_filepath='models/substitute_probs.json', \n",
    "                 insert_weight=1, \n",
    "                 delete_weight=1, \n",
    "                 glue_prob=0.05, \n",
    "                 substitute_weight=2,\n",
    "                 max_replace=0.3,\n",
    "                 is_train=True, \n",
    "                 split=0.9, \n",
    "                 data_type='corpus', \n",
    "                 gbw_corpus=None,\n",
    "                 knn_index=knn_index,\n",
    "                 knn_data=knn_data,\n",
    "                 knn_vocab=knn_vocab,\n",
    "                 knn_num=10,\n",
    "                 proba_synonym=0.1\n",
    "                ):\n",
    "        self.max_replace = max_replace\n",
    "        self.replace_weight = 0 #replace_prob  #Ignore typo dataset\n",
    "        self.substitute_threshold = float(substitute_weight) / (insert_weight + delete_weight + substitute_weight)\n",
    "        self.insert_threshold = self.substitute_threshold + float(insert_weight) / (insert_weight + delete_weight + substitute_weight)\n",
    "        self.delete_threshold = self.insert_threshold + float(delete_weight) / (insert_weight + delete_weight + substitute_weight)\n",
    "        self.glue_prob = glue_prob\n",
    "        self.substitute_dict = json.load(open(substitute_costs_filepath,'r'))\n",
    "        self.split = split\n",
    "        self.data_type = data_type\n",
    "        if self.data_type == 'corpus':\n",
    "            self.text = self._process_text(text_filepath, is_train)\n",
    "        elif self.data_type == 'GBW':\n",
    "            self.gbw_corpus = gbw_corpus\n",
    "        self.knn_index = knn_index\n",
    "        self.knn_data = knn_data\n",
    "        self.knn_vocab = knn_vocab\n",
    "        self.knn_num = knn_num\n",
    "        self.proba_synonym = proba_synonym\n",
    "    \n",
    "    def _process_text(self, filename, is_train):\n",
    "        with open(filename, 'r', encoding='Latin-1') as f:\n",
    "            text = []\n",
    "            for line in f.readlines():\n",
    "                if line != '':\n",
    "                    text.append(line.strip())\n",
    "            \n",
    "            split_index = int(self.split*len(text))\n",
    "            if is_train:\n",
    "                text = text[:split_index]\n",
    "            else:\n",
    "                text = text[split_index:]\n",
    "        return text\n",
    "    \n",
    "    def _replace_synonym(self, line):\n",
    "        processed_line = self._pre_process_line(line)\n",
    "        words = []\n",
    "        num_words = self.knn_num\n",
    "        for i, word in enumerate(processed_line):\n",
    "            draw = random.random()\n",
    "            if word in self.knn_vocab and self.knn_vocab[word] < max_len_vocab and draw < self.proba_synonym and word not in string.punctuation :\n",
    "                index_list = self.knn_index.knn_query(self.knn_data[self.knn_vocab[word]], k=num_words)[0][0]\n",
    "                word = self.knn_vocab.idx_to_token[index_list[random.randint(0,num_words-1)]]\n",
    "            words.append(word)\n",
    "        return self._post_process_line(words)\n",
    "    \n",
    "    def _transform_line(self, line):\n",
    "        \"\"\"\n",
    "        replace words that are in the typo dataset with a typo\n",
    "        with a probability `self.replace_proba`\n",
    "        \"\"\"\n",
    "        output = []\n",
    "        \n",
    "        processed_line = self._pre_process_line(line)\n",
    "        \n",
    "        # We get randomly the index of the modifications\n",
    "        num_chars = len(''.join(processed_line))\n",
    "        if num_chars:\n",
    "            index_modifications = np.random.choice(num_chars, random.randint(0, int(self.max_replace*num_chars)), replace=False)\n",
    "            substitute_letters = []\n",
    "            insert_letters = []\n",
    "            delete_letters = []\n",
    "            # We randomly assign these indices to modifications based on precalculated thresholds\n",
    "            for index in index_modifications:\n",
    "                draw = random.random()\n",
    "                if draw < self.substitute_threshold:\n",
    "                    substitute_letters.append(index)\n",
    "                    continue\n",
    "                if draw < self.insert_threshold:\n",
    "                    insert_letters.append(index)\n",
    "                    continue\n",
    "                else:\n",
    "                    delete_letters.append(index)\n",
    "                            \n",
    "        \n",
    "        j = 0\n",
    "        for i, word in enumerate(processed_line):\n",
    "            \n",
    "            if word != '' and word not in string.punctuation:\n",
    "                \n",
    "                len_word = len(word)\n",
    "                word_ = []\n",
    "                k = j\n",
    "                for letter in word:\n",
    "                    if k in substitute_letters and letter in self.substitute_dict:\n",
    "                        draw = random.random()\n",
    "                        for replace, prob in self.substitute_dict[letter].items():\n",
    "                            if draw < prob:\n",
    "                                letter = replace\n",
    "                                break\n",
    "                    word_.append(letter)\n",
    "                    k += 1\n",
    "                word = ''.join(word_)\n",
    "                                \n",
    "                # Insert random letter\n",
    "                k = j\n",
    "                word_ = []\n",
    "                for letter in word:\n",
    "                    if k in insert_letters:\n",
    "                        word_.append(ALPHABET[random.randint(4, len(ALPHABET)-1)])\n",
    "                    word_.append(letter)\n",
    "                    k += 1\n",
    "                word = ''.join(word_)\n",
    "                \n",
    "                # Delete random letter\n",
    "                k = j\n",
    "                word_ = []\n",
    "                for letter in word:\n",
    "                    if k not in delete_letters:\n",
    "                        word_.append(letter)\n",
    "                    k += 1\n",
    "                word = ''.join(word_)\n",
    "                    \n",
    "                output.append(word)\n",
    "            else:\n",
    "                output.append(word)\n",
    "            j += len(word)\n",
    "\n",
    "        output_ = [\"\"]*len(output)\n",
    "        j = 0\n",
    "        for i, word in enumerate(output):\n",
    "            output_[j] += word\n",
    "            if random.random() > self.glue_prob:\n",
    "                j += 1\n",
    "        \n",
    "        line = self._post_process_line(output_)\n",
    "        return line.strip()\n",
    "    \n",
    "    def _pre_process_line(self, line):\n",
    "        line = line.replace('\\n','').replace('`',\"'\").replace('--',' -- ')\n",
    "        return moses_tokenizer(line)\n",
    "        \n",
    "    def _post_process_line(self, words):\n",
    "        output = ' '.join(moses_detokenizer(words))\n",
    "        return output\n",
    "    \n",
    "    def _match_caps(self, original, typo):\n",
    "        if original.isupper():\n",
    "            return typo.upper()\n",
    "        elif original.istitle():\n",
    "            return typo.capitalize()\n",
    "        else:\n",
    "            return typo\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        if self.data_type == 'GBW':\n",
    "            tokens = moses_detokenizer(self.gbw_corpus[idx][:-1])\n",
    "            if len(tokens) > 6:\n",
    "                start = random.randint(0, len(tokens)-3)\n",
    "                end = random.randint(start, len(tokens))\n",
    "                tokens = tokens[start:end]\n",
    "            line = ' '.join(tokens)\n",
    "        else:\n",
    "            line = self.text[idx]\n",
    "        line = self._replace_synonym(line)\n",
    "        line_typo = self._transform_line(line)\n",
    "        return line_typo, line\n",
    "\n",
    "    def __len__(self):\n",
    "        if self.data_type == 'GBW':\n",
    "            return len(self.gbw_corpus)\n",
    "        else:\n",
    "            return len(self.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_char(text, src=True):\n",
    "    encoded = np.ones(FEATURE_LEN, dtype='float32') * PAD\n",
    "    text = text[:FEATURE_LEN-2]\n",
    "    i = 0\n",
    "    if not src:\n",
    "        encoded[0] = BOS\n",
    "        i = 1\n",
    "    for letter in text:\n",
    "        if letter in ALPHABET_INDEX:\n",
    "            encoded[i] = ALPHABET_INDEX[letter]\n",
    "        i += 1\n",
    "    encoded[i] = EOS\n",
    "    return encoded, np.array([i+1]).astype('float32')\n",
    "\n",
    "def encode_word(text, src=True):\n",
    "    tokens = tokenizer(text)\n",
    "    indices = vocab[tokens]\n",
    "    indices += [vocab['<EOS>']]\n",
    "    indices = [vocab['<BOS>']]+indices\n",
    "    return indices, np.array([len(indices)]).astype('float32')\n",
    "\n",
    "def transform(data, label):\n",
    "    src, src_valid_length = encode_char(data, src=True)\n",
    "    tgt, tgt_valid_length = encode_char(label, src=False)\n",
    "    return src, src_valid_length, tgt, tgt_valid_length, data, label\n",
    "\n",
    "def decode_char(text):\n",
    "    output = []\n",
    "    for val in text:\n",
    "        if val == EOS:\n",
    "            break\n",
    "        elif val == PAD or val == BOS:\n",
    "            continue\n",
    "        output.append(ALPHABET[int(val)])\n",
    "    return \"\".join(output)\n",
    "\n",
    "\n",
    "def decode_word(indices):\n",
    "    return detokenizer([vocab.idx_to_token[int(i)] for i in indices], return_str=True).replace('<PAD>','')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We test our synonym replacer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'decontamination'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word = \"test\"\n",
    "num_words = 200\n",
    "index_list = knn_index.knn_query(knn_data[knn_vocab[word]], k=1000)[0][0]\n",
    "knn_vocab.idx_to_token[index_list[random.randint(0,num_words-1)]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = NoisyTextDataset(text_filepath=text_filepath, glue_prob=0.2, is_train=True).transform(transform)\n",
    "dataset_test = NoisyTextDataset(text_filepath=text_filepath, glue_prob=0.2, is_train=False).transform(transform)\n",
    "\n",
    "# Finetuning on the text from the IAM dataset\n",
    "dataset_train_ft = NoisyTextDataset(text_filepath='dataset/typo/text_train.txt', is_train=True, split=1.0, knn_index=knn_index).transform(transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([42., 25.,  7.,  5., 23., 11., 24.,  4., 19.,  7., 59.,  7.,  5.,\n",
       "        23., 65., 13., 19., 18.,  5., 16.,  4., 11., 16., 19., 17., 20.,\n",
       "        23.,  9., 18.,  4., 19., 10.,  4., 12.,  9., 31., 16., 20., 13.,\n",
       "        18.,  9.,  4., 23., 19., 14., 19., 25., 22., 18., 78.,  4., 18.,\n",
       "        19., 24.,  3.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([55.], dtype=float32),\n",
       " array([ 2., 42., 25.,  7.,  5., 23.,  4., 11., 19., 24.,  4., 19.,  7.,\n",
       "         7.,  5., 23., 13., 19., 18.,  5., 16.,  4., 11., 16., 13., 17.,\n",
       "        20., 23.,  9., 23.,  4., 19., 10.,  4., 24., 12.,  9.,  4., 31.,\n",
       "        16., 20., 13., 18.,  9.,  4., 23., 19., 14., 19., 25., 22., 18.,\n",
       "        78.,  4., 18., 19., 24.,  3.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([58.], dtype=float32),\n",
       " 'Lucasgt oc2cas8ional glompsen of heAlpine sojourn, not',\n",
       " 'Lucas got occasional glimpses of the Alpine sojourn, not')"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_train[random.randint(0, len(dataset_train)-1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We are now trying to release all our books one month --In advance\n",
      "You are now trying to release all our songs one month in advance\n",
      "We are already trying to release everyone our books one month IN advance\n",
      "We are now trying to release all our books one month in advance\n",
      "We are now trying must release all our books one month --in advance\n",
      "We are now trying to release all our books another month in advance\n",
      "We are now trying to release all our books one month in advance\n",
      "We are now trying would release all our novels one month in advance\n",
      "We are now trying to release all our books one month in advance\n",
      "We are now trying to release all our books one month in advance\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(dataset_train[42][5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validation data being the IAM Dataset prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = json.load(open('dataset/typo/validating.json','r'))\n",
    "data_ = []\n",
    "for label, modified in data:\n",
    "    if label.strip() != modified.strip():\n",
    "        data_.append([label, modified])\n",
    "val_dataset_ft = gluon.data.ArrayDataset(list(list(zip(*data_))[1]), list(list(zip(*data_))[0])).transform(transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([24., 27., 19.,  4., 10., 19., 25., 18.,  8., 23.,  4., 17.,  5.,\n",
       "        18., 11.,  9., 78.,  4., 19., 18.,  9.,  4., 12., 13., 23.,  4.,\n",
       "        20., 13., 18.,  9.,  5., 20., 20., 16.,  9., 78.,  4., 19., 18.,\n",
       "         9.,  4., 24., 19., 22., 15.,  4., 12.,  5., 18., 80.,  3.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([51.], dtype=float32),\n",
       " array([ 2., 24., 27., 19.,  4., 20., 19., 25., 18.,  8., 23.,  4., 17.,\n",
       "         5., 22., 11.,  9., 78.,  4., 19., 18.,  9.,  4., 24., 13., 18.,\n",
       "         4., 20., 13., 18.,  9.,  5., 20., 20., 16.,  9., 78.,  4., 19.,\n",
       "        18.,  9.,  4., 55., 19., 22., 15.,  4., 12.,  5., 17., 80., 68.,\n",
       "         3.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([53.], dtype=float32),\n",
       " 'two founds mange, one his pineapple, one tork han.',\n",
       " 'two pounds marge, one tin pineapple, one York ham.\"')"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_dataset_ft[random.randint(0, len(val_dataset_ft)-1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training on GBW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading /home/ubuntu/.mxnet/datasets/gbw/1-billion-word-language-modeling-benchmark-r13output.tar.gz from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/dataset/gbw/1-billion-word-language-modeling-benchmark-r13output.tar.gz...\n"
     ]
    }
   ],
   "source": [
    "gbw_stream = nlp.data.GBWStream(segment='train', skip_empty=True, bos=None, eos='<EOS>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "for e, corpus in enumerate(gbw_stream):\n",
    "    dataset_gbw = NoisyTextDataset(gbw_corpus=corpus, data_type='GBW').transform(transform)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([13., 18., 24., 19.,  4.,  9., 10., 10., 25.,  7., 24.,  4., 10.,\n",
       "        19., 22.,  4.,  7., 19., 16., 13.,  9., 11.,  9.,  4., 10., 19.,\n",
       "        19., 24.,  6.,  5., 16., 16., 73., 23.,  4., 45., 23., 23., 18.,\n",
       "        68., 11.,  9.,  4., 32., 19., 27., 16.,  4., 79.,  4., 18.,  3.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([52.], dtype=float32),\n",
       " array([ 2., 13., 18., 24., 19.,  4.,  9., 10., 10.,  9.,  7., 24.,  4.,\n",
       "        10., 19., 22.,  4.,  7., 19., 16., 16.,  9., 11.,  9.,  4., 10.,\n",
       "        19., 19., 24.,  6.,  5., 16., 16., 73., 23.,  4., 45., 22.,  5.,\n",
       "        18., 11.,  9.,  4., 32., 19., 27., 16.,  4., 79., 79., 19., 18.,\n",
       "         3.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,\n",
       "         1.,  1.,  1.,  1.,  1.,  1.,  1.], dtype=float32),\n",
       " array([53.], dtype=float32),\n",
       " 'into effuct for coliege football\\'s Ossn\"ge Bowl - n',\n",
       " \"into effect for college football's Orange Bowl --on\")"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_gbw[7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### DataLoaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batchify_list(elem):\n",
    "    output = []\n",
    "    for e in elem:\n",
    "        output.append(elem)\n",
    "    return output\n",
    "    \n",
    "batchify = Tuple(Stack(), Stack(), Stack(), Stack(), batchify_list, batchify_list)\n",
    "batchify_word = Tuple(Stack(), Stack(), Pad(), Stack(), batchify_list, batchify_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Process ForkPoolWorker-4:\n",
      "Process ForkPoolWorker-5:\n",
      "Process ForkPoolWorker-2:\n",
      "Process ForkPoolWorker-13:\n",
      "Process ForkPoolWorker-3:\n",
      "Process ForkPoolWorker-8:\n",
      "Process ForkPoolWorker-11:\n",
      "Process ForkPoolWorker-1:\n",
      "Process ForkPoolWorker-6:\n",
      "Process ForkPoolWorker-12:\n",
      "Process ForkPoolWorker-7:\n",
      "Process ForkPoolWorker-10:\n",
      "Process ForkPoolWorker-14:\n",
      "Process ForkPoolWorker-9:\n",
      "Process ForkPoolWorker-15:\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 335, in get\n",
      "    res = self._reader.recv_bytes()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/pool.py\", line 108, in worker\n",
      "    task = get()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 216, in recv_bytes\n",
      "    buf = self._recv_bytes(maxlength)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 335, in get\n",
      "    res = self._reader.recv_bytes()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 335, in get\n",
      "    res = self._reader.recv_bytes()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/queues.py\", line 334, in get\n",
      "    with self._rlock:\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 407, in _recv_bytes\n",
      "    buf = self._recv(4)\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 216, in recv_bytes\n",
      "    buf = self._recv_bytes(maxlength)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 216, in recv_bytes\n",
      "    buf = self._recv_bytes(maxlength)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/synchronize.py\", line 96, in __enter__\n",
      "    return self._semlock.__enter__()\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 407, in _recv_bytes\n",
      "    buf = self._recv(4)\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 379, in _recv\n",
      "    chunk = read(handle, remaining)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 407, in _recv_bytes\n",
      "    buf = self._recv(4)\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 379, in _recv\n",
      "    chunk = read(handle, remaining)\n",
      "  File \"/home/ubuntu/anaconda3/lib/python3.6/multiprocessing/connection.py\", line 379, in _recv\n",
      "    chunk = read(handle, remaining)\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n"
     ]
    }
   ],
   "source": [
    "train_data = gluon.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, last_batch='rollover', batchify_fn=batchify, num_workers=5)\n",
    "test_data = gluon.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=True, last_batch='keep', batchify_fn=batchify, num_workers=5)\n",
    "val_data_ft = gluon.data.DataLoader(val_dataset_ft, batch_size=BATCH_SIZE, shuffle=True, last_batch='keep', batchify_fn=batchify, num_workers=0)\n",
    "train_data_ft = gluon.data.DataLoader(dataset_train_ft, batch_size=BATCH_SIZE, shuffle=True, last_batch='rollover', batchify_fn=batchify, num_workers=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper function to help train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(net, iterator):\n",
    "    loss = 0\n",
    "    for i, (src, src_valid_length, tgt, tgt_valid_length, typo, label) in enumerate(iterator):\n",
    "        src = src.as_in_context(ctx)\n",
    "        tgt = tgt.as_in_context(ctx)\n",
    "        src_valid_length = src_valid_length.as_in_context(ctx).squeeze()\n",
    "        tgt_valid_length = tgt_valid_length.as_in_context(ctx).squeeze()\n",
    "        output = net(src, tgt[:,:-1], src_valid_length, tgt_valid_length-1)\n",
    "        ls = loss_function_test(output, tgt[:,1:], tgt_valid_length).mean()\n",
    "        loss += ls.asscalar()\n",
    "    print(\"[Test Typo     ] {}\".format(decode_char(src[0].asnumpy())))\n",
    "    print(\"[Test Predicted] {}\".format(get_sentence(net, decode_char(src[0].asnumpy()))))\n",
    "    print(\"[Test Correct  ] {}\".format(decode_char(tgt[0].asnumpy())))\n",
    "    return loss / (i+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_epoch(net, epoch, train_iterator, test_iterator, trainer):\n",
    "    loss = 0.\n",
    "    for i, (src, src_valid_length, tgt, tgt_valid_length, typo, label) in enumerate(train_iterator):\n",
    "        src = src.as_in_context(ctx)\n",
    "        tgt = tgt.as_in_context(ctx)\n",
    "        src_valid_length = src_valid_length.as_in_context(ctx).squeeze()\n",
    "        tgt_valid_length = tgt_valid_length.as_in_context(ctx).squeeze()\n",
    "        \n",
    "        with autograd.record():\n",
    "            output = net(src, tgt[:,:-1], src_valid_length, tgt_valid_length-1)\n",
    "            smoothed_label = label_smoothing(tgt[:,1:])\n",
    "            ls = loss_function(output, smoothed_label, tgt_valid_length).mean()\n",
    "        \n",
    "        ls.backward()\n",
    "        trainer.step(src.shape[0])\n",
    "        loss += ls.asscalar()\n",
    "        \n",
    "        if i % send_every_n == 0:\n",
    "            val_loss = evaluate(net, test_iterator)\n",
    "            sw.add_scalar(tag='Val_Loss_it', value={key:val_loss}, global_step=i+e*len(train_iterator))\n",
    "            sw.add_scalar(tag='Train_Loss_it', value={key:loss/(i+1)}, global_step=i+e*len(train_iterator))\n",
    "            print(\"[Iteration {} Train] {}\".format(i, loss / (i+1)))\n",
    "            print(\"[Iteration {} Test ] {}\".format(i, val_loss))\n",
    "            print(\"[Train Typo        ] {}\".format(decode_char(src[0].asnumpy())))\n",
    "            print(\"[Train Predicted   ] {}\".format(decode_char(output[0].asnumpy().argmax(axis=1))))\n",
    "            print(\"[Train Correct     ] {}\".format(decode_char(tgt[0].asnumpy())))\n",
    "            print()\n",
    "            sw.flush()\n",
    "\n",
    "    test_loss = evaluate(net, test_iterator)\n",
    "    print(\"Epoch [{}], Train Loss {:.4f}, Test Loss {:.4f}\".format(e, loss/(i+1), test_loss))\n",
    "    sw.add_scalar(tag='Train_Loss', value={key:loss/(i+1)}, global_step=e)\n",
    "    sw.add_scalar(tag='Test_Loss', value={key:test_loss}, global_step=e)\n",
    "    print()\n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sentence(net, sentence):\n",
    "    scorer = nlp.model.BeamSearchScorer(alpha=0, K=2, from_logits=False)\n",
    "    beam_sampler = nlp.model.BeamSearchSampler(beam_size=5,\n",
    "                                           decoder=net.decode_logprob,\n",
    "                                           eos_id=EOS,\n",
    "                                           scorer=scorer,\n",
    "                                           max_length=150)\n",
    "    src_seq, src_valid_length = encode_char(sentence)\n",
    "    src_seq = mx.nd.array([src_seq], ctx=ctx)\n",
    "    src_valid_length = mx.nd.array(src_valid_length, ctx=ctx)\n",
    "    encoder_outputs, _ = net.encode(src_seq, valid_length=src_valid_length)\n",
    "    states = net.decoder.init_state_from_encoder(encoder_outputs, \n",
    "                                                      encoder_valid_length=src_valid_length)\n",
    "    inputs = mx.nd.full(shape=(1,), ctx=src_seq.context, dtype=np.float32, val=BOS)\n",
    "    samples, scores, valid_lengths = beam_sampler(inputs, states)\n",
    "    samples = samples[0].asnumpy()\n",
    "    scores = scores[0].asnumpy()\n",
    "    valid_lengths = valid_lengths[0].asnumpy()\n",
    "    return decode_char(samples[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_heads = 16\n",
    "embed_size = 512\n",
    "num_layers = 2\n",
    "\n",
    "epochs = 5\n",
    "key = 'language_denoising'\n",
    "best_test_loss = 10e20\n",
    "\n",
    "learning_rate = 0.00004\n",
    "send_every_n = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_test_loss = 10e20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = './logs/text_denoising'\n",
    "checkpoint_dir = \"model_checkpoint\"\n",
    "checkpoint_name = key+\".params\"\n",
    "sw = mxboard.SummaryWriter(logdir=log_dir, flush_secs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Denoiser(alphabet_size=len(ALPHABET), max_src_length=FEATURE_LEN, max_tgt_length=FEATURE_LEN, num_heads=num_heads, embed_size=embed_size, num_layers=num_layers)\n",
    "net.initialize(mx.init.Xavier(), ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preparing the loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dim = len(ALPHABET)\n",
    "label_smoothing = LabelSmoothing(epsilon=0.002, units=output_dim)\n",
    "loss_function_test = SoftmaxCEMaskedLoss(sparse_label=True)\n",
    "loss_function = SoftmaxCEMaskedLoss(sparse_label=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (os.path.isfile(os.path.join(checkpoint_dir, checkpoint_name))):\n",
    "    net.load_parameters(os.path.join(checkpoint_dir, checkpoint_name), ctx=ctx)    \n",
    "    print(\"Loaded parameters\")\n",
    "    best_test_loss = evaluate(net, val_data_ft)\n",
    "    print(best_test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'models/denoiser2.params'\n",
    "if (os.path.isfile(model_path)):\n",
    "    net.load_parameters(model_path, ctx=ctx)    \n",
    "    print(\"Loaded parameters\")\n",
    "    best_test_loss = evaluate(net, val_data_ft)\n",
    "    print(best_test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': 0.0001})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training on the public novel dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 'language_denoising'\n",
    "for e in range(epochs):\n",
    "    test_loss = run_epoch(net, e, train_data, val_data_ft, trainer)\n",
    "    if test_loss < best_test_loss:\n",
    "        print(\"Saving network, previous best test loss {:.6f}, current test loss {:.6f}\".format(best_test_loss, test_loss))\n",
    "        denoiser.save_parameters(os.path.join(checkpoint_dir, checkpoint_name))\n",
    "        best_test_loss = test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training on the GBW dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 'language_denoising_gbw'\n",
    "for e, corpus in enumerate(gbw_stream):\n",
    "    dataset_gbw = NoisyTextDataset(gbw_corpus=corpus, data_type='GBW').transform(transform)\n",
    "    train_data_gbw = gluon.data.DataLoader(dataset_gbw, batch_size=BATCH_SIZE, shuffle=True, last_batch='discard', batchify_fn=batchify, num_workers=5)\n",
    "    test_loss = run_epoch(net, e, train_data_gbw, val_data_ft, trainer)\n",
    "    if test_loss < best_test_loss:\n",
    "        print(\"Saving network, previous best test loss {:.6f}, current test loss {:.6f}\".format(best_test_loss, test_loss))\n",
    "        net.save_parameters(os.path.join(checkpoint_dir, checkpoint_name))\n",
    "        best_test_loss = test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fine-tuning on the IAM training dataset text to learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 'language_denoising_ft'\n",
    "for e in range(epochs):\n",
    "    test_loss = run_epoch(net, e, train_data_ft, val_data_ft, trainer)\n",
    "    if test_loss < best_test_loss:\n",
    "        print(\"Saving network, previous best test loss {:.6f}, current test loss {:.6f}\".format(best_test_loss, test_loss))\n",
    "        net.save_parameters(os.path.join(checkpoint_dir, checkpoint_name))\n",
    "        best_test_loss = test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manual Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence = \"This sentence contains an eror\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_sentence(net, sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Appendix (maybe useful later)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create text file with all vocab words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, vocab = nlp.model.big_rnn_lm_2048_512(dataset_name='gbw', pretrained=True, ctx=mx.cpu())\n",
    "vocab_ = '\\n'.join(vocab.idx_to_token)\n",
    "with open('dataset/typo/vocab.txt', 'w') as f:\n",
    "    f.write(vocab_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create KNN lookup for words"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}