03a_translate_en_to_es_sinkhorn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having trained the English and Spanish LMs (along with the embeddings which constitute part of the encoder), it is now time to attempt translation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from decoder_head.core import *\n",
    "from decoder_head.data import *\n",
    "from fastai2.text.all import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_en= make_vocab(pd.read_pickle('data/en-100_tok/counter.pkl'), max_vocab=4000)\n",
    "vocab_es= make_vocab(pd.read_pickle('data/es-100_tok/counter.pkl'), max_vocab=4000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_es_dict = get_en_es_dict(vocab_en, vocab_es)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2451"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(en_es_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = 'data/en-100_tok/'\n",
    "mult = 4\n",
    "bs = 80\n",
    "seq_len = 70\n",
    "\n",
    "lm = DataBlock(blocks=(TextBlock(vocab=vocab_en, is_lm=True),),\n",
    "                get_x=read_tokenized_file,\n",
    "                get_items=partial(get_text_files, folders=['train', 'valid']),\n",
    "                splitter=FuncSplitter(lambda itm: itm.parent.name == 'valid'))\n",
    "\n",
    "dbunch_lm = lm.databunch(path, path=path, bs=bs, seq_len=seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export core\n",
    "class PermuteEmbeddingsSinkhorn(nn.Module):\n",
    "    def __init__(self, num_embeddings, embedding_dim, padding_idx):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.num_embeddings = num_embeddings\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.padding_idx = padding_idx\n",
    "        \n",
    "        # to conform to nn.Embedding api\n",
    "        self.max_norm=None\n",
    "        self.norm_type=2.0\n",
    "        self.scale_grad_by_freq=False\n",
    "        self.sparse = False\n",
    "\n",
    "        self.weight = nn.Parameter( torch.Tensor(num_embeddings, embedding_dim) )\n",
    "        self.p = nn.Parameter( torch.eye(self.num_embeddings) )\n",
    "        self.p.requires_grad = False\n",
    "        \n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def forward(self, words):        \n",
    "        p1 = self.p / self.p.sum(dim=1).unsqueeze(1)\n",
    "        p2 = p1 / p1.sum(dim=0).unsqueeze(0)\n",
    "            \n",
    "        return F.embedding(words, self.p @ self.weight)\n",
    "    \n",
    "    def reset_parameters(self): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export core\n",
    "class p_sinkhornAWD_LSTM(AWD_LSTM):\n",
    "    def __init__(self, vocab_sz, emb_sz, n_hid, n_layers, pad_token=1, hidden_p=0.2, input_p=0.6, embed_p=0.1,\n",
    "                 weight_p=0.5, bidir=False, packed=False):\n",
    "        store_attr(self, 'emb_sz,n_hid,n_layers,pad_token,packed')\n",
    "        self.bs = 1\n",
    "        self.n_dir = 2 if bidir else 1\n",
    "        self.encoder = PermuteEmbeddingsSinkhorn(vocab_sz, emb_sz, padding_idx=pad_token)\n",
    "\n",
    "        self.encoder_dp = self.encoder # chosing to initally train without embedding dropout\n",
    "        self.rnns = nn.ModuleList([self._one_rnn(emb_sz if l == 0 else n_hid, (n_hid if l != n_layers - 1 else emb_sz)//self.n_dir,\n",
    "                                                 bidir, weight_p, l) for l in range(n_layers)])\n",
    "        self.encoder.weight.data.uniform_(-self.initrange, self.initrange)\n",
    "        self.input_dp = RNNDropout(input_p)\n",
    "        self.hidden_dps = nn.ModuleList([RNNDropout(hidden_p) for l in range(n_layers)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export core\n",
    "import fastai2 \n",
    "fastai2.text.models.core._model_meta[p_sinkhornAWD_LSTM] = fastai2.text.models.core._model_meta[AWD_LSTM]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = language_model_learner(\n",
    "    dbunch_lm,\n",
    "    p_sinkhornAWD_LSTM,\n",
    "    opt_func=opt,\n",
    "    pretrained=False,\n",
    "    config=awd_lstm_lm_config,\n",
    "    drop_mult=0.1,\n",
    "    metrics=[accuracy, top_k_accuracy]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attemp translation from English to Spanish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<fastai2.text.learner.LMLearner at 0x7f087ff01a50>"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.load('pLSTM_en')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "█\r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(#3) [2.8178272247314453,0.42678970098495483,0.6835475564002991]"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.validate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.model[0].encoder.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.optimize_permutation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.model[0].encoder.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn.init.kaiming_normal_(learn.model[0].encoder.p);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.model[0].encoder.weight.data = torch.load('data/embeddings_es.torch').data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.model[0].encoder.p.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.model[0].encoder.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.loss_func = lambda preds, targs: aza_loss(learn, preds, targs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss  accuracy  top_k_accuracy  time    \n",
      "0         4.966329    4.935486    0.265000  0.465976        40:22     \n"
     ]
    }
   ],
   "source": [
    "learn.fit_one_cycle(1, 1e-3,  moms=(0.8, 0.7, 0.8), wd=1e-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.save('pLSTM_sinkorhn_en-es')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<fastai2.text.learner.LMLearner at 0x7f087ff01a50>"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.load('pLSTM_sinkorhn_en-es')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "█\r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(#3) [4.935486316680908,0.26499998569488525,0.4659762382507324]"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learn.validate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def top_n_translation_acc(top_n=1):\n",
    "    p = learn.model[0].encoder.p\n",
    "    hits = 0\n",
    "    for k, v in en_es_dict.items():\n",
    "        idx_en = vocab_en.index(k)\n",
    "        for vv in v:\n",
    "            idx_es = vocab_es.index(vv)\n",
    "            if idx_es in p[idx_en].argsort(descending=True)[:top_n]:\n",
    "                hits += 1\n",
    "#                 print(k, vv)\n",
    "                break\n",
    "    return hits/len(en_es_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.033047735618115054"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_n_translation_acc(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# top_n calculated a take on nearest neighbors, performs better but not by much\n",
    "\n",
    "def top_n_translation_acc(top_n=1):\n",
    "    p = learn.model[0].encoder.p\n",
    "    learned_embeddings = learn.model[0].encoder.p @ learn.model[0].encoder.weight\n",
    "    \n",
    "    hits = 0\n",
    "    for k, v in en_es_dict.items():\n",
    "        idx_en = vocab_en.index(k)\n",
    "        diff = learn.model[0].encoder.weight - learned_embeddings[idx_en]\n",
    "        candidates = diff.sum(-1).argsort(descending=False)[:top_n]\n",
    "        \n",
    "        for vv in v:\n",
    "            idx_es = vocab_es.index(vv)\n",
    "            if idx_es in candidates:\n",
    "                hits += 1\n",
    "                break\n",
    "    return hits/len(en_es_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.03386372909016728"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_n_translation_acc(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, this attempt at translation failed. In the next notebook, we will make another attempt, this time with normalized embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_data.ipynb.\n",
      "Converted 01_train_LM_en.ipynb.\n",
      "Converted 02_train_LM_es.ipynb.\n",
      "Converted 03_translate_en_to_es.ipynb.\n",
      "Converted 99_index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "from nbdev.export import *\n",
    "notebook2script()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}