tree2seq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "from copy import deepcopy\n",
    "\n",
    "PAD = 0\n",
    "UNK = 1\n",
    "BOS = 2\n",
    "EOS = 3\n",
    "\n",
    "PAD_WORD = '<blank>'\n",
    "UNK_WORD = '<unk>'\n",
    "BOS_WORD = '<s>'\n",
    "EOS_WORD = '</s>'\n",
    "\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.utils.data as data\n",
    "\n",
    "class Tree(object):\n",
    "    def __init__(self):\n",
    "        self.parent = None\n",
    "        self.num_children = 0\n",
    "        self.children = list()\n",
    "\n",
    "    def add_child(self, child):\n",
    "        child.parent = self\n",
    "        self.num_children += 1\n",
    "        self.children.append(child)\n",
    "\n",
    "    def size(self):\n",
    "        if getattr(self, '_size'):\n",
    "            return self._size\n",
    "        count = 1\n",
    "        for i in range(self.num_children):\n",
    "            count += self.children[i].size()\n",
    "        self._size = count\n",
    "        return self._size\n",
    "\n",
    "    def depth(self):\n",
    "        if getattr(self, '_depth'):\n",
    "            return self._depth\n",
    "        count = 0\n",
    "        if self.num_children > 0:\n",
    "            for i in range(self.num_children):\n",
    "                child_depth = self.children[i].depth()\n",
    "                if child_depth > count:\n",
    "                    count = child_depth\n",
    "            count += 1\n",
    "        self._depth = count\n",
    "        return self._depth\n",
    "\n",
    "\n",
    "class Dataset(data.Dataset):\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "    def read_trees(self, filename):\n",
    "        with open(filename, 'r') as f:\n",
    "            trees = [self.read_tree(line) for line in tqdm(f.readlines())]\n",
    "        return trees\n",
    "\n",
    "    def read_tree(self, line):\n",
    "        parents = list(map(int, line.split()))\n",
    "        trees = dict()\n",
    "        root = None\n",
    "        for i in range(1, len(parents) + 1):\n",
    "            if i - 1 not in trees.keys() and parents[i - 1] != -1:\n",
    "                idx = i\n",
    "                prev = None\n",
    "                while True:\n",
    "                    parent = parents[idx - 1]\n",
    "                    if parent == -1:\n",
    "                        break\n",
    "                    tree = Tree()\n",
    "                    if prev is not None:\n",
    "                        tree.add_child(prev)\n",
    "                    trees[idx - 1] = tree\n",
    "                    tree.idx = idx - 1\n",
    "                    if parent - 1 in trees.keys():\n",
    "                        trees[parent - 1].add_child(tree)\n",
    "                        break\n",
    "                    elif parent == 0:\n",
    "                        root = tree\n",
    "                        break\n",
    "                    else:\n",
    "                        prev = tree\n",
    "                        idx = parent\n",
    "        return root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable as Var\n",
    "\n",
    "PAD = 0\n",
    "UNK = 1\n",
    "BOS = 2\n",
    "EOS = 3\n",
    "\n",
    "PAD_WORD = '<blank>'\n",
    "UNK_WORD = '<unk>'\n",
    "BOS_WORD = '<s>'\n",
    "EOS_WORD = '</s>'\n",
    "\n",
    "# module for childsumtreelstm\n",
    "class ChildSumTreeLSTM(nn.Module):\n",
    "    def __init__(self, in_dim, mem_dim):\n",
    "        super(ChildSumTreeLSTM, self).__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.mem_dim = mem_dim\n",
    "        self.ioux = nn.Linear(self.in_dim, 3 * self.mem_dim)\n",
    "        self.iouh = nn.Linear(self.mem_dim, 3 * self.mem_dim)\n",
    "        self.fx = nn.Linear(self.in_dim, self.mem_dim)\n",
    "        self.fh = nn.Linear(self.mem_dim, self.mem_dim)\n",
    "\n",
    "    def node_forward(self, inputs, child_c, child_h):\n",
    "        child_h_sum = torch.sum(child_h, dim=0, keepdim=True)\n",
    "\n",
    "        iou = self.ioux(inputs) + self.iouh(child_h_sum)\n",
    "        i, o, u = torch.split(iou, iou.size(1) // 3, dim=1)\n",
    "        i, o, u = F.sigmoid(i), F.sigmoid(o), F.tanh(u)\n",
    "\n",
    "        f = F.sigmoid(\n",
    "            self.fh(child_h) +\n",
    "            self.fx(inputs).repeat(len(child_h), 1)\n",
    "        )\n",
    "        fc = torch.mul(f, child_c)\n",
    "\n",
    "        c = torch.mul(i, u) + torch.sum(fc, dim=0, keepdim=True)\n",
    "        h = torch.mul(o, F.tanh(c))\n",
    "        return c, h\n",
    "\n",
    "    def forward(self, tree, inputs):\n",
    "        for idx in range(tree.num_children):\n",
    "            self.forward(tree.children[idx], inputs)\n",
    "\n",
    "        if tree.num_children == 0:\n",
    "            child_c = Var(inputs[0].data.new(1, self.mem_dim).fill_(0.))\n",
    "            child_h = Var(inputs[0].data.new(1, self.mem_dim).fill_(0.))\n",
    "        else:\n",
    "            child_c, child_h = zip(* map(lambda x: x.state, tree.children))\n",
    "            child_c, child_h = torch.cat(child_c, dim=0), torch.cat(child_h, dim=0)\n",
    "\n",
    "        tree.state = self.node_forward(inputs[tree.idx], child_c, child_h)\n",
    "        return tree.state\n",
    "\n",
    "\n",
    "# module for distance-angle similarity\n",
    "class Similarity(nn.Module):\n",
    "    def __init__(self, mem_dim, hidden_dim, num_classes):\n",
    "        super(Similarity, self).__init__()\n",
    "        self.mem_dim = mem_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_classes = num_classes\n",
    "        self.wh = nn.Linear(2 * self.mem_dim, self.hidden_dim)\n",
    "        self.wp = nn.Linear(self.hidden_dim, self.num_classes)\n",
    "\n",
    "    def forward(self, lvec, rvec):\n",
    "        mult_dist = torch.mul(lvec, rvec)\n",
    "        abs_dist = torch.abs(torch.add(lvec, -rvec))\n",
    "        vec_dist = torch.cat((mult_dist, abs_dist), 1)\n",
    "\n",
    "        out = F.sigmoid(self.wh(vec_dist))\n",
    "        out = F.log_softmax(self.wp(out), dim=1)\n",
    "        return out\n",
    "\n",
    "\n",
    "# putting the whole model together\n",
    "class SimilarityTreeLSTM(nn.Module):\n",
    "    def __init__(self, vocab_size, in_dim, mem_dim, hidden_dim, num_classes, sparsity, freeze):\n",
    "        super(SimilarityTreeLSTM, self).__init__()\n",
    "        self.emb = nn.Embedding(vocab_size, in_dim, padding_idx=PAD, sparse=sparsity)\n",
    "        if freeze:\n",
    "            self.emb.weight.requires_grad = False\n",
    "        self.childsumtreelstm = ChildSumTreeLSTM(in_dim, mem_dim)\n",
    "        self.similarity = Similarity(mem_dim, hidden_dim, num_classes)\n",
    "\n",
    "    def forward(self, ltree, linputs, rtree, rinputs):\n",
    "        linputs = self.emb(linputs)\n",
    "        rinputs = self.emb(rinputs)\n",
    "        lstate, lhidden = self.childsumtreelstm(ltree, linputs)\n",
    "        rstate, rhidden = self.childsumtreelstm(rtree, rinputs)\n",
    "        output = self.similarity(lstate, rstate)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LENGTH = 100\n",
    "class AttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size, dropout_p=0.1, max_length=MAX_LENGTH):\n",
    "        super(AttnDecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.dropout_p = dropout_p\n",
    "        self.max_length = max_length\n",
    "\n",
    "        self.embedding = nn.Embedding(self.output_size, self.hidden_size)\n",
    "        self.attn = nn.Linear(self.hidden_size * 2, self.max_length)\n",
    "        self.attn_combine = nn.Linear(self.hidden_size * 2, self.hidden_size)\n",
    "        self.dropout = nn.Dropout(self.dropout_p)\n",
    "        self.gru = nn.GRU(self.hidden_size, self.hidden_size)\n",
    "        self.out = nn.Linear(self.hidden_size, self.output_size)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        embedded = self.embedding(input).view(1, 1, -1)\n",
    "        embedded = self.dropout(embedded)\n",
    "\n",
    "        attn_weights = F.softmax(\n",
    "            self.attn(torch.cat((embedded[0], hidden[0]), 1)))\n",
    "        attn_applied = torch.bmm(attn_weights.unsqueeze(0),\n",
    "                                 encoder_outputs.unsqueeze(0))\n",
    "\n",
    "        output = torch.cat((embedded[0], attn_applied[0]), 1)\n",
    "        output = self.attn_combine(output).unsqueeze(0)\n",
    "\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "\n",
    "        output = F.log_softmax(self.out(output[0]))\n",
    "        return output, hidden, attn_weights\n",
    "\n",
    "    def initHidden(self):\n",
    "        result = Variable(torch.zeros(1, 1, self.hidden_size))\n",
    "        if use_cuda:\n",
    "            return result.cuda()\n",
    "        else:\n",
    "            return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "teacher_forcing_ratio = 0.5\n",
    "\n",
    "\n",
    "def train(input_variable, target_variable, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, max_length=MAX_LENGTH):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "\n",
    "    input_length = input_variable.size()[0]\n",
    "    target_length = target_variable.size()[0]\n",
    "\n",
    "    encoder_outputs = Variable(torch.zeros(max_length, encoder.hidden_size))\n",
    "    encoder_outputs = encoder_outputs.cuda() if use_cuda else encoder_outputs\n",
    "\n",
    "    loss = 0\n",
    "\n",
    "    for ei in range(input_length):\n",
    "        encoder_output, encoder_hidden = encoder(\n",
    "            input_variable[ei], encoder_hidden)\n",
    "        encoder_outputs[ei] = encoder_output[0][0]\n",
    "\n",
    "    decoder_input = Variable(torch.LongTensor([[SOS_token]]))\n",
    "    decoder_input = decoder_input.cuda() if use_cuda else decoder_input\n",
    "\n",
    "    decoder_hidden = encoder_hidden\n",
    "\n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_variable[di])\n",
    "            decoder_input = target_variable[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            ni = topi[0][0]\n",
    "\n",
    "            decoder_input = Variable(torch.LongTensor([[ni]]))\n",
    "            decoder_input = decoder_input.cuda() if use_cuda else decoder_input\n",
    "\n",
    "            loss += criterion(decoder_output, target_variable[di])\n",
    "            if ni == EOS_token:\n",
    "                break\n",
    "\n",
    "    loss.backward()\n",
    "\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.data[0] / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}