seq2seq_with_attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run seq2seq.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.layer1 = nn.Linear(dim * 2, dim)\n",
    "        self.layer2 = nn.Linear(dim, 1)\n",
    "        \n",
    "    def forward(self, source_h, target_h, source_mask=None):\n",
    "        \n",
    "        num_sentences, num_words = source_h.shape[:-1]\n",
    "        target_h = target_h.repeat(1, num_words).reshape(num_words * num_sentences, -1)\n",
    "        mlp_input = torch.cat((target_h, source_h.reshape(num_words * num_sentences, -1)\n",
    "                              ), dim=1).reshape(num_words * num_sentences, -1)\n",
    "        \n",
    "        h1 = torch.tanh(self.layer1(mlp_input))\n",
    "        h2 = self.layer2(h1)\n",
    "        \n",
    "        raw_weights = h2.reshape(num_sentences, num_words)\n",
    "        \n",
    "        if source_mask is not None:\n",
    "            raw_weights = source_mask * raw_weights\n",
    "        weights = torch.softmax(f.relu(raw_weights), dim=1)\n",
    "        \n",
    "        return torch.sum(weights.unsqueeze(-1) * source_h, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "source_h = torch.Tensor(np.random.randint(0, 10, (2, 4, 10)))\n",
    "target_h = torch.Tensor(np.random.randint(0, 10, (2, 10)))\n",
    "a = Attention(10)\n",
    "a(source_h, target_h)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderAttention(nn.Module):\n",
    "    def __init__(self, num_embeddings, embedding_dim, hidden_dim, num_layers, padding_dim=0, start_dim=1):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(num_embeddings, embedding_dim, padding_dim)\n",
    "        self.decoder = nn.LSTM(embedding_dim + hidden_dim, hidden_dim, num_layers, batch_first=True)\n",
    "        self.word_predictor = nn.Linear(hidden_dim, num_embeddings)\n",
    "        self.start_dim = start_dim\n",
    "        \n",
    "        self.attention = Attention(hidden_dim)\n",
    "        \n",
    "    def forward(self, output, h, c, input=None, max_sen_len=20, source_mask=None):\n",
    "        \n",
    "        source_h = output\n",
    "        num_sentences, num_words, hidden_dim = output.shape\n",
    "        words_selected = torch.cuda.LongTensor([[self.start_dim] for _ in range(num_sentences)])\n",
    "\n",
    "        decoder_context = torch.zeros(num_sentences, hidden_dim).cuda()\n",
    "        pred = []\n",
    "        \n",
    "        if input is not None:\n",
    "            max_sen_len = input.shape[1]\n",
    "            teacher_words = input.t()\n",
    "        \n",
    "        for i in range(max_sen_len):\n",
    "            embeddings = self.embedding(words_selected)\n",
    "            decoder_input = torch.cat((embeddings, decoder_context.unsqueeze(1)), dim=-1)\n",
    "            y_t, (h, c) = self.decoder(decoder_input, (h, c))\n",
    "            orig_shape = y_t.shape\n",
    "            pred_t = self.word_predictor(y_t.reshape(orig_shape[0]*orig_shape[1], \n",
    "                                                     -1)).reshape(*orig_shape[:-1], -1).squeeze(1)\n",
    "            pred.append(pred_t)\n",
    "            decoder_context = self.attention(source_h, h[-1], source_mask)\n",
    "            \n",
    "            if input is None:\n",
    "                words_selected = torch.max(pred_t, dim=1)[1].unsqueeze(1)\n",
    "            else:\n",
    "                words_selected = teacher_words[i].unsqueeze(1)\n",
    "\n",
    "\n",
    "        pred = torch.stack(pred, dim=1)\n",
    "        return torch.log_softmax(pred, dim=-1), None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqWithAttention(nn.Module):\n",
    "    def __init__(self, source_vocab_len, target_vocab_len, embedding_dim, hidden_dim, num_layers=1,\n",
    "                       padding_dim_source=0, start_dim_target=1, padding_dim_target=0):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(source_vocab_len, embedding_dim, hidden_dim, num_layers, padding_dim_source)\n",
    "        self.decoder = DecoderAttention(target_vocab_len, embedding_dim, hidden_dim, num_layers, padding_dim_target, start_dim_target)\n",
    "        \n",
    "    def forward(self, source_input, target_input=None, source_mask=None):\n",
    "        output, h, c = self.encoder(source_input)\n",
    "        return self.decoder(output, h, c, target_input, source_mask=source_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "s = Seq2SeqWithAttention(10, 12, 9, 8)\n",
    "s.cuda()\n",
    "input = torch.cuda.LongTensor(np.random.randint(0, 10, (3, 6)))\n",
    "target = torch.cuda.LongTensor(np.random.randint(0, 12, (3, 10))\n",
    "\n",
    "# without teacher forcing\n",
    "s(input)\n",
    "\n",
    "# with teacher forcing\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AI",
   "language": "python",
   "name": "ai"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}